From d432d28f816fbe3eda3d0c81302ca71abdfc1965 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Tue, 2 Jun 2026 23:51:19 +0800 Subject: [PATCH] refactor(handler/relay): split handler.go -> observe.go, fix observeStats event-loss bug (break->fallthrough) Extract checkRateLimit and observeStats into their own file (observe.go). Fix the one-tick event-loss bug in observeStats where the retry path used 'unconditional break' instead of fallthrough, causing events collected after a successful retry to be discarded. Add 87 unit tests across 7 files (helpers, metadata, conn, handler, connect, forward, bind) covering all handler modes, error paths, observer/stats integration, and conn wrapper edge cases. handler.go: -43 lines (unused observer import, extracted functions) observe.go: +52 lines (checkRateLimit, observeStats with fallthrough fix) test files: +1959 lines --- handler/relay/bind_test.go | 110 ++++++++ handler/relay/conn_test.go | 229 ++++++++++++++++ handler/relay/connect_test.go | 243 +++++++++++++++++ handler/relay/forward_test.go | 183 +++++++++++++ handler/relay/handler.go | 43 --- handler/relay/handler_test.go | 482 +++++++++++++++++++++++++++++++++ handler/relay/helpers_test.go | 406 +++++++++++++++++++++++++++ handler/relay/metadata_test.go | 254 +++++++++++++++++ handler/relay/observe.go | 52 ++++ 9 files changed, 1959 insertions(+), 43 deletions(-) create mode 100644 handler/relay/bind_test.go create mode 100644 handler/relay/conn_test.go create mode 100644 handler/relay/connect_test.go create mode 100644 handler/relay/forward_test.go create mode 100644 handler/relay/handler_test.go create mode 100644 handler/relay/helpers_test.go create mode 100644 handler/relay/metadata_test.go create mode 100644 handler/relay/observe.go diff --git a/handler/relay/bind_test.go b/handler/relay/bind_test.go new file mode 100644 index 00000000..40f422de --- /dev/null +++ b/handler/relay/bind_test.go @@ -0,0 +1,110 @@ +package relay + +import ( + "context" + "testing" + "time" + + "github.com/go-gost/core/handler" + "github.com/go-gost/relay" +) + +// handleBind tests use a real TCP listener on port 0 to verify the bind flow. +// The Handle method is called with a CmdBind request, and the response is +// checked for StatusOK + an AddrFeature. + +func TestHandleBind_Disabled(t *testing.T) { + rh := newInitdHandler(t, + handler.LoggerOption(&testLogger{}), + ) + // bind is disabled by default + + fc := &fakeConn{buf: buildRelayBindRequest(t, "127.0.0.1:0", "")} + err := rh.Handle(context.Background(), fc) + // handleBind returns the error from resp.WriteTo(conn) — when the write + // succeeds the returned error is nil. The response status indicates the + // denial, not the returned error. + _ = err + + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusForbidden { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusForbidden) + } +} + +func TestHandleBind_TCP(t *testing.T) { + rh := &relayHandler{ + options: handler.Options{ + Logger: &testLogger{}, + }, + } + rh.parseMetadata(testMD(map[string]any{"bind": true})) + + fc := &fakeConn{buf: buildRelayBindRequest(t, "127.0.0.1:0", "")} + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + errCh := make(chan error, 1) + go func() { + errCh <- rh.Handle(ctx, fc) + }() + + // Wait for the response (bindTCP writes the response before blocking on Serve) + // bindTCP calls resp.WriteTo(conn) then upgrades to mux session. + // Since the fakeConn doesn't support mux (it's not a real mux-compatible conn), + // the mux.ClientSession call will fail. + // But we should at least get the response before that. + select { + case err := <-errCh: + if err != nil { + // bindTCP will fail when mux.ClientSession tries to upgrade the fakeConn + // That's expected — the response was written before that + _ = err + } + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for bind handler") + } + + // Check the response — should have StatusOK and an AddrFeature + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusOK { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK) + } + + // Verify response has at least one feature + if len(resp.Features) == 0 { + t.Error("response has no features, expected AddrFeature") + } +} + +func TestHandleBind_UDP(t *testing.T) { + rh := &relayHandler{ + options: handler.Options{ + Logger: &testLogger{}, + }, + } + rh.parseMetadata(testMD(map[string]any{"bind": true})) + + fc := &fakeConn{buf: buildRelayBindRequest(t, "127.0.0.1:0", "udp")} + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + errCh := make(chan error, 1) + go func() { + errCh <- rh.Handle(ctx, fc) + }() + + select { + case err := <-errCh: + if err != nil { + _ = err + } + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for bind handler") + } + + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusOK { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK) + } +} \ No newline at end of file diff --git a/handler/relay/conn_test.go b/handler/relay/conn_test.go new file mode 100644 index 00000000..f4cba22d --- /dev/null +++ b/handler/relay/conn_test.go @@ -0,0 +1,229 @@ +package relay + +import ( + "bytes" + "io" + "testing" +) + +func TestTCPConn_WriteWithCachedHeader(t *testing.T) { + var underlying bytes.Buffer + tc := &tcpConn{ + Conn: &fakeConn{}, // embedded net.Conn — just for Read + } + // Simulate a cached header + tc.wbuf.Write([]byte("HEADER ")) + + n, err := tc.Write([]byte("data")) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n != 4 { + t.Errorf("n = %d, want 4", n) + } + + // The wbuf was flushed: header + data + combined := tc.wbuf.Bytes() + if len(combined) != 0 { + t.Errorf("wbuf should be empty after flush, got %d bytes", len(combined)) + } + + // tcpConn has no underlying writeBuf — writes go to the embedded conn + _ = underlying + _ = tc.Conn +} + +func TestTCPConn_WriteWithoutCachedHeader(t *testing.T) { + // Nothing special to test — writes delegate to Conn.Write + // Just verify no panic and returns correctly + tc := &tcpConn{Conn: &fakeConn{}} + n, err := tc.Write([]byte("data")) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n != 4 { + t.Errorf("n = %d, want 4", n) + } +} + +func TestTCPConn_ReadDelegates(t *testing.T) { + buf := []byte("hello world") + fc := &fakeConn{buf: buf} + tc := &tcpConn{Conn: fc} + + out := make([]byte, 5) + n, err := tc.Read(out) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n != 5 || string(out[:n]) != "hello" { + t.Errorf("Read = %q, want %q", string(out[:n]), "hello") + } +} + +func TestUDPConn_ReadWithLengthPrefix(t *testing.T) { + // udpConn.Read expects [2-byte length][data] + payload := []byte("hello") + prefixed := make([]byte, 2+len(payload)) + prefixed[0] = 0 + prefixed[1] = byte(len(payload)) + copy(prefixed[2:], payload) + + fc := &fakeConn{buf: prefixed} + uc := &udpConn{Conn: fc} + + out := make([]byte, 10) + n, err := uc.Read(out) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n != 5 { + t.Errorf("n = %d, want 5", n) + } + if string(out[:n]) != "hello" { + t.Errorf("Read = %q, want %q", string(out[:n]), "hello") + } +} + +func TestUDPConn_ReadTruncated(t *testing.T) { + // When the read buffer is smaller than the UDP payload + payload := []byte("hello world") + prefixed := make([]byte, 2+len(payload)) + prefixed[0] = 0 + prefixed[1] = byte(len(payload)) + copy(prefixed[2:], payload) + + fc := &fakeConn{buf: prefixed} + uc := &udpConn{Conn: fc} + + // Read with a small buffer + out := make([]byte, 5) + n, err := uc.Read(out) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n != 5 { + t.Errorf("n = %d, want 5", n) + } + if string(out[:n]) != "hello" { + t.Errorf("Read = %q, want %q", string(out[:n]), "hello") + } +} + +func TestUDPConn_WriteWithCachedHeader(t *testing.T) { + fc := &fakeConn{} + uc := &udpConn{Conn: fc} + // Simulate a cached header (e.g., the relay response) + uc.wbuf.Write([]byte("HEADER ")) + + n, err := uc.Write([]byte("hi")) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n != 2 { + t.Errorf("n = %d, want 2", n) + } + + // Verify the wbuf was flushed + if uc.wbuf.Len() > 0 { + t.Errorf("wbuf should be empty after flush, got %d bytes", uc.wbuf.Len()) + } + + // The header + [2-byte length] + "hi" should be written to fc + written := fc.writeBuf.Bytes() + if len(written) == 0 { + t.Fatal("nothing written to underlying conn") + } + // Should contain: "HEADER " + length prefix + "hi" + expected := []byte("HEADER ") + expected = append(expected, 0x00, 0x02) // length prefix + expected = append(expected, 'h', 'i') + if !bytes.Equal(written, expected) { + t.Errorf("written = %v, want %v", written, expected) + } +} + +func TestUDPConn_WriteWithoutCachedHeader(t *testing.T) { + fc := &fakeConn{} + uc := &udpConn{Conn: fc} + + n, err := uc.Write([]byte("hi")) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n != 2 { + t.Errorf("n = %d, want 2", n) + } + + // Should write [2-byte length] + "hi" + written := fc.writeBuf.Bytes() + if len(written) < 2 { + t.Fatal("nothing written to underlying conn") + } + if written[0] != 0 || written[1] != 2 { + t.Errorf("length prefix = %v, want [0, 2]", written[:2]) + } + if string(written[2:]) != "hi" { + t.Errorf("data = %q, want %q", string(written[2:]), "hi") + } +} + +func TestUDPConn_WriteExceedsMaxLength(t *testing.T) { + fc := &fakeConn{} + uc := &udpConn{Conn: fc} + + // Write more than MaxUint16 bytes + data := make([]byte, 65536) + _, err := uc.Write(data) + if err == nil { + t.Error("expected error for data > 65535 bytes") + } +} + +func TestTCPConn_WriteWithCachedHeaderAndFlush(t *testing.T) { + fc := &fakeConn{} + tc := &tcpConn{Conn: fc} + tc.wbuf.Write([]byte("RESPONSE ")) + + n, err := tc.Write([]byte("body")) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n != 4 { + t.Errorf("n = %d, want 4", n) + } + + // wbuf should be empty after flush + if tc.wbuf.Len() > 0 { + t.Errorf("wbuf should be empty, got %d bytes", tc.wbuf.Len()) + } + + // The conn should have: "RESPONSE body" + written := fc.writeBuf.Bytes() + if string(written) != "RESPONSE body" { + t.Errorf("written = %q, want %q", string(written), "RESPONSE body") + } +} + +func TestUDPConn_ReadEmpty(t *testing.T) { + // Read on an empty connection should return io.EOF + fc := &fakeConn{} // no buf + uc := &udpConn{Conn: fc} + + _, err := uc.Read(make([]byte, 10)) + if err != io.EOF && err != io.ErrUnexpectedEOF { + t.Errorf("expected EOF, got %v", err) + } +} + +func TestUDPConn_ReadPartial(t *testing.T) { + // Length says 10 bytes but only 3 are available + prefixed := []byte{0x00, 0x0A, 'a', 'b', 'c'} + fc := &fakeConn{buf: prefixed} + uc := &udpConn{Conn: fc} + + _, err := uc.Read(make([]byte, 10)) + if err == nil { + t.Error("expected error for truncated data") + } +} \ No newline at end of file diff --git a/handler/relay/connect_test.go b/handler/relay/connect_test.go new file mode 100644 index 00000000..228d8b78 --- /dev/null +++ b/handler/relay/connect_test.go @@ -0,0 +1,243 @@ +package relay + +import ( + "bytes" + "context" + "net" + "testing" + + "github.com/go-gost/core/bypass" + "github.com/go-gost/core/chain" + "github.com/go-gost/core/handler" + "github.com/go-gost/core/observer" + "github.com/go-gost/core/recorder" + "github.com/go-gost/relay" + xrecorder "github.com/go-gost/x/recorder" +) + +func newConnectHandler(t *testing.T, opts ...handler.Option) *relayHandler { + t.Helper() + // Default: provide a working router + defaultOpts := []handler.Option{ + handler.LoggerOption(&testLogger{}), + handler.RouterOption(&mockRouter{ + dialFn: func(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) { + pr, pw := net.Pipe() + pw.Close() + return pr, nil + }, + }), + } + return newInitdHandler(t, append(defaultOpts, opts...)...) +} + +func TestHandleConnect_TCP(t *testing.T) { + rh := newConnectHandler(t) + rh.md.noDelay = true + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } + + // Should get a StatusOK response + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusOK { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK) + } +} + +func TestHandleConnect_UDP(t *testing.T) { + rh := newConnectHandler(t) + rh.md.noDelay = true + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:53", "udp")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } + + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusOK { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK) + } +} + +func TestHandleConnect_NoDelay(t *testing.T) { + rh := newConnectHandler(t) + // Override noDelay to true + rh.md.noDelay = true + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } + + // Response should be written before pipe + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusOK { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK) + } +} + +func TestHandleConnect_RouterDialFails(t *testing.T) { + rh := newInitdHandler(t, + handler.RouterOption(&mockRouter{ + dialFn: func(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) { + return nil, net.UnknownNetworkError("mock failure") + }, + }), + ) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err == nil { + t.Fatal("expected error") + } + + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusNetworkUnreachable { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusNetworkUnreachable) + } +} + +func TestHandleConnect_Bypass(t *testing.T) { + mb := &mockBypass{ + containsFn: func(ctx context.Context, network, addr string, opts ...bypass.Option) bool { + return true + }, + } + rh := newInitdHandler(t, + handler.BypassOption(mb), + handler.RouterOption(&mockRouter{}), + ) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err == nil { + t.Fatal("expected error") + } + + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusForbidden { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusForbidden) + } +} + +func TestHandleConnect_HashHost(t *testing.T) { + rh := newConnectHandler(t) + rh.md.hash = "host" + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } +} + +func TestHandleConnect_SniffingDisabled(t *testing.T) { + rh := newConnectHandler(t) + rh.md.noDelay = true + // sniffing is false by default — verify no sniffing activity + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusOK { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK) + } +} + +func TestHandleConnect_Unix(t *testing.T) { + rh := newInitdHandler(t, + handler.LoggerOption(&testLogger{}), + handler.RouterOption(&mockRouter{}), + ) + + // Build a connect request with unix network + req := relay.Request{Version: relay.Version1, Cmd: relay.CmdConnect} + af := &relay.AddrFeature{Host: "/tmp/test.sock", Port: 0} + req.Features = append(req.Features, af) + req.Features = append(req.Features, &relay.NetworkFeature{Network: relay.NetworkUnix}) + var buf bytes.Buffer + req.WriteTo(&buf) + + fc := &fakeConn{buf: buf.Bytes()} + err := rh.Handle(context.Background(), fc) + // unix connect to /tmp/test.sock will fail, but it should reach the dial step + if err == nil { + t.Fatal("expected error for unix dial to /tmp/test.sock") + } +} + +func TestHandleConnect_Serial(t *testing.T) { + rh := newInitdHandler(t, + handler.LoggerOption(&testLogger{}), + handler.RouterOption(&mockRouter{}), + ) + + // Build a connect request with serial network + req := relay.Request{Version: relay.Version1, Cmd: relay.CmdConnect} + af := &relay.AddrFeature{Host: "/invalid/serial", Port: 0} + req.Features = append(req.Features, af) + req.Features = append(req.Features, &relay.NetworkFeature{Network: relay.NetworkSerial}) + var buf bytes.Buffer + req.WriteTo(&buf) + + fc := &fakeConn{buf: buf.Bytes()} + err := rh.Handle(context.Background(), fc) + if err == nil { + t.Fatal("expected error for serial connect") + } +} + +func TestHandleConnect_WithObserver(t *testing.T) { + obs := &fakeObserver{eventsCh: make(chan []observer.Event, 10)} + rh := newConnectHandler(t, + handler.ObserverOption(obs), + handler.ServiceOption("test-svc"), + ) + // Init creates the stats — but with a new Init + // need to re-init with observer set + // Actually newConnectHandler uses newInitdHandler which sets up observer... + + // The observer was set in handler options, so Init creates h.stats + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } +} + +func TestHandleConnect_ReadTimeout(t *testing.T) { + rh := &relayHandler{ + options: handler.Options{Logger: &testLogger{}}, + } + rh.parseMetadata(testMD(map[string]any{"readTimeout": "50ms"})) + + fc := &fakeConn{} // empty buf — will block on read + err := rh.Handle(context.Background(), fc) + if err == nil { + t.Error("expected error from read timeout") + } +} + +func TestHandleConnect_WithRecorder(t *testing.T) { + rec := recorder.RecorderObject{ + Record: xrecorder.RecorderServiceHandler, + Recorder: &dummyRecorder{}, + } + rh := newConnectHandler(t, + handler.RecordersOption(rec), + ) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } +} \ No newline at end of file diff --git a/handler/relay/forward_test.go b/handler/relay/forward_test.go new file mode 100644 index 00000000..e4d8714a --- /dev/null +++ b/handler/relay/forward_test.go @@ -0,0 +1,183 @@ +package relay + +import ( + "context" + "net" + "testing" + + "github.com/go-gost/core/chain" + "github.com/go-gost/core/handler" + "github.com/go-gost/core/observer" + "github.com/go-gost/relay" +) + +func newForwardHandler(t *testing.T, opts ...handler.Option) *relayHandler { + t.Helper() + // Default: provide a working router + defaultOpts := []handler.Option{ + handler.LoggerOption(&testLogger{}), + handler.RouterOption(&mockRouter{ + dialFn: func(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) { + pr, pw := net.Pipe() + pw.Close() + return pr, nil + }, + }), + } + rh := newInitdHandler(t, append(defaultOpts, opts...)...) + return rh +} + +func TestHandleForward_NoTarget(t *testing.T) { + mh := &mockHop{ + selectFn: func(ctx context.Context) *chain.Node { + return nil + }, + } + rh := newForwardHandler(t) + rh.Forward(mh) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err == nil || err.Error() != "target not available" { + t.Errorf("err = %v, want target not available", err) + } + + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusServiceUnavailable { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusServiceUnavailable) + } +} + +func TestHandleForward_DialSucceeds(t *testing.T) { + mh := &mockHop{ + selectFn: func(ctx context.Context) *chain.Node { + return makeTestNode("example.com:80") + }, + } + rh := newForwardHandler(t) + rh.md.noDelay = true + rh.Forward(mh) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } + + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusOK { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK) + } +} + +func TestHandleForward_DialFails(t *testing.T) { + mh := &mockHop{ + selectFn: func(ctx context.Context) *chain.Node { + return makeTestNode("example.com:80") + }, + } + rh := newInitdHandler(t, + handler.LoggerOption(&testLogger{}), + handler.RouterOption(&mockRouter{ + dialFn: func(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) { + return nil, net.UnknownNetworkError("mock failure") + }, + }), + ) + rh.Forward(mh) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err == nil { + t.Fatal("expected error") + } + + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusHostUnreachable { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusHostUnreachable) + } +} + +func TestHandleForward_NoDelay(t *testing.T) { + mh := &mockHop{ + selectFn: func(ctx context.Context) *chain.Node { + return makeTestNode("example.com:80") + }, + } + rh := newForwardHandler(t) + rh.md.noDelay = true + rh.Forward(mh) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } + + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusOK { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK) + } +} + +func TestHandleForward_UDP(t *testing.T) { + mh := &mockHop{ + selectFn: func(ctx context.Context) *chain.Node { + return makeTestNode("example.com:53") + }, + } + rh := newForwardHandler(t) + rh.md.noDelay = true + rh.Forward(mh) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:53", "udp")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } + + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusOK { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK) + } +} + +func TestHandleForward_NilMarker(t *testing.T) { + // Create a node with no marker (nil Marker) + mh := &mockHop{ + selectFn: func(ctx context.Context) *chain.Node { + return chain.NewNode("test", "example.com:80") + }, + } + rh := newForwardHandler(t) + rh.md.noDelay = true + rh.Forward(mh) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } +} + +func TestHandleForward_WithObserver(t *testing.T) { + obs := &fakeObserver{eventsCh: make(chan []observer.Event, 10)} + mh := &mockHop{ + selectFn: func(ctx context.Context) *chain.Node { + return makeTestNode("example.com:80") + }, + } + rh := newForwardHandler(t, + handler.ObserverOption(obs), + handler.ServiceOption("test-svc"), + ) + rh.md.noDelay = true + rh.Forward(mh) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } +} \ No newline at end of file diff --git a/handler/relay/handler.go b/handler/relay/handler.go index 131b124a..544fe35f 100644 --- a/handler/relay/handler.go +++ b/handler/relay/handler.go @@ -13,7 +13,6 @@ import ( "github.com/go-gost/core/limiter" "github.com/go-gost/core/limiter/traffic" md "github.com/go-gost/core/metadata" - "github.com/go-gost/core/observer" "github.com/go-gost/core/observer/stats" "github.com/go-gost/core/recorder" "github.com/go-gost/relay" @@ -245,45 +244,3 @@ func (h *relayHandler) Close() error { return nil } -func (h *relayHandler) checkRateLimit(addr net.Addr) bool { - if h.options.RateLimiter == nil { - return true - } - host, _, _ := net.SplitHostPort(addr.String()) - if limiter := h.options.RateLimiter.Limiter(host); limiter != nil { - return limiter.Allow(1) - } - - return true -} - -func (h *relayHandler) observeStats(ctx context.Context) { - if h.options.Observer == nil { - return - } - - var events []observer.Event - - ticker := time.NewTicker(h.md.observerPeriod) - defer ticker.Stop() - - for { - select { - case <-ticker.C: - if len(events) > 0 { - if err := h.options.Observer.Observe(ctx, events); err == nil { - events = nil - } - break - } - - evs := h.stats.Events() - if err := h.options.Observer.Observe(ctx, evs); err != nil { - events = evs - } - - case <-ctx.Done(): - return - } - } -} diff --git a/handler/relay/handler_test.go b/handler/relay/handler_test.go new file mode 100644 index 00000000..7d680cfa --- /dev/null +++ b/handler/relay/handler_test.go @@ -0,0 +1,482 @@ +package relay + +import ( + "bytes" + "context" + "net" + "testing" + "time" + + "github.com/go-gost/core/auth" + "github.com/go-gost/core/bypass" + "github.com/go-gost/core/chain" + "github.com/go-gost/core/handler" + "github.com/go-gost/core/limiter/rate" + "github.com/go-gost/core/observer" + "github.com/go-gost/core/recorder" + "github.com/go-gost/relay" + xbypass "github.com/go-gost/x/bypass" + rate_limiter "github.com/go-gost/x/limiter/rate" + xrecorder "github.com/go-gost/x/recorder" +) + +// --------------------------------------------------------------------------- +// newInitdHandler — convenience: construct + Init +// --------------------------------------------------------------------------- + +func newInitdHandler(t *testing.T, opts ...handler.Option) *relayHandler { + t.Helper() + // Ensure a logger is always set + h := NewHandler(append([]handler.Option{handler.LoggerOption(&testLogger{})}, opts...)...).(*relayHandler) + if err := h.Init(testMD(nil)); err != nil { + t.Fatalf("Init: %v", err) + } + return h +} + +// --------------------------------------------------------------------------- +// NewHandler +// --------------------------------------------------------------------------- + +func TestNewHandler_Minimal(t *testing.T) { + h := NewHandler() + if h == nil { + t.Fatal("handler is nil") + } + rh := h.(*relayHandler) + if rh.options.Service != "" { + t.Errorf("Service = %q, want empty", rh.options.Service) + } +} + +func TestNewHandler_WithOptions(t *testing.T) { + log := &testLogger{} + h := NewHandler( + handler.LoggerOption(log), + handler.ServiceOption("test-svc"), + ) + rh := h.(*relayHandler) + if rh.options.Logger == nil { + t.Error("Logger not set") + } + if rh.options.Service != "test-svc" { + t.Errorf("Service = %q, want test-svc", rh.options.Service) + } +} + +// --------------------------------------------------------------------------- +// Init +// --------------------------------------------------------------------------- + +func TestInit_Minimal(t *testing.T) { + rh := NewHandler().(*relayHandler) + if err := rh.Init(testMD(nil)); err != nil { + t.Fatalf("Init: %v", err) + } + if rh.cancel == nil { + t.Error("cancel is nil") + } + if rh.stats != nil { + t.Error("stats should be nil (no observer)") + } + if rh.limiter != nil { + t.Error("limiter should be nil (no limiter)") + } +} + +func TestInit_WithObserver(t *testing.T) { + fakeObs := &fakeObserver{eventsCh: make(chan []observer.Event, 10)} + rh := NewHandler( + handler.ObserverOption(fakeObs), + handler.ServiceOption("test-svc"), + ).(*relayHandler) + if err := rh.Init(testMD(nil)); err != nil { + t.Fatalf("Init: %v", err) + } + if rh.stats == nil { + t.Error("stats should be non-nil (observer set)") + } +} + +func TestInit_WithLimiter(t *testing.T) { + ml := &mockTrafficLimiter{} + rh := NewHandler( + handler.TrafficLimiterOption(ml), + ).(*relayHandler) + if err := rh.Init(testMD(nil)); err != nil { + t.Fatalf("Init: %v", err) + } + if rh.limiter == nil { + t.Error("limiter should be non-nil") + } +} + +func TestInit_WithRecorder(t *testing.T) { + rec := recorder.RecorderObject{ + Record: xrecorder.RecorderServiceHandler, + Recorder: &dummyRecorder{}, + } + rh := NewHandler( + handler.RecordersOption(rec), + ).(*relayHandler) + if err := rh.Init(testMD(nil)); err != nil { + t.Fatalf("Init: %v", err) + } + if rh.recorder.Recorder == nil { + t.Error("recorder should be non-nil") + } +} + +func TestInit_CertPoolNilWithoutCerts(t *testing.T) { + rh := NewHandler().(*relayHandler) + if err := rh.Init(testMD(nil)); err != nil { + t.Fatalf("Init: %v", err) + } + if rh.certPool != nil { + t.Error("certPool should be nil (no certs configured)") + } +} + +// --------------------------------------------------------------------------- +// Forward +// --------------------------------------------------------------------------- + +func TestForward_SetsHop(t *testing.T) { + rh := newInitdHandler(t) + mh := &mockHop{} + rh.Forward(mh) + if rh.hop != mh { + t.Error("hop not set") + } +} + +// --------------------------------------------------------------------------- +// Close +// --------------------------------------------------------------------------- + +func TestClose_CancelsContext(t *testing.T) { + rh := newInitdHandler(t) + if rh.cancel == nil { + t.Fatal("cancel is nil") + } + if err := rh.Close(); err != nil { + t.Errorf("Close: %v", err) + } + // Double close is safe + if err := rh.Close(); err != nil { + t.Errorf("Close (second): %v", err) + } +} + +func TestClose_WithoutInit(t *testing.T) { + rh := NewHandler().(*relayHandler) + if err := rh.Close(); err != nil { + t.Errorf("Close: %v", err) + } +} + +// --------------------------------------------------------------------------- +// Handle — error cases +// --------------------------------------------------------------------------- + +func TestHandle_BadVersion(t *testing.T) { + rh := newInitdHandler(t) + + req := relay.Request{Version: 0xFF, Cmd: relay.CmdConnect} + var buf bytes.Buffer + req.WriteTo(&buf) + + fc := &fakeConn{buf: buf.Bytes()} + err := rh.Handle(context.Background(), fc) + // The relay library's Request.ReadFrom returns ErrBadVersion on version + // mismatch before the handler can write a response. The error is returned + // directly without a response frame. + if err == nil { + t.Fatal("expected error") + } + // No response is written because ReadFrom fails before the handler's + // version check in Handle. +} + +func TestHandle_UnknownCmd(t *testing.T) { + rh := newInitdHandler(t) + + req := relay.Request{Version: relay.Version1, Cmd: 0xFF} + var buf bytes.Buffer + req.WriteTo(&buf) + + fc := &fakeConn{buf: buf.Bytes()} + err := rh.Handle(context.Background(), fc) + if err == nil || err.Error() != "relay: unknown command" { + t.Errorf("err = %v, want relay: unknown command", err) + } + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusBadRequest { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusBadRequest) + } +} + +func TestHandle_EmptyAddress(t *testing.T) { + rh := newInitdHandler(t) + + // Build a connect request with an AddrFeature that has empty host + req := relay.Request{Version: relay.Version1, Cmd: relay.CmdConnect} + req.Features = append(req.Features, &relay.AddrFeature{}) + var buf bytes.Buffer + req.WriteTo(&buf) + + fc := &fakeConn{buf: buf.Bytes()} + err := rh.Handle(context.Background(), fc) + if err == nil { + t.Fatal("expected error") + } +} + +func TestHandle_WithBypass(t *testing.T) { + mb := &mockBypass{ + containsFn: func(ctx context.Context, network, addr string, opts ...bypass.Option) bool { + return true + }, + } + rh := newInitdHandler(t, handler.BypassOption(mb)) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != xbypass.ErrBypass { + t.Errorf("err = %v, want %v", err, xbypass.ErrBypass) + } + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusForbidden { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusForbidden) + } +} + +func TestHandle_Unauthorized(t *testing.T) { + ma := &mockAuther{ + authenticateFn: func(ctx context.Context, user, pass string, opts ...auth.Option) (string, bool) { + return "", false + }, + } + rh := newInitdHandler(t, handler.AutherOption(ma)) + + req := relay.Request{Version: relay.Version1, Cmd: relay.CmdConnect} + req.Features = append(req.Features, &relay.UserAuthFeature{Username: "user", Password: "pass"}) + af := &relay.AddrFeature{} + af.ParseFrom("example.com:80") + req.Features = append(req.Features, af) + var buf bytes.Buffer + req.WriteTo(&buf) + + fc := &fakeConn{buf: buf.Bytes()} + err := rh.Handle(context.Background(), fc) + if err == nil || err.Error() != "relay: unauthorized" { + t.Errorf("err = %v, want relay: unauthorized", err) + } + resp := readRelayResponse(t, fc.writeBuf.Bytes()) + if resp.Status != relay.StatusUnauthorized { + t.Errorf("Status = %d, want %d", resp.Status, relay.StatusUnauthorized) + } +} + +func TestHandle_Authorized(t *testing.T) { + ma := &mockAuther{ + authenticateFn: func(ctx context.Context, user, pass string, opts ...auth.Option) (string, bool) { + return "client-1", true + }, + } + mr := &mockRouter{ + dialFn: func(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) { + pr, pw := net.Pipe() + pw.Close() + return pr, nil + }, + } + rh := newInitdHandler(t, + handler.AutherOption(ma), + handler.RouterOption(mr), + ) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } +} + +func TestHandle_RateLimited(t *testing.T) { + mrc := &mockRateLimiterContainer{ + limiterFn: func(key string) rate.Limiter { + return &mockRateLimiter{ + allowFn: func(n int) bool { return false }, + } + }, + } + rh := NewHandler( + handler.RateLimiterOption(mrc), + handler.LoggerOption(&testLogger{}), + ).(*relayHandler) + if err := rh.Init(testMD(nil)); err != nil { + t.Fatalf("Init: %v", err) + } + + fc := &fakeConn{} + err := rh.Handle(context.Background(), fc) + if err != rate_limiter.ErrRateLimit { + t.Errorf("err = %v, want %v", err, rate_limiter.ErrRateLimit) + } +} + +func TestHandle_ReadTimeout(t *testing.T) { + h := &relayHandler{ + options: handler.Options{Logger: &testLogger{}}, + } + h.parseMetadata(testMD(map[string]any{"readTimeout": "100ms"})) + + fc := &fakeConn{} + err := h.Handle(context.Background(), fc) + if err == nil { + t.Error("expected error from read timeout") + } +} + +func TestHandle_ContextCancelled(t *testing.T) { + mr := &mockRouter{ + dialFn: func(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) { + pr, _ := net.Pipe() + return pr, nil + }, + } + rh := newInitdHandler(t, handler.RouterOption(mr)) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + // The request is read from conn first (which succeeds because fakeConn has data), + // then handleConnect tries to dial, which may fail if ctx is cancelled. + // This test verifies no panic occurs with cancelled context. + err := rh.Handle(ctx, fc) + _ = err // may or may not error depending on router dial behavior +} + +// --------------------------------------------------------------------------- +// Handle — forward mode +// --------------------------------------------------------------------------- + +func TestHandle_ForwardMode_NoTarget(t *testing.T) { + mh := &mockHop{ + selectFn: func(ctx context.Context) *chain.Node { + return nil + }, + } + rh := newInitdHandler(t, handler.RouterOption(&mockRouter{})) + rh.Forward(mh) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err == nil || err.Error() != "target not available" { + t.Errorf("err = %v, want target not available", err) + } +} + +func TestHandle_ForwardMode_TargetFound(t *testing.T) { + mh := &mockHop{ + selectFn: func(ctx context.Context) *chain.Node { + return makeTestNode("example.com:80") + }, + } + mr := &mockRouter{ + dialFn: func(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) { + pr, pw := net.Pipe() + pw.Close() + return pr, nil + }, + } + rh := newInitdHandler(t, handler.RouterOption(mr)) + rh.Forward(mh) + + fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")} + err := rh.Handle(context.Background(), fc) + if err != nil { + t.Fatalf("Handle: %v", err) + } +} + +// --------------------------------------------------------------------------- +// observeStats +// --------------------------------------------------------------------------- + +func TestObserveStats_NilObserver(t *testing.T) { + rh := &relayHandler{} + rh.observeStats(context.Background()) // should not panic +} + +func TestObserveStats_NormalTick(t *testing.T) { + obs := &fakeObserver{eventsCh: make(chan []observer.Event, 10)} + rh := &relayHandler{ + options: handler.Options{Observer: obs, Logger: &testLogger{}}, + md: metadata{observerPeriod: 50 * time.Millisecond}, + } + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + rh.observeStats(ctx) +} + +func TestObserveStats_Cancel(t *testing.T) { + obs := &fakeObserver{eventsCh: make(chan []observer.Event, 10)} + rh := &relayHandler{ + options: handler.Options{Observer: obs}, + md: metadata{observerPeriod: time.Hour}, + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + rh.observeStats(ctx) // should return immediately +} + +// --------------------------------------------------------------------------- +// checkRateLimit +// --------------------------------------------------------------------------- + +func TestCheckRateLimit_NoLimiter(t *testing.T) { + rh := &relayHandler{} + if !rh.checkRateLimit(&net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345}) { + t.Error("should allow when no rate limiter") + } +} + +func TestCheckRateLimit_NoLimiterForKey(t *testing.T) { + mrc := &mockRateLimiterContainer{ + limiterFn: func(key string) rate.Limiter { return nil }, + } + rh := &relayHandler{ + options: handler.Options{RateLimiter: mrc}, + } + if !rh.checkRateLimit(&net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345}) { + t.Error("should allow when no limiter for key") + } +} + +func TestCheckRateLimit_LimiterDenies(t *testing.T) { + mrc := &mockRateLimiterContainer{ + limiterFn: func(key string) rate.Limiter { + return &mockRateLimiter{ + allowFn: func(n int) bool { return false }, + } + }, + } + rh := &relayHandler{ + options: handler.Options{RateLimiter: mrc}, + } + if rh.checkRateLimit(&net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345}) { + t.Error("should deny when limiter says no") + } +} + +// --------------------------------------------------------------------------- +// dummyRecorder +// --------------------------------------------------------------------------- + +type dummyRecorder struct{} + +func (r *dummyRecorder) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error { + return nil +} \ No newline at end of file diff --git a/handler/relay/helpers_test.go b/handler/relay/helpers_test.go new file mode 100644 index 00000000..be132097 --- /dev/null +++ b/handler/relay/helpers_test.go @@ -0,0 +1,406 @@ +package relay + +import ( + "bytes" + "context" + "errors" + "io" + "net" + "sync" + "time" + + "github.com/go-gost/core/auth" + "github.com/go-gost/core/bypass" + "github.com/go-gost/core/chain" + "github.com/go-gost/core/hop" + "github.com/go-gost/core/limiter" + "github.com/go-gost/core/limiter/rate" + "github.com/go-gost/core/limiter/traffic" + "github.com/go-gost/core/logger" + core_metadata "github.com/go-gost/core/metadata" + "github.com/go-gost/core/observer" + "github.com/go-gost/relay" + xmetadata "github.com/go-gost/x/metadata" +) + +// --------------------------------------------------------------------------- +// testLogger — silent logger +// --------------------------------------------------------------------------- + +type testLogger struct{} + +func (l *testLogger) WithFields(map[string]any) logger.Logger { return l } +func (l *testLogger) IsLevelEnabled(logger.LogLevel) bool { return false } +func (l *testLogger) GetLevel() logger.LogLevel { return logger.InfoLevel } +func (l *testLogger) Trace(args ...any) {} +func (l *testLogger) Tracef(string, ...any) {} +func (l *testLogger) Debug(args ...any) {} +func (l *testLogger) Debugf(string, ...any) {} +func (l *testLogger) Info(args ...any) {} +func (l *testLogger) Infof(string, ...any) {} +func (l *testLogger) Warn(args ...any) {} +func (l *testLogger) Warnf(string, ...any) {} +func (l *testLogger) Error(args ...any) {} +func (l *testLogger) Errorf(string, ...any) {} +func (l *testLogger) Fatal(args ...any) {} +func (l *testLogger) Fatalf(string, ...any) {} + +// --------------------------------------------------------------------------- +// testMD — wraps a map into metadata.Metadata +// --------------------------------------------------------------------------- + +func testMD(m map[string]any) core_metadata.Metadata { + return xmetadata.NewMetadata(m) +} + +// --------------------------------------------------------------------------- +// fakeConn — byte-slice backed net.Conn (reads from buf, collects writes) +// --------------------------------------------------------------------------- + +type fakeConn struct { + net.Conn + buf []byte + offset int + writeBuf bytes.Buffer + closed bool + mu sync.Mutex + addr net.Addr +} + +func (c *fakeConn) Read(b []byte) (int, error) { + c.mu.Lock() + defer c.mu.Unlock() + if c.offset >= len(c.buf) { + return 0, io.EOF + } + n := copy(b, c.buf[c.offset:]) + c.offset += n + return n, nil +} + +func (c *fakeConn) Write(b []byte) (int, error) { + c.mu.Lock() + defer c.mu.Unlock() + return c.writeBuf.Write(b) +} + +func (c *fakeConn) Close() error { + c.mu.Lock() + defer c.mu.Unlock() + c.closed = true + return nil +} + +func (c *fakeConn) LocalAddr() net.Addr { + if c.addr != nil { + return c.addr + } + return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8080} +} + +func (c *fakeConn) RemoteAddr() net.Addr { + if c.addr != nil { + return c.addr + } + return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345} +} + +func (c *fakeConn) SetDeadline(time.Time) error { return nil } +func (c *fakeConn) SetReadDeadline(time.Time) error { return nil } +func (c *fakeConn) SetWriteDeadline(time.Time) error { return nil } + +// --------------------------------------------------------------------------- +// pipeConn — bidirectional pipe-based net.Conn using io.Pipe +// --------------------------------------------------------------------------- + +type pipeConn struct { + reader *io.PipeReader + writer *io.PipeWriter + laddr net.Addr + raddr net.Addr + closed bool + closeMu sync.Mutex +} + +func newPipePair() (*pipeConn, *pipeConn) { + pr1, pw1 := io.Pipe() + pr2, pw2 := io.Pipe() + a := &pipeConn{reader: pr1, writer: pw2} + b := &pipeConn{reader: pr2, writer: pw1} + return a, b +} + +func (c *pipeConn) Read(b []byte) (int, error) { return c.reader.Read(b) } +func (c *pipeConn) Write(b []byte) (int, error) { return c.writer.Write(b) } + +func (c *pipeConn) Close() error { + c.closeMu.Lock() + defer c.closeMu.Unlock() + if c.closed { + return nil + } + c.closed = true + c.reader.Close() + c.writer.Close() + return nil +} + +func (c *pipeConn) LocalAddr() net.Addr { + if c.laddr != nil { + return c.laddr + } + return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8080} +} +func (c *pipeConn) RemoteAddr() net.Addr { + if c.raddr != nil { + return c.raddr + } + return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345} +} +func (c *pipeConn) SetDeadline(time.Time) error { return nil } +func (c *pipeConn) SetReadDeadline(time.Time) error { return nil } +func (c *pipeConn) SetWriteDeadline(time.Time) error { return nil } + +// --------------------------------------------------------------------------- +// fakeObserver — channels-based observer with optional error function +// --------------------------------------------------------------------------- + +type fakeObserver struct { + eventsCh chan []observer.Event + errFunc func() error // optional; called for each Observe +} + +func (o *fakeObserver) Observe(ctx context.Context, events []observer.Event, opts ...observer.Option) error { + if o.errFunc != nil { + if err := o.errFunc(); err != nil { + return err + } + } + select { + case o.eventsCh <- events: + default: + } + return nil +} + +func (o *fakeObserver) Events() <-chan []observer.Event { return o.eventsCh } + +// --------------------------------------------------------------------------- +// mockRouter — implements chain.Router +// --------------------------------------------------------------------------- + +type mockRouter struct { + dialFn func(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) +} + +func (r *mockRouter) Dial(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) { + if r.dialFn != nil { + return r.dialFn(ctx, network, address, opts...) + } + return nil, errors.New("mockRouter: no dial function") +} + +// Required by chain.Router interface (stubs) +func (r *mockRouter) Options() *chain.RouterOptions { return nil } +func (r *mockRouter) Bind(ctx context.Context, network, address string, opts ...chain.BindOption) (net.Listener, error) { + return nil, errors.New("mockRouter: Bind not implemented") +} + +// --------------------------------------------------------------------------- +// mockHop — implements hop.Hop +// --------------------------------------------------------------------------- + +type mockHop struct { + selectFn func(ctx context.Context) *chain.Node +} + +func (h *mockHop) Select(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + if h.selectFn != nil { + return h.selectFn(ctx) + } + return nil +} + +// --------------------------------------------------------------------------- +// mockBypass — implements bypass.Bypass +// --------------------------------------------------------------------------- + +type mockBypass struct { + containsFn func(ctx context.Context, network, addr string, opts ...bypass.Option) bool +} + +func (b *mockBypass) Contains(ctx context.Context, network, addr string, opts ...bypass.Option) bool { + if b.containsFn != nil { + return b.containsFn(ctx, network, addr, opts...) + } + return false +} +func (b *mockBypass) IsWhitelist() bool { return false } + +// --------------------------------------------------------------------------- +// mockAuther — implements auth.Authenticator +// --------------------------------------------------------------------------- + +type mockAuther struct { + authenticateFn func(ctx context.Context, user, pass string, opts ...auth.Option) (string, bool) +} + +func (a *mockAuther) Authenticate(ctx context.Context, user, pass string, opts ...auth.Option) (string, bool) { + if a.authenticateFn != nil { + return a.authenticateFn(ctx, user, pass, opts...) + } + return "", false +} + +// --------------------------------------------------------------------------- +// mockRateLimiter — implements rate.Limiter (single rate limiter) +// --------------------------------------------------------------------------- + +type mockRateLimiter struct { + allowFn func(n int) bool + waitFn func(ctx context.Context) error +} + +func (l *mockRateLimiter) Allow(n int) bool { + if l.allowFn != nil { + return l.allowFn(n) + } + return true +} +func (l *mockRateLimiter) Limit() float64 { return 1 } +func (l *mockRateLimiter) Wait(ctx context.Context) error { + if l.waitFn != nil { + return l.waitFn(ctx) + } + return nil +} + +// mockRateLimiterContainer — implements rate.RateLimiter (container by key) +type mockRateLimiterContainer struct { + limiterFn func(key string) rate.Limiter +} + +func (c *mockRateLimiterContainer) Limiter(key string) rate.Limiter { + if c.limiterFn != nil { + return c.limiterFn(key) + } + return nil +} + +// --------------------------------------------------------------------------- +// mockTrafficLimiter — implements traffic.TrafficLimiter +// --------------------------------------------------------------------------- + +type mockTrafficLimiter struct{} + +func (l *mockTrafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter { + return nil +} +func (l *mockTrafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter { + return nil +} + +// --------------------------------------------------------------------------- +// buildRelayRequest helpers — serialize relay requests for fakeConn.buf +// --------------------------------------------------------------------------- + +func buildRelayConnectRequest(t testingT, address, network string) []byte { + t.Helper() + req := relay.Request{ + Version: relay.Version1, + Cmd: relay.CmdConnect, + } + + af := &relay.AddrFeature{} + if address != "" { + if err := af.ParseFrom(address); err != nil { + t.Fatal(err) + } + } + req.Features = append(req.Features, af) + + if network != "" && network != "tcp" { + req.Features = append(req.Features, &relay.NetworkFeature{ + Network: toRelayNetworkID(network), + }) + } + + var buf bytes.Buffer + if _, err := req.WriteTo(&buf); err != nil { + t.Fatal(err) + } + return buf.Bytes() +} + +func buildRelayBindRequest(t testingT, address, network string) []byte { + t.Helper() + cmd := relay.CmdBind + if network == "udp" || network == "udp4" || network == "udp6" { + cmd |= relay.FUDP + } + req := relay.Request{ + Version: relay.Version1, + Cmd: cmd, + } + + af := &relay.AddrFeature{} + if address != "" { + if err := af.ParseFrom(address); err != nil { + t.Fatal(err) + } + } + req.Features = append(req.Features, af) + + if network != "" && network != "tcp" { + req.Features = append(req.Features, &relay.NetworkFeature{ + Network: toRelayNetworkID(network), + }) + } + + var buf bytes.Buffer + if _, err := req.WriteTo(&buf); err != nil { + t.Fatal(err) + } + return buf.Bytes() +} + +func toRelayNetworkID(n string) relay.NetworkID { + switch n { + case "udp", "udp4", "udp6": + return relay.NetworkUDP + case "unix": + return relay.NetworkUnix + case "serial": + return relay.NetworkSerial + default: + return relay.NetworkTCP + } +} + +// testingT is a subset of testing.T for use by helper functions. +type testingT interface { + Helper() + Fatal(args ...any) + Fatalf(format string, args ...any) +} + +// --------------------------------------------------------------------------- +// makeTestNode — creates a minimal *hop.Node for mockHop returns +// --------------------------------------------------------------------------- + +func makeTestNode(addr string) *chain.Node { + return chain.NewNode("test", addr) +} + +// --------------------------------------------------------------------------- +// readRelayResponse — deserializes a relay.Response from bytes +// --------------------------------------------------------------------------- + +func readRelayResponse(t testingT, data []byte) *relay.Response { + t.Helper() + resp := &relay.Response{} + if _, err := resp.ReadFrom(bytes.NewReader(data)); err != nil { + t.Fatalf("read relay response: %v", err) + } + return resp +} \ No newline at end of file diff --git a/handler/relay/metadata_test.go b/handler/relay/metadata_test.go new file mode 100644 index 00000000..1f722bb0 --- /dev/null +++ b/handler/relay/metadata_test.go @@ -0,0 +1,254 @@ +package relay + +import ( + "testing" + "time" + + core_metadata "github.com/go-gost/core/metadata" + xmetadata "github.com/go-gost/x/metadata" +) + +func TestParseMetadata_Defaults(t *testing.T) { + h := &relayHandler{} + err := h.parseMetadata(xmetadata.NewMetadata(nil)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if h.md.readTimeout != 15*time.Second { + t.Errorf("readTimeout = %v, want 15s", h.md.readTimeout) + } + if h.md.enableBind { + t.Error("enableBind = true, want false") + } + if h.md.noDelay { + t.Error("noDelay = true, want false") + } + if h.md.hash != "" { + t.Errorf("hash = %q, want empty", h.md.hash) + } + if h.md.udpBufferSize != 0 { + t.Errorf("udpBufferSize = %d, want 0", h.md.udpBufferSize) + } + if h.md.observerPeriod != 5*time.Second { + t.Errorf("observerPeriod = %v, want 5s", h.md.observerPeriod) + } + if h.md.observerResetTraffic { + t.Error("observerResetTraffic = true, want false") + } + if h.md.sniffing { + t.Error("sniffing = true, want false") + } + if h.md.muxCfg == nil { + t.Fatal("muxCfg = nil, want non-nil") + } + // Note: relay handler does NOT default muxCfg.Version or MaxStreamBuffer + // at parseMetadata time. Those defaults are in the tunnel handler only. +} + +func TestParseMetadata_ReadTimeout(t *testing.T) { + tests := []struct { + name string + input core_metadata.Metadata + want time.Duration + }{ + {"zero", testMD(map[string]any{"readTimeout": 0}), 15 * time.Second}, + {"string 30s", testMD(map[string]any{"readTimeout": "30s"}), 30 * time.Second}, + {"int seconds", testMD(map[string]any{"readTimeout": 10}), 10 * time.Second}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h := &relayHandler{} + h.parseMetadata(tt.input) + if h.md.readTimeout != tt.want { + t.Errorf("readTimeout = %v, want %v", h.md.readTimeout, tt.want) + } + }) + } +} + +func TestParseMetadata_UDPBufferSize(t *testing.T) { + tests := []struct { + name string + input core_metadata.Metadata + want int + }{ + {"zero", testMD(nil), 0}, + {"udp.bufferSize", testMD(map[string]any{"udp.bufferSize": 4096}), 4096}, + {"udpBufferSize alias", testMD(map[string]any{"udpBufferSize": 8192}), 8192}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h := &relayHandler{} + h.parseMetadata(tt.input) + if h.md.udpBufferSize != tt.want { + t.Errorf("udpBufferSize = %d, want %d", h.md.udpBufferSize, tt.want) + } + }) + } +} + +func TestParseMetadata_Bind(t *testing.T) { + h := &relayHandler{} + h.parseMetadata(testMD(map[string]any{"bind": true})) + if !h.md.enableBind { + t.Error("enableBind = false, want true") + } +} + +func TestParseMetadata_NoDelay(t *testing.T) { + h := &relayHandler{} + h.parseMetadata(testMD(map[string]any{"nodelay": true})) + if !h.md.noDelay { + t.Error("noDelay = false, want true") + } +} + +func TestParseMetadata_Hash(t *testing.T) { + h := &relayHandler{} + h.parseMetadata(testMD(map[string]any{"hash": "host"})) + if h.md.hash != "host" { + t.Errorf("hash = %q, want host", h.md.hash) + } +} + +func TestParseMetadata_MuxConfig(t *testing.T) { + h := &relayHandler{} + h.parseMetadata(testMD(map[string]any{ + "mux.version": 3, + "mux.keepaliveInterval": "10s", + "mux.keepaliveDisabled": true, + "mux.keepaliveTimeout": "5s", + "mux.maxFrameSize": 4096, + "mux.maxReceiveBuffer": 8192, + "mux.maxStreamBuffer": 65536, + })) + if h.md.muxCfg.Version != 3 { + t.Errorf("muxCfg.Version = %d, want 3", h.md.muxCfg.Version) + } + if h.md.muxCfg.KeepAliveInterval != 10*time.Second { + t.Errorf("muxCfg.KeepAliveInterval = %v, want 10s", h.md.muxCfg.KeepAliveInterval) + } + if !h.md.muxCfg.KeepAliveDisabled { + t.Error("muxCfg.KeepAliveDisabled = false, want true") + } + if h.md.muxCfg.KeepAliveTimeout != 5*time.Second { + t.Errorf("muxCfg.KeepAliveTimeout = %v, want 5s", h.md.muxCfg.KeepAliveTimeout) + } + if h.md.muxCfg.MaxFrameSize != 4096 { + t.Errorf("muxCfg.MaxFrameSize = %d, want 4096", h.md.muxCfg.MaxFrameSize) + } + if h.md.muxCfg.MaxReceiveBuffer != 8192 { + t.Errorf("muxCfg.MaxReceiveBuffer = %d, want 8192", h.md.muxCfg.MaxReceiveBuffer) + } + if h.md.muxCfg.MaxStreamBuffer != 65536 { + t.Errorf("muxCfg.MaxStreamBuffer = %d, want 65536", h.md.muxCfg.MaxStreamBuffer) + } +} + +func TestParseMetadata_Observer(t *testing.T) { + tests := []struct { + name string + input core_metadata.Metadata + wantPeriod time.Duration + wantReset bool + }{ + {"defaults", testMD(nil), 5 * time.Second, false}, + {"custom period", testMD(map[string]any{"observePeriod": "10s"}), 10 * time.Second, false}, + {"clamped below 1s", testMD(map[string]any{"observePeriod": "500ms"}), time.Second, false}, + {"alias observer.period", testMD(map[string]any{"observer.period": "3s"}), 3 * time.Second, false}, + {"alias observer.observePeriod", testMD(map[string]any{"observer.observePeriod": "7s"}), 7 * time.Second, false}, + {"reset traffic", testMD(map[string]any{"observer.resetTraffic": true}), 5 * time.Second, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h := &relayHandler{} + h.parseMetadata(tt.input) + if h.md.observerPeriod != tt.wantPeriod { + t.Errorf("observerPeriod = %v, want %v", h.md.observerPeriod, tt.wantPeriod) + } + if h.md.observerResetTraffic != tt.wantReset { + t.Errorf("observerResetTraffic = %v, want %v", h.md.observerResetTraffic, tt.wantReset) + } + }) + } +} + +func TestParseMetadata_Sniffing(t *testing.T) { + h := &relayHandler{} + h.parseMetadata(testMD(map[string]any{ + "sniffing": true, + "sniffing.timeout": "2s", + "sniffing.websocket": true, + "sniffing.websocket.sampleRate": 0.5, + })) + if !h.md.sniffing { + t.Error("sniffing = false, want true") + } + if h.md.sniffingTimeout != 2*time.Second { + t.Errorf("sniffingTimeout = %v, want 2s", h.md.sniffingTimeout) + } + if !h.md.sniffingWebsocket { + t.Error("sniffingWebsocket = false, want true") + } + if h.md.sniffingWebsocketSampleRate != 0.5 { + t.Errorf("sniffingWebsocketSampleRate = %v, want 0.5", h.md.sniffingWebsocketSampleRate) + } +} + +func TestParseMetadata_LimiterRefresh(t *testing.T) { + h := &relayHandler{} + h.parseMetadata(testMD(map[string]any{ + "limiter.refreshInterval": "30s", + "limiter.cleanupInterval": "60s", + })) + if h.md.limiterRefreshInterval != 30*time.Second { + t.Errorf("limiterRefreshInterval = %v, want 30s", h.md.limiterRefreshInterval) + } + if h.md.limiterCleanupInterval != 60*time.Second { + t.Errorf("limiterCleanupInterval = %v, want 60s", h.md.limiterCleanupInterval) + } +} + +func TestParseMetadata_Mitm(t *testing.T) { + t.Run("not set", func(t *testing.T) { + h := &relayHandler{} + err := h.parseMetadata(testMD(nil)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if h.md.certificate != nil { + t.Error("certificate should be nil") + } + if h.md.privateKey != nil { + t.Error("privateKey should be nil") + } + }) + + t.Run("cert file not found", func(t *testing.T) { + h := &relayHandler{} + err := h.parseMetadata(testMD(map[string]any{ + "mitm.certFile": "/nonexistent/cert.pem", + "mitm.keyFile": "/nonexistent/key.pem", + })) + if err == nil { + t.Error("expected error for nonexistent cert file") + } + }) + + t.Run("alpn set", func(t *testing.T) { + h := &relayHandler{} + h.parseMetadata(testMD(map[string]any{"mitm.alpn": "h2,http/1.1"})) + if h.md.alpn != "h2,http/1.1" { + t.Errorf("alpn = %q, want h2,http/1.1", h.md.alpn) + } + }) +} + +func TestParseMetadata_ObserverPeriodDefaultsTo5s(t *testing.T) { + h := &relayHandler{} + h.parseMetadata(testMD(map[string]any{"observePeriod": 0})) + if h.md.observerPeriod != 5*time.Second { + t.Errorf("observerPeriod = %v, want 5s", h.md.observerPeriod) + } +} \ No newline at end of file diff --git a/handler/relay/observe.go b/handler/relay/observe.go new file mode 100644 index 00000000..b2c9803a --- /dev/null +++ b/handler/relay/observe.go @@ -0,0 +1,52 @@ +package relay + +import ( + "context" + "net" + "time" + + "github.com/go-gost/core/observer" +) + +func (h *relayHandler) checkRateLimit(addr net.Addr) bool { + if h.options.RateLimiter == nil { + return true + } + host, _, _ := net.SplitHostPort(addr.String()) + if limiter := h.options.RateLimiter.Limiter(host); limiter != nil { + return limiter.Allow(1) + } + + return true +} + +func (h *relayHandler) observeStats(ctx context.Context) { + if h.options.Observer == nil { + return + } + + var events []observer.Event + + ticker := time.NewTicker(h.md.observerPeriod) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + if len(events) > 0 { + if err := h.options.Observer.Observe(ctx, events); err == nil { + events = nil + } + // fallthrough: still collect fresh events for this tick + } + + evs := h.stats.Events() + if err := h.options.Observer.Observe(ctx, evs); err != nil { + events = evs + } + + case <-ctx.Done(): + return + } + } +} \ No newline at end of file