zeromicro/go-zero
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor(rest): keep rest log collector context key private (#3407)

Browse Source
This commit is contained in:
cong
2023-07-10 09:52:26 +08:00
committed by GitHub
parent b71453985c
commit 61e562d0c7
4 changed files with 39 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
rest/handler/loghandler.go
View File

@@ -3,7 +3,6 @@ package handler
import (
"bufio"
"bytes"
"context"
"errors"
"fmt"
"io"
@@ -44,7 +43,7 @@ func LogHandler(next http.Handler) http.Handler {
var dup io.ReadCloser
r.Body, dup = iox.DupReadCloser(r.Body)
next.ServeHTTP(&lrw, r.WithContext(context.WithValue(r.Context(), internal.LogContext, logs)))
next.ServeHTTP(&lrw, r.WithContext(internal.WithLogCollector(r.Context(), logs)))
r.Body = dup
logBrief(r, lrw.Code, timer, logs)
})
@@ -102,7 +101,7 @@ func DetailedLogHandler(next http.Handler) http.Handler {
var dup io.ReadCloser
r.Body, dup = iox.DupReadCloser(r.Body)
logs := new(internal.LogCollector)
next.ServeHTTP(lrw, r.WithContext(context.WithValue(r.Context(), internal.LogContext, logs)))
next.ServeHTTP(lrw, r.WithContext(internal.WithLogCollector(r.Context(), logs)))
r.Body = dup
logDetails(r, lrw, timer, logs)
})

4
rest/handler/loghandler_test.go
View File

@@ -22,7 +22,7 @@ func TestLogHandler(t *testing.T) {
for _, logHandler := range handlers {
req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
handler := logHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
r.Context().Value(internal.LogContext).(*internal.LogCollector).Append("anything")
internal.LogCollectorFromContext(r.Context()).Append("anything")
w.Header().Set("X-Test", "test")
w.WriteHeader(http.StatusServiceUnavailable)
_, err := w.Write([]byte("content"))
@@ -49,7 +49,7 @@ func TestLogHandlerVeryLong(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, "http://localhost", &buf)
handler := LogHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
r.Context().Value(internal.LogContext).(*internal.LogCollector).Append("anything")
internal.LogCollectorFromContext(r.Context()).Append("anything")
_, _ = io.Copy(io.Discard, r.Body)
w.Header().Set("X-Test", "test")
w.WriteHeader(http.StatusServiceUnavailable)