增加 replayer 中间件

didi · Aug 16, 2021 · 3d9fbde · 3d9fbde
1 parent f3c1350
commit 3d9fbde
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/replayer/middleware.go b/replayer/middleware.go
@@ -0,0 +1,16 @@
+package replayer
+
+import (
+	"context"
+	"net/http"
+)
+
+func AddHeader(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		const traceIDKey = "Sharingan-Replayer-TraceID"
+		if s := r.Header.Get(traceIDKey); s != "" {
+			r = r.WithContext(context.WithValue(r.Context(), traceIDKey, s))
+		}
+		next.ServeHTTP(w, r)
+	})
+}