feat: log 404 requests with traceid (#1554)
This commit is contained in:
@@ -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
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
Reference in New Issue
Block a user