gemserve/cmd/gemserve/gemserve.go

package main

import (
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"gemserve/lib/apperrors"
	"gemserve/lib/logging"

	"gemserve/config"
	"gemserve/server"

	"git.antanst.com/antanst/uid"
)

func main() {
	config.CONFIG = *config.GetConfig()
	logging.SetupLogging(config.CONFIG.LogLevel)
	logger := logging.Logger
	ctx := logging.WithLogger(context.Background(), logger)
	err := runApp(ctx)
	if err != nil {
		logger.Error("Fatal Error", "err", err)
		panic(fmt.Sprintf("Fatal Error: %v", err))
	}
	os.Exit(0)
}

func runApp(ctx context.Context) error {
	logger := logging.FromContext(ctx)
	logger.Info("Starting up. Press Ctrl+C to exit")

	listenAddr := config.CONFIG.ListenAddr

	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)

	// Only this file should send to this channel.
	// All external functions should return errors.
	fatalErrors := make(chan error)

	// Root context, used to cancel
	// connections and graceful shutdown.
	serverCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	// WaitGroup to track active connections
	// in order to be able to wait until
	// they are properly dropped
	var wg sync.WaitGroup

	// Spawn server on the background.
	// Returned errors are considered fatal.
	go func() {
		err := startServer(serverCtx, listenAddr, &wg, fatalErrors)
		if err != nil {
			fatalErrors <- apperrors.NewFatalError(fmt.Errorf("server startup failed: %w", err))
		}
	}()

	for {
		select {
		case <-signals:
			logger.Warn("Received SIGINT or SIGTERM signal, shutting down gracefully")
			cancel()
			wg.Wait()
			return nil
		case fatalError := <-fatalErrors:
			cancel()
			wg.Wait()
			return fatalError
		}
	}
}

func startServer(ctx context.Context, listenAddr string, wg *sync.WaitGroup, fatalErrors chan<- error) (err error) {
	logger := logging.FromContext(ctx)

	cert, err := tls.LoadX509KeyPair(config.CONFIG.TLSCert, config.CONFIG.TLSKey)
	if err != nil {
		return apperrors.NewFatalError(fmt.Errorf("failed to load TLS certificate/key: %w", err))
	}

	logger.Debug("Using TLS cert", "path", config.CONFIG.TLSCert)
	logger.Debug("Using TLS key", "path", config.CONFIG.TLSKey)

	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert},
		MinVersion:   tls.VersionTLS12,
	}

	listener, err := tls.Listen("tcp", listenAddr, tlsConfig)
	if err != nil {
		return apperrors.NewFatalError(err)
	}

	defer func(listener net.Listener) {
		_ = listener.Close()
	}(listener)

	// If context is cancelled, close listener
	// to unblock Accept() inside main loop.
	go func() {
		<-ctx.Done()
		_ = listener.Close()
	}()

	logger.Info("Server listening", "address", listenAddr)

	for {
		conn, err := listener.Accept()
		if err != nil {
			if ctx.Err() != nil {
				return nil
			} // ctx cancellation
			logger.Info("Failed to accept connection: %v", "error", err)
			continue
		}

		// At this point we have a new connection.
		wg.Add(1)
		go func() {
			defer wg.Done()

			// Type assert the connection to TLS connection
			tlsConn, ok := conn.(*tls.Conn)
			if !ok {
				logger.Error("Connection is not a TLS connection")
				_ = conn.Close()
				return
			}

			remoteAddr := conn.RemoteAddr().String()
			connId := uid.UID()

			// Create cancellable connection context
			// with connection ID.
			connLogger := logging.WithAttr(logger, "id", connId)
			connLogger = logging.WithAttr(connLogger, "remoteAddr", remoteAddr)
			connCtx := context.WithValue(ctx, server.CtxConnIdKey, connId)
			connCtx = context.WithValue(connCtx, logging.CtxLoggerKey, connLogger)
			connCtx, cancel := context.WithTimeout(connCtx, time.Duration(config.CONFIG.ResponseTimeout)*time.Second)
			defer cancel()

			err := server.HandleConnection(connCtx, tlsConn)
			if err != nil {
				if apperrors.IsFatal(err) {
					fatalErrors <- err
					return
				}
				connLogger.Info("Connection failed", "error", err)
			}
		}()
	}
}