Repository: rs/zerolog
Pull request #197 (2019-11-01 17:31:31)

Proposal for hlog, add custom header middleware
When using L7 loadbalancer like nginx, It is common to add some headers to the request like "X-Request-ID". It will be useful if we can utilize this value. 
This would cover more general use cases.

Thanks for the great library.


Patch:
diff --git a/array.go b/array.go
index aa4f623..d88c497 100644
--- a/array.go
+++ b/array.go
@@ -85,6 +85,12 @@ func (a *Array) Hex(val []byte) *Array {
 	return a
 }
 
+// RawJSON adds already encoded JSON to the array.
+func (a *Array) RawJSON(val []byte) *Array {
+	a.buf = appendJSON(enc.AppendArrayDelim(a.buf), val)
+	return a
+}
+
 // Err serializes and appends the err to the array.
 func (a *Array) Err(err error) *Array {
 	marshaled := ErrorMarshalFunc(err)
diff --git a/hlog/hlog.go b/hlog/hlog.go
index f7c3a5a..37f4336 100644
--- a/hlog/hlog.go
+++ b/hlog/hlog.go
@@ -171,6 +171,22 @@ func RequestIDHandler(fieldKey, headerName string) func(next http.Handler) http.
 	}
 }
 
+// CustomHeaderHandler adds given header from request's header as a field to
+// the context's logger using fieldKey as field key.
+func CustomHeaderHandler(fieldKey, header string) func(next http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if val := r.Header.Get(header); val != "" {
+				log := zerolog.Ctx(r.Context())
+				log.UpdateContext(func(c zerolog.Context) zerolog.Context {
+					return c.Str(fieldKey, val)
+				})
+			}
+			next.ServeHTTP(w, r)
+		})
+	}
+}
+
 // AccessHandler returns a handler that call f after each request.
 func AccessHandler(f func(r *http.Request, status, size int, duration time.Duration)) func(next http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {


Test patch:
diff --git a/array_test.go b/array_test.go
index 66c6668..e8857e6 100644
--- a/array_test.go
+++ b/array_test.go
@@ -24,10 +24,11 @@ func TestArray(t *testing.T) {
 		Str("a").
 		Bytes([]byte("b")).
 		Hex([]byte{0x1f}).
+		RawJSON([]byte(`{"some":"json"}`)).
 		Time(time.Time{}).
 		IPAddr(net.IP{192, 168, 0, 10}).
 		Dur(0)
-	want := `[true,1,2,3,4,5,6,7,8,9,10,11.98122,12.987654321,"a","b","1f","0001-01-01T00:00:00Z","192.168.0.10",0]`
+	want := `[true,1,2,3,4,5,6,7,8,9,10,11.98122,12.987654321,"a","b","1f",{"some":"json"},"0001-01-01T00:00:00Z","192.168.0.10",0]`
 	if got := decodeObjectToStr(a.write([]byte{})); got != want {
 		t.Errorf("Array.write()\ngot:  %s\nwant: %s", got, want)
 	}
diff --git a/hlog/hlog_test.go b/hlog/hlog_test.go
index 975ce38..2967eb2 100644
--- a/hlog/hlog_test.go
+++ b/hlog/hlog_test.go
@@ -182,6 +182,24 @@ func TestRequestIDHandler(t *testing.T) {
 	h.ServeHTTP(httptest.NewRecorder(), r)
 }
 
+func TestCustomHeaderHandler(t *testing.T) {
+	out := &bytes.Buffer{}
+	r := &http.Request{
+		Header: http.Header{
+			"X-Request-Id": []string{"514bbe5bb5251c92bd07a9846f4a1ab6"},
+		},
+	}
+	h := CustomHeaderHandler("reqID", "X-Request-Id")(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		l := FromRequest(r)
+		l.Log().Msg("")
+	}))
+	h = NewHandler(zerolog.New(out))(h)
+	h.ServeHTTP(nil, r)
+	if want, got := `{"reqID":"514bbe5bb5251c92bd07a9846f4a1ab6"}`+"\n", decodeIfBinary(out); want != got {
+		t.Errorf("Invalid log output, got: %s, want: %s", got, want)
+	}
+}
+
 func TestCombinedHandlers(t *testing.T) {
 	out := &bytes.Buffer{}
 	r := &http.Request{