Repository: gomods/athens
Pull request #1655 (2020-07-27 03:09:57)

pkg/requestid: add Request ID to Athens requests and logs
This PR adds a new `Athens-Request-ID` header that if present would propagate throughout logs and requests. And if not present, a new request id will be created. 

Fixes https://github.com/gomods/athens/issues/1654


Patch:
diff --git a/cmd/proxy/actions/app.go b/cmd/proxy/actions/app.go
index 22ffdbf8..a0c0de8a 100644
--- a/cmd/proxy/actions/app.go
+++ b/cmd/proxy/actions/app.go
@@ -51,6 +51,7 @@ func App(conf *config.Config) (http.Handler, error) {
 	lggr := log.New(conf.CloudRuntime, logLvl)
 
 	r := mux.NewRouter()
+	r.Use(mw.WithRequestID)
 	r.Use(mw.LogEntryMiddleware(lggr))
 	r.Use(mw.RequestLogger)
 	r.Use(secure.New(secure.Options{
diff --git a/pkg/middleware/log_entry.go b/pkg/middleware/log_entry.go
index 4d2acee6..88b084a9 100644
--- a/pkg/middleware/log_entry.go
+++ b/pkg/middleware/log_entry.go
@@ -4,6 +4,7 @@ import (
 	"net/http"
 
 	"github.com/gomods/athens/pkg/log"
+	"github.com/gomods/athens/pkg/requestid"
 	"github.com/gorilla/mux"
 	"github.com/sirupsen/logrus"
 )
@@ -13,12 +14,13 @@ import (
 func LogEntryMiddleware(lggr *log.Logger) mux.MiddlewareFunc {
 	return func(h http.Handler) http.Handler {
 		f := func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
 			ent := lggr.WithFields(logrus.Fields{
 				"http-method": r.Method,
 				"http-path":   r.URL.Path,
+				"request-id":  requestid.FromContext(ctx),
 			})
-
-			ctx := log.SetEntryInContext(r.Context(), ent)
+			ctx = log.SetEntryInContext(ctx, ent)
 			r = r.WithContext(ctx)
 			h.ServeHTTP(w, r)
 		}
diff --git a/pkg/middleware/requestid.go b/pkg/middleware/requestid.go
new file mode 100644
index 00000000..d4850f94
--- /dev/null
+++ b/pkg/middleware/requestid.go
@@ -0,0 +1,23 @@
+package middleware
+
+import (
+	"net/http"
+
+	"github.com/gomods/athens/pkg/requestid"
+	"github.com/google/uuid"
+)
+
+// WithRequestID ensures a request id is in the
+// request context by either the incoming header
+// or creating a new one
+func WithRequestID(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requestID := r.Header.Get(requestid.HeaderKey)
+		if requestID == "" {
+			requestID = uuid.New().String()
+		}
+		ctx := requestid.SetInContext(r.Context(), requestID)
+		r = r.WithContext(ctx)
+		h.ServeHTTP(w, r)
+	})
+}
diff --git a/pkg/requestid/requestid.go b/pkg/requestid/requestid.go
new file mode 100644
index 00000000..fc3a8b92
--- /dev/null
+++ b/pkg/requestid/requestid.go
@@ -0,0 +1,21 @@
+package requestid
+
+import "context"
+
+// HeaderKey is the header key that athens uses
+// to pass request ids into logs and outbound requests
+const HeaderKey = "Athens-Request-ID"
+
+type key struct{}
+
+// SetInContext sets the given requestID into the context
+func SetInContext(ctx context.Context, id string) context.Context {
+	return context.WithValue(ctx, key{}, id)
+}
+
+// FromContext returns a requestID from the context or an empty
+// string if not found
+func FromContext(ctx context.Context) string {
+	id, _ := ctx.Value(key{}).(string)
+	return id
+}


Test patch:
diff --git a/pkg/middleware/log_entry_test.go b/pkg/middleware/log_entry_test.go
index 14e88d8b..0f6a76dd 100644
--- a/pkg/middleware/log_entry_test.go
+++ b/pkg/middleware/log_entry_test.go
@@ -34,6 +34,6 @@ func TestLogContext(t *testing.T) {
 	req, _ := http.NewRequest("GET", "/test", nil)
 	r.ServeHTTP(w, req)
 
-	expected := `{"http-method":"GET","http-path":"/test","level":"info","msg":"test"}`
+	expected := `{"http-method":"GET","http-path":"/test","level":"info","msg":"test","request-id":""}`
 	assert.True(t, strings.Contains(buf.String(), expected), fmt.Sprintf("%s should contain: %s", buf.String(), expected))
 }
diff --git a/pkg/middleware/requestid_test.go b/pkg/middleware/requestid_test.go
new file mode 100644
index 00000000..4cbe53ba
--- /dev/null
+++ b/pkg/middleware/requestid_test.go
@@ -0,0 +1,31 @@
+package middleware
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/gomods/athens/pkg/requestid"
+	"github.com/google/uuid"
+)
+
+func TestWithRequestID(t *testing.T) {
+	var givenRequestID string
+	h := WithRequestID(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		givenRequestID = requestid.FromContext(r.Context())
+	}))
+	req := httptest.NewRequest("GET", "/", nil)
+	expectedRequestID := uuid.New().String()
+	req.Header.Set(requestid.HeaderKey, expectedRequestID)
+	w := httptest.NewRecorder()
+	h.ServeHTTP(w, req)
+	if givenRequestID != expectedRequestID {
+		t.Fatalf("expected request id to be %q but got %q", expectedRequestID, givenRequestID)
+	}
+	req = httptest.NewRequest("GET", "/", nil)
+	w = httptest.NewRecorder()
+	h.ServeHTTP(w, req)
+	if givenRequestID == "" {
+		t.Fatal("expected a request id to be created when a request id header is empty")
+	}
+}