output trace in slog and override correlation header name (#1986)

bobcallaway · web-flow · commit 19cd558bdb34 · 2024-02-02T17:43:30.000-05:00
Signed-off-by: Bob Callaway &lt;bcallaway@google.com&gt;
diff --git a/cmd/rekor-server/app/root.go b/cmd/rekor-server/app/root.go
@@ -22,6 +22,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/go-chi/chi/middleware"
 	homedir "github.com/mitchellh/go-homedir"
 	"github.com/sigstore/rekor/pkg/log"
 	"github.com/spf13/cobra"
@@ -131,6 +132,9 @@ Memory and file-based signers should only be used for testing.`)
 	rootCmd.PersistentFlags().Int("search_index.mysql.max_open_connections", 0, "maximum open connections")
 	rootCmd.PersistentFlags().Int("search_index.mysql.max_idle_connections", 0, "maximum idle connections")
 
+	rootCmd.PersistentFlags().String("http-request-id-header-name", middleware.RequestIDHeader, "name of HTTP Request Header to use as request correlation ID")
+	rootCmd.PersistentFlags().String("trace-string-prefix", "", "if set, this will be used to prefix the 'trace' field when outputting structured logs")
+
 	if err := viper.BindPFlags(rootCmd.PersistentFlags()); err != nil {
 		log.Logger.Fatal(err)
 	}
diff --git a/cmd/rekor-server/app/serve.go b/cmd/rekor-server/app/serve.go
@@ -21,6 +21,7 @@ import (
 	"time"
 
 	"cloud.google.com/go/profiler"
+	"github.com/go-chi/chi/middleware"
 	"github.com/go-openapi/loads"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/spf13/cobra"
@@ -63,7 +64,7 @@ var serveCmd = &cobra.Command{
 	Long:  `Starts a http server and serves the configured api`,
 	Run: func(cmd *cobra.Command, args []string) {
 		// Setup the logger to dev/prod
-		log.ConfigureLogger(viper.GetString("log_type"))
+		log.ConfigureLogger(viper.GetString("log_type"), viper.GetString("trace-string-prefix"))
 
 		// workaround for https://github.com/sigstore/rekor/issues/68
 		// from https://github.com/golang/glog/commit/fca8c8854093a154ff1eb580aae10276ad6b1b5f
@@ -76,6 +77,9 @@ var serveCmd = &cobra.Command{
 		}
 		log.Logger.Infof("starting rekor-server @ %v", viStr)
 
+		// overrides the correlation ID printed in logs, if config is set
+		middleware.RequestIDHeader = viper.GetString("http-request-id-header-name")
+
 		if viper.GetBool("gcp_cloud_profiling.enabled") {
 			cfg := profiler.Config{
 				Service:           viper.GetString("gcp_cloud_profiling.service"),
diff --git a/pkg/log/log.go b/pkg/log/log.go
@@ -17,6 +17,7 @@ package log
 
 import (
 	"context"
+	"fmt"
 	"log"
 
 	"github.com/go-chi/chi/middleware"
@@ -27,11 +28,13 @@ import (
 // Logger set the default logger to development mode
 var Logger *zap.SugaredLogger
 
+var traceStringPrefix string
+
 func init() {
-	ConfigureLogger("dev")
+	ConfigureLogger("dev", "")
 }
 
-func ConfigureLogger(logType string) {
+func ConfigureLogger(logType, traceStrPrefix string) {
 	var cfg zap.Config
 	if logType == "prod" {
 		cfg = zap.NewProductionConfig()
@@ -51,6 +54,10 @@ func ConfigureLogger(logType string) {
 		log.Fatalln("createLogger", err)
 	}
 	Logger = logger.Sugar()
+
+	if traceStrPrefix != "" {
+		traceStringPrefix = traceStrPrefix
+	}
 }
 
 func encodeLevel() zapcore.LevelEncoder {
@@ -109,6 +116,9 @@ func ContextLogger(ctx context.Context) *zap.SugaredLogger {
 		if ctxRequestID, ok := ctx.Value(middleware.RequestIDKey).(string); ok {
 			requestID := operation{ctxRequestID}
 			proposedLogger = proposedLogger.With(zap.Object("operation", requestID))
+			if traceStringPrefix != "" {
+				proposedLogger = proposedLogger.With(zap.String("trace", fmt.Sprintf("%s/%s", traceStringPrefix, ctxRequestID)))
+			}
 		}
 	}
 	return proposedLogger
