Tap/raw recorder

handler/forward/local/conn.go

@@ -0,0 +1,52 @@
+package local
+
+import (
+	"bytes"
+	"context"
+	"encoding/hex"
+	"net"
+	"time"
+
+	"github.com/go-gost/core/recorder"
+)
+
+type recorderConn struct {
+	net.Conn
+	recorder recorder.RecorderObject
+}
+
+func (c *recorderConn) record(ctx context.Context, direction byte, b []byte) {
+	if len(b) == 0 || c.recorder.Recorder == nil {
+		return
+	}
+
+	var buf bytes.Buffer
+	if c.recorder.Options != nil && c.recorder.Options.Direction {
+		buf.WriteByte(direction)
+	}
+	if c.recorder.Options != nil && c.recorder.Options.TimestampFormat != "" {
+		buf.WriteString(time.Now().Format(c.recorder.Options.TimestampFormat))
+	}
+	if buf.Len() > 0 {
+		buf.WriteByte('\n')
+	}
+
+	if c.recorder.Options != nil && c.recorder.Options.Hexdump {
+		buf.WriteString(hex.Dump(b))
+	} else {
+		buf.Write(b)
+	}
+
+	_ = c.recorder.Recorder.Record(ctx, buf.Bytes())
+}
+
+func (c *recorderConn) Read(b []byte) (n int, err error) {
+	n, err = c.Conn.Read(b)
+	c.record(context.Background(), '>', b[:n])
+	return
+}
+
+func (c *recorderConn) Write(b []byte) (int, error) {
+	c.record(context.Background(), '<', b)
+	return c.Conn.Write(b)
+}

handler/forward/local/conn_test.go

@@ -0,0 +1,81 @@
+package local
+
+import (
+	"context"
+	"errors"
+	"net"
+	"testing"
+
+	"github.com/go-gost/core/recorder"
+)
+
+type captureRecorder struct {
+	records [][]byte
+	err     error
+}
+
+func (r *captureRecorder) Record(_ context.Context, b []byte, _ ...recorder.RecordOption) error {
+	r.records = append(r.records, append([]byte(nil), b...))
+	return r.err
+}
+
+func TestRecorderConnRecordsPayload(t *testing.T) {
+	client, upstream := net.Pipe()
+	defer client.Close()
+	defer upstream.Close()
+
+	rec := &captureRecorder{}
+	conn := &recorderConn{
+		Conn: client,
+		recorder: recorder.RecorderObject{
+			Recorder: rec,
+			Options:  &recorder.Options{Direction: true},
+		},
+	}
+
+	go func() {
+		buf := make([]byte, 7)
+		_, _ = upstream.Read(buf)
+		_, _ = upstream.Write([]byte("reply"))
+	}()
+
+	if _, err := conn.Write([]byte("payload")); err != nil {
+		t.Fatal(err)
+	}
+	buf := make([]byte, 5)
+	if _, err := conn.Read(buf); err != nil {
+		t.Fatal(err)
+	}
+
+	if got := string(rec.records[0]); got != "<\npayload" {
+		t.Fatalf("write record = %q, want %q", got, "<\npayload")
+	}
+	if got := string(rec.records[1]); got != ">\nreply" {
+		t.Fatalf("read record = %q, want %q", got, ">\nreply")
+	}
+}
+
+func TestRecorderConnRecordFailureDoesNotStopWrite(t *testing.T) {
+	client, upstream := net.Pipe()
+	defer client.Close()
+	defer upstream.Close()
+
+	conn := &recorderConn{
+		Conn: client,
+		recorder: recorder.RecorderObject{
+			Recorder: &captureRecorder{err: errors.New("sink down")},
+		},
+	}
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		buf := make([]byte, 7)
+		_, _ = upstream.Read(buf)
+	}()
+
+	if _, err := conn.Write([]byte("payload")); err != nil {
+		t.Fatal(err)
+	}
+	<-done
+}

handler/forward/local/handler.go

@@ -35,11 +35,12 @@ func init() {
 }
 
 type forwardHandler struct {
-	hop      hop.Hop
-	md       metadata
-	options  handler.Options
-	recorder recorder.RecorderObject
-	certPool tls_util.CertPool
+	hop         hop.Hop
+	md          metadata
+	options     handler.Options
+	recorder    recorder.RecorderObject
+	rawRecorder recorder.RecorderObject
+	certPool    tls_util.CertPool
 }
 
 func NewHandler(opts ...handler.Option) handler.Handler {
@@ -59,9 +60,11 @@ func (h *forwardHandler) Init(md md.Metadata) (err error) {
 	}
 
 	for _, ro := range h.options.Recorders {
-		if ro.Record == xrecorder.RecorderServiceHandler {
+		switch ro.Record {
+		case xrecorder.RecorderServiceHandler:
 			h.recorder = ro
-			break
+		case xrecorder.RecorderServiceHandlerRaw:
+			h.rawRecorder = ro
 		}
 	}
 
@@ -110,6 +113,13 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
 	})
 	log.Infof("%s <> %s", conn.RemoteAddr(), conn.LocalAddr())
 
+	if network == "tcp" && h.rawRecorder.Recorder != nil {
+		conn = &recorderConn{
+			Conn:     conn,
+			recorder: h.rawRecorder,
+		}
+	}
+
 	pStats := xstats.Stats{}
 	conn = stats_wrapper.WrapConn(conn, &pStats)
 

recorder/recorder.go

@@ -13,6 +13,7 @@ const (
 	RecorderServiceHandler       = "recorder.service.handler"
 	RecorderServiceHandlerSerial = "recorder.service.handler.serial"
 	RecorderServiceHandlerTunnel = "recorder.service.handler.tunnel"
+	RecorderServiceHandlerRaw    = "recorder.service.handler.raw"
 )
 
 type HTTPRequestRecorderObject struct {