feat: log 404 requests with traceid (#1554)

This commit is contained in:
Kevin Wan
2022-02-19 20:50:33 +08:00
committed by GitHub
parent aa29036cb3
commit 842656aa90
16 changed files with 279 additions and 188 deletions

View File

@@ -1,42 +0,0 @@
package security
import (
"bufio"
"net"
"net/http"
)
// A WithCodeResponseWriter is a helper to delay sealing a http.ResponseWriter on writing code.
type WithCodeResponseWriter struct {
Writer http.ResponseWriter
Code int
}
// Flush flushes the response writer.
func (w *WithCodeResponseWriter) Flush() {
if flusher, ok := w.Writer.(http.Flusher); ok {
flusher.Flush()
}
}
// Header returns the http header.
func (w *WithCodeResponseWriter) Header() http.Header {
return w.Writer.Header()
}
// Hijack implements the http.Hijacker interface.
// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
func (w *WithCodeResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
return w.Writer.(http.Hijacker).Hijack()
}
// Write writes bytes into w.
func (w *WithCodeResponseWriter) Write(bytes []byte) (int, error) {
return w.Writer.Write(bytes)
}
// WriteHeader writes code into w, and not sealing the writer.
func (w *WithCodeResponseWriter) WriteHeader(code int) {
w.Writer.WriteHeader(code)
w.Code = code
}

View File

@@ -1,33 +0,0 @@
package security
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/stretchr/testify/assert"
)
func TestWithCodeResponseWriter(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
cw := &WithCodeResponseWriter{Writer: w}
cw.Header().Set("X-Test", "test")
cw.WriteHeader(http.StatusServiceUnavailable)
assert.Equal(t, cw.Code, http.StatusServiceUnavailable)
_, err := cw.Write([]byte("content"))
assert.Nil(t, err)
flusher, ok := http.ResponseWriter(cw).(http.Flusher)
assert.True(t, ok)
flusher.Flush()
})
resp := httptest.NewRecorder()
handler.ServeHTTP(resp, req)
assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
assert.Equal(t, "test", resp.Header().Get("X-Test"))
assert.Equal(t, "content", resp.Body.String())
}