emissary/internal/agent/controller/app/server.go

package app

import (
	"context"
	"net/http"
	"strings"
	"sync"
	"time"

	"forge.cadoles.com/Cadoles/emissary/internal/agent/controller/app/spec"
	appSpec "forge.cadoles.com/Cadoles/emissary/internal/agent/controller/app/spec"
	"forge.cadoles.com/Cadoles/emissary/internal/proxy/wildcard"
	edgeHTTP "forge.cadoles.com/arcad/edge/pkg/http"
	authHTTP "forge.cadoles.com/arcad/edge/pkg/module/auth/http"
	"gitlab.com/wpetit/goweb/logger"

	"forge.cadoles.com/arcad/edge/pkg/bundle"
	"github.com/go-chi/chi/middleware"
	"github.com/go-chi/chi/v5"
	"github.com/lestrrat-go/jwx/v2/jwa"
	"github.com/lestrrat-go/jwx/v2/jwk"
	"github.com/pkg/errors"

	_ "forge.cadoles.com/Cadoles/emissary/internal/imports/passwd"
)

const defaultCookieDuration time.Duration = 24 * time.Hour

type Server struct {
	bundle         bundle.Bundle
	handlerOptions []edgeHTTP.HandlerOptionFunc
	server         *http.Server
	serverMutex    sync.RWMutex
	auth           *appSpec.Auth
}

func (s *Server) Start(ctx context.Context, addr string) (err error) {
	if s.server != nil {
		if err := s.Stop(); err != nil {
			return errors.WithStack(err)
		}
	}

	router := chi.NewRouter()

	router.Use(middleware.Logger)

	handler := edgeHTTP.NewHandler(s.handlerOptions...)
	if err := handler.Load(s.bundle); err != nil {
		return errors.Wrap(err, "could not load app bundle")
	}

	if err := s.configureAuth(router, s.auth); err != nil {
		return errors.WithStack(err)
	}

	router.Handle("/*", handler)

	server := &http.Server{
		Addr:    addr,
		Handler: router,
	}

	go func() {
		defer func() {
			if recovered := recover(); recovered != nil {
				if err, ok := recovered.(error); ok {
					logger.Error(ctx, err.Error(), logger.E(errors.WithStack(err)))

					return
				}

				panic(recovered)
			}
		}()

		defer func() {
			if err := s.Stop(); err != nil {
				panic(errors.WithStack(err))
			}
		}()

		if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			panic(errors.WithStack(err))
		}
	}()

	s.serverMutex.Lock()
	s.server = server
	s.serverMutex.Unlock()

	return nil
}

func (s *Server) Running() bool {
	s.serverMutex.RLock()
	defer s.serverMutex.RUnlock()

	return s.server != nil
}

func (s *Server) Stop() error {
	if s.server == nil {
		return nil
	}

	defer func() {
		s.serverMutex.Lock()
		s.server = nil
		s.serverMutex.Unlock()
	}()

	if err := s.server.Close(); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (s *Server) configureAuth(router chi.Router, auth *spec.Auth) error {
	if auth == nil {
		return nil
	}

	switch {
	case auth.Local != nil:
		var rawKey any = s.auth.Local.Key
		if strKey, ok := rawKey.(string); ok {
			rawKey = []byte(strKey)
		}

		key, err := jwk.FromRaw(rawKey)
		if err != nil {
			return errors.WithStack(err)
		}

		cookieDuration := defaultCookieDuration
		if s.auth.Local.CookieDuration != "" {
			cookieDuration, err = time.ParseDuration(s.auth.Local.CookieDuration)
			if err != nil {
				return errors.WithStack(err)
			}
		}

		if s.auth.Local.CookieDomain != "" {
			router.Use(invalidCookieDomainRedirect(s.auth.Local.CookieDomain))
		}

		router.Handle("/auth/*", authHTTP.NewLocalHandler(
			jwa.HS256, key,
			authHTTP.WithRoutePrefix("/auth"),
			authHTTP.WithAccounts(s.auth.Local.Accounts...),
			authHTTP.WithCookieOptions(s.auth.Local.CookieDomain, cookieDuration),
		))
	}

	return nil
}

func NewServer(bundle bundle.Bundle, auth *appSpec.Auth, handlerOptions ...edgeHTTP.HandlerOptionFunc) *Server {
	return &Server{
		bundle:         bundle,
		auth:           auth,
		handlerOptions: handlerOptions,
	}
}

func invalidCookieDomainRedirect(cookieDomain string) func(http.Handler) http.Handler {
	domain := strings.TrimPrefix(cookieDomain, ".")
	hostPattern := "*" + domain

	return func(h http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			hostParts := strings.SplitN(r.Host, ":", 2)

			if !wildcard.Match(hostParts[0], hostPattern) {
				url := r.URL

				newHost := domain
				if len(hostParts) > 1 {
					newHost += ":" + hostParts[1]
				}

				url.Host = newHost

				http.Redirect(w, r, url.String(), http.StatusTemporaryRedirect)

				return
			}

			h.ServeHTTP(w, r)
		}

		return http.HandlerFunc(fn)
	}
}