vibeproxy/internal/proxy/server.go

package proxy

import (
	"context"
	"crypto/tls"
	"database/sql"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"time"

	"github.com/owen/vibeproxy/internal/certificates"
	"github.com/owen/vibeproxy/internal/logging"
	"github.com/owen/vibeproxy/internal/metrics"
	"github.com/owen/vibeproxy/internal/runtime"
)

// Server owns the plain-HTTP and TLS listeners that serve all proxy hosts.
type Server struct {
	httpSrv  *http.Server
	httpsSrv *http.Server
	log      *slog.Logger
}

type Options struct {
	HTTPAddr  string // e.g. ":80"
	HTTPSAddr string // e.g. ":443"
}

func NewServer(opts Options, store *runtime.Store, certMgr *certificates.Manager, db *sql.DB, log *slog.Logger) *Server {
	reqLog := logging.NewRequestLogger(db, log)
	go reqLog.Run(context.Background())

	handler := NewHandler(store, certMgr, reqLog, db, log)

	tlsConfig := &tls.Config{
		MinVersion: tls.VersionTLS12,
		NextProtos: []string{"h2", "http/1.1"},
		GetCertificate: func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
			cert, err := certMgr.GetCertificateForDomain(hello.ServerName)
			if err != nil {
				metrics.TLSHandshakesTotal.WithLabelValues("error").Inc()
				return nil, err
			}
			metrics.TLSHandshakesTotal.WithLabelValues("ok").Inc()
			return cert, nil
		},
	}

	httpAddr := opts.HTTPAddr
	if httpAddr == "" {
		httpAddr = ":80"
	}
	httpsAddr := opts.HTTPSAddr
	if httpsAddr == "" {
		httpsAddr = ":443"
	}

	return &Server{
		log: log,
		httpSrv: &http.Server{
			Addr:         httpAddr,
			Handler:      handler,
			ReadTimeout:  0, // per-request timeouts are enforced inside the handler
			WriteTimeout: 0,
			IdleTimeout:  120 * time.Second,
			ErrorLog:     slog.NewLogLogger(log.Handler(), slog.LevelWarn),
		},
		httpsSrv: &http.Server{
			Addr:        httpsAddr,
			Handler:     handler,
			TLSConfig:   tlsConfig,
			IdleTimeout: 120 * time.Second,
			ErrorLog:    slog.NewLogLogger(log.Handler(), slog.LevelWarn),
		},
	}
}

// Start runs both listeners; it blocks until either fails or ctx is done.
func (s *Server) Start(ctx context.Context) error {
	errCh := make(chan error, 2)

	go func() {
		s.log.Info("http listener starting", "addr", s.httpSrv.Addr)
		if err := s.httpSrv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			errCh <- fmt.Errorf("http listener: %w", err)
		}
	}()

	go func() {
		s.log.Info("https listener starting", "addr", s.httpsSrv.Addr)
		// Certificates are served dynamically via GetCertificate, so we
		// don't pass cert/key file paths here.
		if err := s.httpsSrv.ListenAndServeTLS("", ""); err != nil && !errors.Is(err, http.ErrServerClosed) {
			errCh <- fmt.Errorf("https listener: %w", err)
		}
	}()

	select {
	case <-ctx.Done():
		return s.Shutdown(context.Background())
	case err := <-errCh:
		s.Shutdown(context.Background())
		return err
	}
}

// Shutdown gracefully drains in-flight requests/connections before
// closing the listeners.
func (s *Server) Shutdown(ctx context.Context) error {
	shutdownCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	err1 := s.httpSrv.Shutdown(shutdownCtx)
	err2 := s.httpsSrv.Shutdown(shutdownCtx)
	if err1 != nil {
		return err1
	}
	return err2
}