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
This commit is contained in:
ginuerzh
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user