From afe431c94d732397bbe132ddd5c725d2bdd95cfe Mon Sep 17 00:00:00 2001
From: Dan Sosedoff <dan.sosedoff@gmail.com>
Date: Fri, 16 Dec 2022 12:38:57 -0600
Subject: [PATCH] Request logging additions (request id, forwarded user) (#618)

* Add request-id logging
* Missing test file
* Add option to log forwarded user information if available via X-Forwarded-Header
* Format
---
 pkg/api/logger.go      | 24 ++++++++++++++++++++++++
 pkg/api/logger_test.go | 31 +++++++++++++++++++++++++++++++
 pkg/command/options.go |  1 +
 3 files changed, 56 insertions(+)
 create mode 100644 pkg/api/logger_test.go

diff --git a/pkg/api/logger.go b/pkg/api/logger.go
index 230d45b..1441a18 100644
--- a/pkg/api/logger.go
+++ b/pkg/api/logger.go
@@ -8,6 +8,8 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/sirupsen/logrus"
+
+	"github.com/sosedoff/pgweb/pkg/command"
 )
 
 var (
@@ -29,6 +31,7 @@ func SetLogger(l *logrus.Logger) {
 
 func RequestLogger(logger *logrus.Logger) gin.HandlerFunc {
 	debug := logger.Level > logrus.InfoLevel
+	logForwardedUser := command.Opts.LogForwardedUser
 
 	return func(c *gin.Context) {
 		start := time.Now()
@@ -58,6 +61,19 @@ func RequestLogger(logger *logrus.Logger) gin.HandlerFunc {
 			"path":        path,
 		}
 
+		if reqID := getRequestID(c); reqID != "" {
+			fields["id"] = reqID
+		}
+
+		if logForwardedUser {
+			if forwardedUser := c.GetHeader("X-Forwarded-User"); forwardedUser != "" {
+				fields["forwarded_user"] = forwardedUser
+			}
+			if forwardedEmail := c.GetHeader("X-Forwarded-Email"); forwardedEmail != "" {
+				fields["forwarded_email"] = forwardedEmail
+			}
+		}
+
 		if err := c.Errors.Last(); err != nil {
 			fields["error"] = err.Error()
 		}
@@ -88,3 +104,11 @@ func RequestLogger(logger *logrus.Logger) gin.HandlerFunc {
 func sanitizeLogPath(str string) string {
 	return reConnectToken.ReplaceAllString(str, "/connect/REDACTED")
 }
+
+func getRequestID(c *gin.Context) string {
+	id := c.GetHeader("x-request-id")
+	if id == "" {
+		id = c.GetHeader("x-amzn-trace-id")
+	}
+	return id
+}
diff --git a/pkg/api/logger_test.go b/pkg/api/logger_test.go
new file mode 100644
index 0000000..18566e8
--- /dev/null
+++ b/pkg/api/logger_test.go
@@ -0,0 +1,31 @@
+package api
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_getRequestID(t *testing.T) {
+	examples := []struct {
+		headers map[string]string
+		result  string
+	}{
+		{map[string]string{}, ""},
+		{map[string]string{"X-Request-ID": "foo"}, "foo"},
+		{map[string]string{"x-request-id": "foo"}, "foo"},
+		{map[string]string{"x-request-id": "foo"}, "foo"},
+		{map[string]string{"x-request-id": "foo", "x-amzn-trace-id": "amz"}, "foo"},
+	}
+
+	for _, ex := range examples {
+		req := &http.Request{Header: http.Header{}}
+		for k, v := range ex.headers {
+			req.Header.Set(k, v)
+		}
+
+		assert.Equal(t, ex.result, getRequestID(&gin.Context{Request: req}))
+	}
+}
diff --git a/pkg/command/options.go b/pkg/command/options.go
index ac94248..0bdf19c 100644
--- a/pkg/command/options.go
+++ b/pkg/command/options.go
@@ -23,6 +23,7 @@ type Options struct {
 	Debug                        bool   `short:"d" long:"debug" description:"Enable debugging mode"`
 	LogLevel                     string `long:"log-level" description:"Logging level" default:"info"`
 	LogFormat                    string `long:"log-format" description:"Logging output format" default:"text"`
+	LogForwardedUser             bool   `long:"log-forwarded-user" description:"Log user information available in X-Forwarded-User/Email headers"`
 	URL                          string `long:"url" description:"Database connection string"`
 	Host                         string `long:"host" description:"Server hostname or IP" default:"localhost"`
 	Port                         int    `long:"port" description:"Server port" default:"5432"`