feat: TLS, HTTP2, pre-start functions, logging improvements

cli.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
+	"go.uber.org/zap/zapcore"
 )
 
 // AddGlobalFlag will make a new global flag on the root command.
@@ -46,13 +47,34 @@ func (r *Router) setupCLI() {
 		Use:     filepath.Base(os.Args[0]),
 		Version: r.api.Version,
 		Run: func(cmd *cobra.Command, args []string) {
+			// Call any pre-start functions.
+			for _, f := range r.prestart {
+				f()
+			}
+
 			if viper.GetBool("debug") {
-
+				if logLevel != nil {
+					logLevel.SetLevel(zapcore.DebugLevel)
+				}
 			}
 
-			if err := r.Listen(fmt.Sprintf("%s:%v", viper.Get("host"), viper.Get("port"))); err != nil {
-				panic(err)
+			// Start either an HTTP or HTTPS server based on whether TLS cert/key
+			// paths were given.
+			cert := viper.GetString("cert")
+			key := viper.GetString("key")
+			if cert == "" && key == "" {
+				if err := r.Listen(fmt.Sprintf("%s:%v", viper.Get("host"), viper.Get("port"))); err != nil {
+					panic(err)
+				}
 			}
+
+			if cert != "" && key != "" {
+				if err := r.ListenTLS(fmt.Sprintf("%s:%v", viper.Get("host"), viper.Get("port")), cert, key); err != nil {
+					panic(err)
+				}
+			}
+
+			panic("must pass key and cert for TLS")
 		},
 	}
 
@@ -79,5 +101,7 @@ func (r *Router) setupCLI() {
 
 	r.AddGlobalFlag("host", "", "Hostname", "0.0.0.0")
 	r.AddGlobalFlag("port", "p", "Port", 8888)
+	r.AddGlobalFlag("cert", "", "SSL certificate file path", "")
+	r.AddGlobalFlag("key", "", "SSL key file path", "")
 	r.AddGlobalFlag("debug", "d", "Enable debug logs", false)
 }

middelware.go

@@ -11,21 +11,33 @@ import (
 	"go.uber.org/zap/zapcore"
 )
 
+var logLevel *zap.AtomicLevel
+
 // LogMiddleware creates a new middleware to set a tagged `*zap.SugarLogger` in the
-// Gin context under the `log` key. It debug logs request info. If the current
-// terminal is a TTY, it will try to use colored output automatically.
-func LogMiddleware() func(*gin.Context) {
-	var l *zap.Logger
+// Gin context under the `log` key. It debug logs request info. If passed `nil`
+// for the logger, then it creates one. If the current terminal is a TTY, it
+// will try to use colored output automatically.
+func LogMiddleware(l *zap.Logger, tags map[string]string) func(*gin.Context) {
 	var err error
-	if isatty.IsTerminal(os.Stdout.Fd()) || isatty.IsCygwinTerminal(os.Stdout.Fd()) {
-		config := zap.NewDevelopmentConfig()
-		config.EncoderConfig.EncodeLevel = zapcore.CapitalColorLevelEncoder
-		l, err = config.Build()
-	} else {
-		l, err = zap.NewProduction()
+	if l == nil {
+		if isatty.IsTerminal(os.Stdout.Fd()) || isatty.IsCygwinTerminal(os.Stdout.Fd()) {
+			config := zap.NewDevelopmentConfig()
+			logLevel = &config.Level
+			config.EncoderConfig.EncodeLevel = zapcore.CapitalColorLevelEncoder
+			l, err = config.Build()
+		} else {
+			config := zap.NewProductionConfig()
+			logLevel = &config.Level
+			l, err = config.Build()
+		}
+		if err != nil {
+			panic(err)
+		}
 	}
-	if err != nil {
-		panic(err)
+
+	// Add any additional tags that were passed.
+	for k, v := range tags {
+		l = l.With(zap.String(k, v))
 	}
 
 	return func(c *gin.Context) {

router.go

@@ -125,9 +125,10 @@ func getRequestBody(c *gin.Context, t reflect.Type, op *Operation) (interface{},
 
 // Router handles API requests.
 type Router struct {
-	api    *OpenAPI
-	engine *gin.Engine
-	root   *cobra.Command
+	api      *OpenAPI
+	engine   *gin.Engine
+	root     *cobra.Command
+	prestart []func()
 }
 
 // NewRouter creates a new Huma router for handling API requests with
@@ -138,7 +139,7 @@ func NewRouter(api *OpenAPI) *Router {
 	g := gin.New()
 	g.Use(gin.Recovery())
 	g.Use(cors.Default())
-	g.Use(LogMiddleware())
+	g.Use(LogMiddleware(nil, nil))
 	return NewRouterWithGin(g, api)
 }
 
@@ -150,8 +151,9 @@ func NewRouterWithGin(engine *gin.Engine, api *OpenAPI) *Router {
 	}
 
 	r := &Router{
-		api:    api,
-		engine: engine,
+		api:      api,
+		engine:   engine,
+		prestart: []func(){},
 	}
 
 	r.setupCLI()
@@ -187,6 +189,11 @@ func (r *Router) GinEngine() *gin.Engine {
 	return r.engine
 }
 
+// PreStart registers a function to run before server start.
+func (r *Router) PreStart(f func()) {
+	r.prestart = append(r.prestart, f)
+}
+
 // Use attaches middleware to the router.
 func (r *Router) Use(middleware ...gin.HandlerFunc) {
 	r.engine.Use(middleware...)
@@ -353,6 +360,11 @@ func (r *Router) Listen(addr string) error {
 	return r.engine.Run(addr)
 }
 
+// ListenTLS listens for new connections using HTTPS & HTTP2
+func (r *Router) ListenTLS(addr, certFile, keyFile string) error {
+	return r.engine.RunTLS(addr, certFile, keyFile)
+}
+
 // Run executes the router command.
 func (r *Router) Run() {
 	if err := r.root.Execute(); err != nil {