add traffic sniffing for handler

This commit is contained in:
ginuerzh
2024-09-27 15:55:28 +08:00
parent dfb6cb95d0
commit 3c8add4b82
43 changed files with 4518 additions and 1839 deletions
+79 -15
View File
@@ -9,6 +9,8 @@ import (
"github.com/go-gost/core/handler"
md "github.com/go-gost/core/metadata"
"github.com/go-gost/core/recorder"
xrecorder "github.com/go-gost/x/recorder"
"github.com/go-gost/x/registry"
)
@@ -17,11 +19,12 @@ func init() {
}
type fileHandler struct {
handler http.Handler
server *http.Server
ln *singleConnListener
md metadata
options handler.Options
handler http.Handler
server *http.Server
ln *singleConnListener
md metadata
options handler.Options
recorder recorder.RecorderObject
}
func NewHandler(opts ...handler.Option) handler.Handler {
@@ -45,6 +48,13 @@ func (h *fileHandler) Init(md md.Metadata) (err error) {
Handler: http.HandlerFunc(h.handleFunc),
}
for _, ro := range h.options.Recorders {
if ro.Record == xrecorder.RecorderServiceHandler {
h.recorder = ro
break
}
}
h.ln = &singleConnListener{
conn: make(chan net.Conn),
done: make(chan struct{}),
@@ -70,8 +80,54 @@ func (h *fileHandler) Close() error {
}
func (h *fileHandler) handleFunc(w http.ResponseWriter, r *http.Request) {
start := time.Now()
ro := &xrecorder.HandlerRecorderObject{
Service: h.options.Service,
RemoteAddr: r.RemoteAddr,
Network: "tcp",
Host: r.Host,
Proto: "http",
HTTP: &xrecorder.HTTPRecorderObject{
Host: r.Host,
Method: r.Method,
Proto: r.Proto,
Scheme: r.URL.Scheme,
URI: r.RequestURI,
Request: xrecorder.HTTPRequestRecorderObject{
ContentLength: r.ContentLength,
Header: r.Header,
},
},
Time: start,
}
ro.ClientIP, _, _ = net.SplitHostPort(r.RemoteAddr)
log := h.options.Logger.WithFields(map[string]any{
"remote": r.RemoteAddr,
})
rw := &responseWriter{ResponseWriter: w, statusCode: http.StatusOK}
defer func() {
ro.Duration = time.Since(start)
ro.HTTP.StatusCode = rw.statusCode
ro.HTTP.Response = xrecorder.HTTPResponseRecorderObject{
ContentLength: rw.contentLength,
Header: rw.Header(),
}
if err := ro.Record(context.Background(), h.recorder.Recorder); err != nil {
log.Errorf("record: %v", err)
}
log.WithFields(map[string]any{
"duration": time.Since(start),
}).Infof("%s %s %d %d", r.Method, r.RequestURI, rw.statusCode, rw.contentLength)
}()
if auther := h.options.Auther; auther != nil {
u, p, _ := r.BasicAuth()
ro.ClientID = u
if _, ok := auther.Authenticate(r.Context(), u, p); !ok {
w.Header().Set("WWW-Authenticate", "Basic")
w.WriteHeader(http.StatusUnauthorized)
@@ -79,16 +135,7 @@ func (h *fileHandler) handleFunc(w http.ResponseWriter, r *http.Request) {
}
}
log := h.options.Logger
start := time.Now()
h.handler.ServeHTTP(w, r)
log = log.WithFields(map[string]any{
"remote": r.RemoteAddr,
"duration": time.Since(start),
})
log.Infof("%s %s", r.Method, r.RequestURI)
h.handler.ServeHTTP(rw, r)
}
type singleConnListener struct {
@@ -132,3 +179,20 @@ func (l *singleConnListener) send(conn net.Conn) {
return
}
}
type responseWriter struct {
http.ResponseWriter
statusCode int
contentLength int64
}
func (w *responseWriter) Write(p []byte) (int, error) {
n, err := w.ResponseWriter.Write(p)
w.contentLength += int64(n)
return n, err
}
func (w *responseWriter) WriteHeader(statusCode int) {
w.statusCode = statusCode
w.ResponseWriter.WriteHeader(statusCode)
}