add traffic sniffing for handler
This commit is contained in:
+79
-15
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user