diff --git a/README.md b/README.md index 8469cd2..16264f8 100644 --- a/README.md +++ b/README.md @@ -118,6 +118,39 @@ $ xcaddy build --with github.com/imgk/caddy-trojan curl -X POST -H "Content-Type: application/json" -d '{"password": "test1234"}' http://localhost:2019/trojan/users/add ``` +## Rate Limiting / DoS Hardening + +Password validation applies a uniform 250ms delay on every connection (hit and +miss alike, to keep validation timing-constant). This means an unauthenticated +peer can hold a server goroutine + file descriptor for ~250ms per connection by +sending 58 bytes after the TLS handshake. Caddy itself has no built-in +per-IP connection limiting, so add one at the edge: + +1. Bound connection lifetime in the server block so stalled peers are dropped: +``` +servers :443 { + listener_wrappers { + trojan + } + read_timeout 30s + write_timeout 30s + idle_timeout 1m +} +``` +2. For per-IP connection concurrency limits, build caddy with a limiting +plugin, e.g. [mholt/caddy-limit](https://github.com/mholt/caddy-limit): +``` +xcaddy build --with github.com/mholt/caddy-limit +``` +``` +servers :443 { + ... + conn_limit 20 +} +``` +3. As a last line of defense, use a firewall / fail2ban at the host level to +cap new connections per source IP on port 443. + ## Docker ``` diff --git a/app/proxy.go b/app/proxy.go index d1844bb..712fe6c 100644 --- a/app/proxy.go +++ b/app/proxy.go @@ -9,6 +9,7 @@ import ( "io" "net" "net/http" + "strings" "sync" "golang.org/x/net/proxy" @@ -17,6 +18,7 @@ import ( "github.com/caddyserver/caddy/v2/caddyconfig" "github.com/imgk/caddy-trojan/pkgs/domaintree" + "github.com/imgk/caddy-trojan/pkgs/rawconn" "github.com/imgk/caddy-trojan/pkgs/trojan" "github.com/imgk/caddy-trojan/pkgs/x" ) @@ -145,7 +147,7 @@ func (p *EnvProxy) Provision(ctx caddy.Context) error { if p.ProxyRaw != nil { mod, err := ctx.LoadModule(p, "ProxyRaw") if err != nil { - return nil + return err } p.proxy = mod.(Proxy) p.dialer = proxy.FromEnvironmentUsing(p.proxy) @@ -210,7 +212,7 @@ func (p *SocksProxy) Provision(ctx caddy.Context) error { if p.ProxyRaw != nil { mod, err := ctx.LoadModule(p, "ProxyRaw") if err != nil { - return nil + return err } p.proxy = mod.(Proxy) } else { @@ -291,7 +293,7 @@ func (p *HttpProxy) Provision(ctx caddy.Context) error { if p.ProxyRaw != nil { mod, err := ctx.LoadModule(p, "ProxyRaw") if err != nil { - return nil + return err } p.proxy = mod.(Proxy) } else { @@ -310,7 +312,21 @@ func (p *HttpProxy) Close() error { return p.proxy.Close() } +// validateProxyTarget rejects target addresses that could break out of the +// HTTP CONNECT request line (CR/LF header injection) or smuggle bytes into +// the upstream proxy connection. +func validateProxyTarget(addr string) error { + if strings.ContainsAny(addr, "\r\n\x00 \t") { + return errors.New("invalid target address: control characters are not allowed") + } + return nil +} + func (p *HttpProxy) Dial(network, addr string) (net.Conn, error) { + if err := validateProxyTarget(addr); err != nil { + return nil, err + } + conn, err := p.proxy.Dial("tcp", p.Server) if err != nil { return nil, err @@ -318,7 +334,8 @@ func (p *HttpProxy) Dial(network, addr string) (net.Conn, error) { req, err := http.NewRequest(http.MethodConnect, p.tcpURL, nil) if err != nil { - return conn, fmt.Errorf("request error: %w", err) + _ = conn.Close() + return nil, fmt.Errorf("request error: %w", err) } req.URL.Opaque = addr if p.basicAuth != "" { @@ -327,15 +344,26 @@ func (p *HttpProxy) Dial(network, addr string) (net.Conn, error) { err = req.WriteProxy(conn) if err != nil { - return conn, fmt.Errorf("write request error: %w", err) + _ = conn.Close() + return nil, fmt.Errorf("write request error: %w", err) } - resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + br := bufio.NewReader(conn) + resp, err := http.ReadResponse(br, nil) if err != nil { - return conn, fmt.Errorf("read response error: %w", err) + _ = conn.Close() + return nil, fmt.Errorf("read response error: %w", err) } if resp.StatusCode != http.StatusOK { - return conn, fmt.Errorf("server status code error: %v", resp.StatusCode) + _ = conn.Close() + return nil, fmt.Errorf("server status code error: %v", resp.StatusCode) + } + // bufio may have read past the response head into the tunnel payload; + // replay those bytes so the first tunneled packet is not lost. + if n := br.Buffered(); n > 0 { + if b, err := br.Peek(n); err == nil { + return rawconn.NewConn(conn, b), nil + } } return conn, nil } diff --git a/app/proxy_http_dial_test.go b/app/proxy_http_dial_test.go new file mode 100644 index 0000000..19fa330 --- /dev/null +++ b/app/proxy_http_dial_test.go @@ -0,0 +1,319 @@ +package app + +import ( + "context" + "encoding/json" + "errors" + "net" + "net/http" + "sync/atomic" + "testing" + "time" + + "github.com/caddyserver/caddy/v2" +) + +// trackingProxy is a test-only caddy module that returns a countingConn on +// Dial so tests can assert HttpProxy.Dial called Close() on the underlying +// connection. The module is registered in init() below; tests load it as +// HttpProxy's pre_proxy and read closeCount back via the *trackingProxy +// stored in h.proxy after Provision. +type trackingProxy struct { + Server string `json:"server"` + closed atomic.Int32 + connClosed atomic.Int32 // conns that closed via our wrapper +} + +func init() { + caddy.RegisterModule(&trackingProxy{}) +} + +func (*trackingProxy) CaddyModule() caddy.ModuleInfo { + return caddy.ModuleInfo{ + ID: "trojan.proxy.tracking", + New: func() caddy.Module { return new(trackingProxy) }, + } +} + +func (p *trackingProxy) Provision(_ caddy.Context) error { return nil } + +func (p *trackingProxy) Dial(network, _ string) (net.Conn, error) { + c, err := net.Dial(network, p.Server) + if err != nil { + return nil, err + } + return &countingConn{Conn: c, parent: p}, nil +} + +func (p *trackingProxy) ListenPacket(_, _ string) (net.PacketConn, error) { + return nil, errors.New("not supported") +} + +func (p *trackingProxy) Close() error { return nil } + +type countingConn struct { + net.Conn + parent *trackingProxy +} + +func (c *countingConn) Close() error { + c.parent.closed.Add(1) + return c.Conn.Close() +} + +// newHttpProxyWithTracking provisions an HttpProxy that uses trackingProxy +// as its pre_proxy dialer. Returns the HttpProxy and the *trackingProxy so +// tests can read closeCount after Dial returns. +func newHttpProxyWithTracking(t *testing.T, serverAddr string) (*HttpProxy, *trackingProxy) { + t.Helper() + ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()}) + t.Cleanup(cancel) + // Embed the server address in the ProxyRaw JSON so the loaded + // trackingProxy unmarshals it into its Server field. The pre-built tp + // here is only used for sanity checks; the actual proxy that HttpProxy + // stores in p.proxy is a fresh instance created by caddy.LoadModule. + tp := &trackingProxy{Server: serverAddr} + raw := []byte(`{"proxy":"tracking","server":"` + serverAddr + `"}`) + h := &HttpProxy{ + Server: serverAddr, + ProxyRaw: json.RawMessage(raw), + } + if err := h.Provision(ctx); err != nil { + t.Fatalf("HttpProxy.Provision: %v", err) + } + return h, tp +} + +// fakeProxy is a tiny TCP server used only by HttpProxy.Dial tests. Each +// accepted connection is handed to onConn, which writes whatever response +// the test wants. +type fakeProxy struct { + listener net.Listener + onConn func(c net.Conn) +} + +func newFakeProxy(t *testing.T, onConn func(c net.Conn)) *fakeProxy { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + fp := &fakeProxy{listener: ln, onConn: onConn} + go func() { + for { + c, err := ln.Accept() + if err != nil { + return + } + go onConn(c) + } + }() + return fp +} + +func (fp *fakeProxy) addr() string { return fp.listener.Addr().String() } +func (fp *fakeProxy) close() { fp.listener.Close() } + +// newHttpProxyForTest provisions an HttpProxy pointing at the fake upstream +// directly (no pre_proxy chain, so NoProxy.Dial is used). +func newHttpProxyForTest(t *testing.T, serverAddr string) *HttpProxy { + t.Helper() + ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()}) + t.Cleanup(cancel) + p := &HttpProxy{Server: serverAddr} + if err := p.Provision(ctx); err != nil { + t.Fatalf("Provision: %v", err) + } + return p +} + +// readExact reads exactly len(buf) bytes from c within timeout, or fails. +func readExact(c net.Conn, buf []byte, timeout time.Duration) error { + type result struct { + err error + } + done := make(chan result, 1) + go func() { + read := 0 + for read < len(buf) { + n, err := c.Read(buf[read:]) + read += n + if err != nil { + done <- result{err} + return + } + } + done <- result{nil} + }() + select { + case r := <-done: + return r.err + case <-time.After(timeout): + return errors.New("read timeout") + } +} + +// TestHttpProxyDialClosesConnOnNonOKStatus covers the "server status code +// error" branch in HttpProxy.Dial: when the upstream returns anything but +// 200, the proxy conn must be closed (previously leaked). +func TestHttpProxyDialClosesConnOnNonOKStatus(t *testing.T) { + t.Parallel() + + fp := newFakeProxy(t, func(c net.Conn) { + // Drain the CONNECT request line + headers. + buf := make([]byte, 4096) + _, _ = c.Read(buf) + // Reply with a non-200 status. + _, _ = c.Write([]byte("HTTP/1.1 407 Proxy Authentication Required\r\nContent-Length: 0\r\n\r\n")) + // Hold the conn open so the test can read whatever the client writes. + _, _ = c.Read(make([]byte, 1)) + }) + defer fp.close() + + h, _ := newHttpProxyWithTracking(t, fp.addr()) + + conn, err := h.Dial("tcp", "example.com:443") + if err == nil { + if conn != nil { + conn.Close() + } + t.Fatal("Dial returned no error; want server status code error") + } + if conn != nil { + t.Fatal("Dial returned non-nil conn on error; want nil") + } + + tp := h.proxy.(*trackingProxy) + if got := tp.closed.Load(); got != 1 { + t.Errorf("trackingProxy.closed = %d, want 1 (HttpProxy.Dial must Close conn on non-200 path)", got) + } +} + +// TestHttpProxyDialClosesConnOnWriteError covers the "write request error" +// branch in HttpProxy.Dial: when req.WriteProxy fails, the dialer must +// explicitly close the conn. A fake proxy that reads the CONNECT request then +// closes immediately makes WriteProxy fail with a broken pipe; a +// trackingProxy counts Close calls to assert the fix. +func TestHttpProxyDialClosesConnOnWriteError(t *testing.T) { + t.Parallel() + + // Use SetLinger=0 so Close sends RST, guaranteeing WriteProxy fails + // rather than succeeding (kernel buffer would absorb FIN without RST). + fp := newFakeProxy(t, func(c net.Conn) { + // Wait for the client to start sending the CONNECT request before + // RSTing. On Linux an immediate linger-0 close can reset the + // connection during connect() itself (net.Dial returns ECONNRESET), + // so the post-dial error path would never be exercised. Reading one + // byte guarantees the client's connect completed and WriteProxy + // began; the subsequent RST then fails either a remaining write + // (EPIPE → write-error path) or the response read (ECONNRESET → + // read-error path), both of which must close the conn. + buf := make([]byte, 1) + _, _ = c.Read(buf) + if tcp, ok := c.(*net.TCPConn); ok { + _ = tcp.SetLinger(0) + } + c.Close() + }) + defer fp.close() + + h, _ := newHttpProxyWithTracking(t, fp.addr()) + + conn, err := h.Dial("tcp", "example.com:443") + if err == nil { + if conn != nil { + conn.Close() + } + t.Fatal("Dial returned no error; want write request error") + } + if conn != nil { + t.Fatal("Dial returned non-nil conn on error; want nil") + } + + tp := h.proxy.(*trackingProxy) + if got := tp.closed.Load(); got != 1 { + t.Errorf("trackingProxy.closed = %d, want 1 (HttpProxy.Dial must Close conn on write-error path); Dial err = %v", got, err) + } +} + +// TestHttpProxyDialClosesConnOnReadError covers the "read response error" +// branch in HttpProxy.Dial: when the upstream never sends a valid HTTP +// response, the dialer must close the conn. +func TestHttpProxyDialClosesConnOnReadError(t *testing.T) { + t.Parallel() + + fp := newFakeProxy(t, func(c net.Conn) { + // Drain the CONNECT request, then send garbage that isn't a valid + // HTTP response so http.ReadResponse surfaces an error. + buf := make([]byte, 4096) + _, _ = c.Read(buf) + _, _ = c.Write([]byte("not-a-valid-http-response\r\n\r\n")) + c.Close() + }) + defer fp.close() + + h, _ := newHttpProxyWithTracking(t, fp.addr()) + + conn, err := h.Dial("tcp", "example.com:443") + if err == nil { + if conn != nil { + conn.Close() + } + t.Fatal("Dial returned no error; want read response error") + } + if conn != nil { + t.Fatal("Dial returned non-nil conn on error; want nil") + } + + tp := h.proxy.(*trackingProxy) + if got := tp.closed.Load(); got != 1 { + t.Errorf("trackingProxy.closed = %d, want 1 (HttpProxy.Dial must Close conn on read-error path)", got) + } +} + +// TestHttpProxyDialReplaysBufferedBytes is a regression test for the bufio +// replay fix in HttpProxy.Dial: when bufio.NewReader(conn) reads past the +// HTTP response head into tunneled payload (because the proxy responded +// 200 OK and started sending tunnel data in the same TCP segment), those +// buffered bytes must be replayed to the caller. Without the replay, the +// first tunneled packet would be silently dropped. +func TestHttpProxyDialReplaysBufferedBytes(t *testing.T) { + t.Parallel() + + const payload = "tunneled-bytes-immediately-after-200-ok" + + fp := newFakeProxy(t, func(c net.Conn) { + // Drain the CONNECT request, then send 200 OK + payload in a single + // Write so they arrive in one TCP segment and bufio (4 KiB default + // buffer) reads past the response head into the payload. + buf := make([]byte, 4096) + _, _ = c.Read(buf) + resp := "HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n" + payload + _, _ = c.Write([]byte(resp)) + // Hold the conn open until the test is done. + _, _ = c.Read(make([]byte, 1)) + }) + defer fp.close() + + h := newHttpProxyForTest(t, fp.addr()) + + conn, err := h.Dial("tcp", "example.com:443") + if err != nil { + t.Fatalf("Dial returned error: %v", err) + } + defer conn.Close() + + // The returned conn must expose the payload on its first Read. rawconn + // wraps the real conn and replays the buffered bytes before falling + // through to the underlying Read. + got := make([]byte, len(payload)) + if err := readExact(conn, got, 2*time.Second); err != nil { + t.Fatalf("read exact from Dial conn: %v", err) + } + if string(got) != payload { + t.Errorf("replayed bytes = %q, want %q", got, payload) + } +} + +// keep the http import used (for completeness in the test file) +var _ = http.MethodConnect diff --git a/app/proxy_provision_test.go b/app/proxy_provision_test.go new file mode 100644 index 0000000..26d2231 --- /dev/null +++ b/app/proxy_provision_test.go @@ -0,0 +1,41 @@ +package app + +import ( + "context" + "encoding/json" + "testing" + + "github.com/caddyserver/caddy/v2" +) + +// TestProxyProvisionPropagatesLoadModuleError is a regression test for +// swallowed LoadModule errors in EnvProxy/SocksProxy/HttpProxy.Provision: +// each used to return nil on a bad pre_proxy config, silently starting up +// with a half-configured proxy that would later crash on first dial. After +// the fix, LoadModule errors are propagated so misconfiguration is caught at +// startup. +func TestProxyProvisionPropagatesLoadModuleError(t *testing.T) { + t.Parallel() + + ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()}) + defer cancel() + + bad := json.RawMessage(`{"proxy":"this-module-does-not-exist"}`) + + for _, tc := range []struct { + name string + mod caddy.Module + }{ + {"EnvProxy", &EnvProxy{ProxyRaw: bad}}, + {"SocksProxy", &SocksProxy{ProxyRaw: bad}}, + {"HttpProxy", &HttpProxy{ProxyRaw: bad}}, + } { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + if err := tc.mod.(caddy.Provisioner).Provision(ctx); err == nil { + t.Fatalf("%s.Provision with unknown pre_proxy module = nil, want error", tc.name) + } + }) + } +} diff --git a/app/proxy_shadowsocks.go b/app/proxy_shadowsocks.go index c161699..7d59752 100644 --- a/app/proxy_shadowsocks.go +++ b/app/proxy_shadowsocks.go @@ -54,14 +54,15 @@ func (p *ShadowsocksProxy) Provision(ctx caddy.Context) error { if p.ProxyRaw != nil { mod, err := ctx.LoadModule(p, "ProxyRaw") if err != nil { - return nil + return err } p.proxy = mod.(Proxy) - return nil } else { p.proxy = &NoProxy{} } + // The cipher must always be initialized, including with a pre_proxy chain; + // otherwise Dial/ListenPacket would call methods on a nil core.Cipher. var err error p.cipher, err = core.PickCipher(p.Method, nil, p.Password) if err != nil { diff --git a/app/proxy_shadowsocks_provision_test.go b/app/proxy_shadowsocks_provision_test.go new file mode 100644 index 0000000..373aecb --- /dev/null +++ b/app/proxy_shadowsocks_provision_test.go @@ -0,0 +1,59 @@ +package app + +import ( + "context" + "encoding/json" + "testing" + + "github.com/caddyserver/caddy/v2" +) + +// TestShadowsocksProvisionInitializesCipherWithPreProxy is a regression test +// for a startup-crash bug: Provision's pre_proxy branch used to return +// before the cipher was initialized, so the first Dial/ListenPacket on a +// ShadowsocksProxy with a pre_proxy chain would panic with a nil +// core.Cipher, taking down the whole caddy process. After the fix, the +// cipher is initialized regardless of which branch is taken. +func TestShadowsocksProvisionInitializesCipherWithPreProxy(t *testing.T) { + t.Parallel() + + ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()}) + defer cancel() + + p := &ShadowsocksProxy{ + Method: "aes-256-gcm", + Password: "testpassword123", + // Inline pre_proxy pointing at NoProxy (registered as + // "trojan.proxy.none" via app/proxy.go's init()). + ProxyRaw: json.RawMessage(`{"proxy":"none"}`), + } + if err := p.Provision(ctx); err != nil { + t.Fatalf("Provision error: %v", err) + } + if p.cipher == nil { + t.Fatal("p.cipher is nil after Provision with pre_proxy \u2014 the cipher must be initialized regardless of pre_proxy presence") + } +} + +// TestShadowsocksProvisionPropagatesLoadModuleError is a regression test for +// swallowed LoadModule errors: Provision used to return nil on a bad +// pre_proxy config, silently starting up with a half-configured proxy that +// would later crash on first dial. After the fix, LoadModule errors are +// propagated so misconfiguration is caught at startup. +func TestShadowsocksProvisionPropagatesLoadModuleError(t *testing.T) { + t.Parallel() + + ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()}) + defer cancel() + + p := &ShadowsocksProxy{ + Method: "aes-256-gcm", + Password: "testpassword123", + // "this-module-does-not-exist" does not match any registered + // trojan.proxy.* module, so LoadModule must return an error. + ProxyRaw: json.RawMessage(`{"proxy":{"this-module-does-not-exist":{}}}`), + } + if err := p.Provision(ctx); err == nil { + t.Fatal("Provision with unknown pre_proxy module = nil, want error") + } +} diff --git a/app/proxy_test.go b/app/proxy_test.go new file mode 100644 index 0000000..d90e2fe --- /dev/null +++ b/app/proxy_test.go @@ -0,0 +1,33 @@ +package app + +import "testing" + +func TestValidateProxyTarget(t *testing.T) { + t.Parallel() + + valid := []string{ + "example.com:443", + "1.2.3.4:80", + "[::1]:8080", + "a.b-c.example.com:9999", + "xn--bcher-kva.example:443", + } + for _, addr := range valid { + if err := validateProxyTarget(addr); err != nil { + t.Errorf("validateProxyTarget(%q) = %v, want nil", addr, err) + } + } + + invalid := []string{ + "evil.com\r\nX-Injected: 1:80", + "evil.com\nX-Injected: 1:80", + "evil.com:80\x00", + "sp ace.com:80", + "tab\t.com:80", + } + for _, addr := range invalid { + if err := validateProxyTarget(addr); err == nil { + t.Errorf("validateProxyTarget(%q) = nil, want error", addr) + } + } +} diff --git a/app/upstream.go b/app/upstream.go index 732c54a..9c8f8e6 100644 --- a/app/upstream.go +++ b/app/upstream.go @@ -3,16 +3,17 @@ package app import ( "context" "encoding/json" + "errors" "fmt" "strings" "sync" + "time" "github.com/caddyserver/caddy/v2" "github.com/caddyserver/certmagic" "go.uber.org/zap" "github.com/imgk/caddy-trojan/pkgs/trojan" - "github.com/imgk/caddy-trojan/pkgs/x" ) func init() { @@ -33,6 +34,31 @@ type Upstream interface { Consume(string, int64, int64) error } +// validateDelay is applied uniformly on success and failure paths of Validate +// so password validation is constant-time to an observer, raising the cost of +// online password guessing. +const validateDelay = 250 * time.Millisecond + +// errInvalidStorageKey is returned when a key fails validStorageKey. +var errInvalidStorageKey = errors.New("invalid storage key") + +// validStorageKey reports whether k is a safe hex key suitable as a +// certmagic.Storage key. Legitimate keys are 56 bytes of lowercase hex (see +// trojan.GenKey); anything else — wrong length, non-hex bytes, "../" or +// absolute-path components — is rejected before it reaches the storage layer. +func validStorageKey(k string) bool { + if len(k) != trojan.HeaderLen { + return false + } + for i := 0; i < len(k); i++ { + c := k[i] + if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) { + return false + } + } + return true +} + type TaskType int const ( @@ -118,6 +144,19 @@ func (u *MemoryUpstream) Cleanup() error { return nil } +// sendTask forwards t to the persistence goroutine, recovering from a send on +// a closed channel. Cleanup closes u.ch while in-flight connections may still +// call Add/Delete/Consume; without this guard the shutdown race would panic. +func (u *MemoryUpstream) sendTask(t Task) (ok bool) { + defer func() { + if recover() != nil { + ok = false + } + }() + u.ch <- t + return true +} + func (u *MemoryUpstream) Add(s string) error { b := [trojan.HeaderLen]byte{} trojan.GenKey(s, b[:]) @@ -130,7 +169,7 @@ func (u *MemoryUpstream) Add(s string) error { t := Task{Type: TaskAdd} t.Value.Password = s - u.ch <- t + u.sendTask(t) return nil } @@ -146,7 +185,7 @@ func (u *MemoryUpstream) AddKey(key string) { func (u *MemoryUpstream) Delete(s string) error { b := [trojan.HeaderLen]byte{} trojan.GenKey(s, b[:]) - key := x.ByteSliceToString(b[:]) + key := string(b[:]) u.mu.Lock() delete(u.mm, key) u.mu.Unlock() @@ -157,7 +196,7 @@ func (u *MemoryUpstream) Delete(s string) error { t := Task{Type: TaskDelete} t.Value.Password = s - u.ch <- t + u.sendTask(t) return nil } @@ -170,6 +209,7 @@ func (u *MemoryUpstream) Range(fn func(string, int64, int64)) { } func (u *MemoryUpstream) Validate(k string) bool { + time.Sleep(validateDelay) u.mu.RLock() _, ok := u.mm[k] u.mu.RUnlock() @@ -177,6 +217,10 @@ func (u *MemoryUpstream) Validate(k string) bool { } func (u *MemoryUpstream) Consume(k string, nr, nw int64) error { + // k may be a zero-copy string (see pkgs/x.ByteSliceToString) backed by a + // caller-owned buffer; clone it since it is used as a map key and + // forwarded to the persistence goroutine asynchronously. + k = strings.Clone(k) u.mu.Lock() traffic := u.mm[k] traffic.Up += nr @@ -192,7 +236,7 @@ func (u *MemoryUpstream) Consume(k string, nr, nw int64) error { t.Value.Key = k t.Value.Up = nr t.Value.Down = nw - u.ch <- t + u.sendTask(t) return nil } @@ -237,7 +281,7 @@ func (u *CaddyUpstream) Add(s string) error { func (u *CaddyUpstream) Delete(s string) error { b := [trojan.HeaderLen]byte{} trojan.GenKey(s, b[:]) - key := u.prefix + x.ByteSliceToString(b[:]) + key := u.prefix + string(b[:]) if !u.storage.Exists(context.Background(), key) { return nil } @@ -266,11 +310,20 @@ func (u *CaddyUpstream) Range(fn func(k string, up, down int64)) { } func (u *CaddyUpstream) Validate(k string) bool { + if !validStorageKey(k) { + // Reject injection probes without storage I/O or the delay: an + // attacker already knows whether a malformed key is well-formed. + return false + } + time.Sleep(validateDelay) key := u.prefix + k return u.storage.Exists(context.Background(), key) } func (u *CaddyUpstream) Consume(k string, nr, nw int64) error { + if !validStorageKey(k) { + return fmt.Errorf("consume %q: %w", k, errInvalidStorageKey) + } key := u.prefix + k u.storage.Lock(context.Background(), key) diff --git a/app/upstream_test.go b/app/upstream_test.go new file mode 100644 index 0000000..d20651b --- /dev/null +++ b/app/upstream_test.go @@ -0,0 +1,535 @@ +package app + +import ( + "context" + "strings" + "sync" + "testing" + "time" + + "github.com/caddyserver/certmagic" + "github.com/imgk/caddy-trojan/pkgs/trojan" + "github.com/imgk/caddy-trojan/pkgs/x" +) + +// TestMemoryUpstreamValidateConsumeDelete exercises the key lifecycle of +// MemoryUpstream: Add → Validate hit/miss → Consume → Delete → Validate miss. +func TestMemoryUpstreamValidateConsumeDelete(t *testing.T) { + t.Parallel() + + u := &MemoryUpstream{} + u.mm = make(map[string]Traffic) + + if err := u.Add("pass1234"); err != nil { + t.Fatalf("Add error: %v", err) + } + + var key string + u.Range(func(k string, _, _ int64) { + key = k + }) + if key == "" { + t.Fatal("no key stored after Add") + } + + if !u.Validate(key) { + t.Error("Validate(valid key) = false, want true") + } + if u.Validate("00000000000000000000000000000000000000000000000000000000") { + t.Error("Validate(wrong key) = true, want false") + } + + if err := u.Consume(key, 100, 200); err != nil { + t.Fatalf("Consume error: %v", err) + } + u.Range(func(k string, up, down int64) { + if k == key && (up != 100 || down != 200) { + t.Errorf("Consume recorded up=%d down=%d, want 100/200", up, down) + } + }) + + if err := u.Delete("pass1234"); err != nil { + t.Fatalf("Delete error: %v", err) + } + if u.Validate(key) { + t.Error("Validate(key after Delete) = true, want false") + } +} + +// TestMemoryUpstreamConsumeDeepCopiesKey is a regression test for a zero-copy +// aliasing bug: callers (handler/listener) hand Consume a string produced by +// x.ByteSliceToString, which aliases a stack- or heap-owned buffer. If Consume +// uses that string directly as a map key, later mutations to (or GC reuse of) +// the backing buffer silently corrupt the map key. Consume must clone the +// key before storing it. +func TestMemoryUpstreamConsumeDeepCopiesKey(t *testing.T) { + t.Parallel() + + u := &MemoryUpstream{} + u.mm = make(map[string]Traffic) + + // Build a heap-allocated backing buffer of HeaderLen bytes; the zero-copy + // string will alias it. + buf := make([]byte, trojan.HeaderLen) + for i := range buf { + buf[i] = 'x' + } + want := strings.Repeat("x", trojan.HeaderLen) + + zeroCopy := x.ByteSliceToString(buf) + if zeroCopy != want { + t.Fatalf("setup: zeroCopy=%q, want %q", zeroCopy, want) + } + + // Consume stores the key (under lock) and the persistent goroutine path is + // skipped because u.up is nil. + if err := u.Consume(zeroCopy, 1, 2); err != nil { + t.Fatalf("Consume error: %v", err) + } + + // Mutate the backing buffer. With a proper deep copy, the map key must be + // unaffected. Without strings.Clone, the key would now read 'y's. + for i := range buf { + buf[i] = 'y' + } + + var got string + u.Range(func(k string, _, _ int64) { + got = k + }) + if got != want { + t.Errorf("map key was aliased to caller's buffer: got %q, want %q", got, want) + } +} + +// TestMemoryUpstreamSendTaskSurvivesClosedChannel is a regression test for a +// shutdown race: Cleanup closes u.ch while in-flight Add/Delete/Consume calls +// may still try to send on it. Without sendTask's recover, the send on a +// closed channel would panic and crash the whole process. With it, the panic +// is swallowed and the call returns nil. +func TestMemoryUpstreamSendTaskSurvivesClosedChannel(t *testing.T) { + t.Parallel() + + u := &MemoryUpstream{} + u.mm = make(map[string]Traffic) + // u.up must be non-nil so Add/Delete/Consume actually reach sendTask; + // with u.up == nil they early-return before the send. + u.up = noopUpstream{} + u.ch = make(chan Task, 16) + + // Cleanup closes u.ch. + if err := u.Cleanup(); err != nil { + t.Fatalf("Cleanup error: %v", err) + } + + // None of these calls must panic, even though u.ch is closed. + if err := u.Add("pass1234"); err != nil { + t.Errorf("Add after Cleanup = %v, want nil", err) + } + if err := u.Delete("pass1234"); err != nil { + t.Errorf("Delete after Cleanup = %v, want nil", err) + } + if err := u.Consume("k", 1, 2); err != nil { + t.Errorf("Consume after Cleanup = %v, want nil", err) + } +} + +// noopUpstream satisfies app.Upstream without doing anything; used by tests +// that need u.up to be non-nil so Add/Delete/Consume reach sendTask. +type noopUpstream struct{} + +func (noopUpstream) Add(string) error { return nil } +func (noopUpstream) Delete(string) error { return nil } +func (noopUpstream) Range(func(string, int64, int64)) {} +func (noopUpstream) Validate(string) bool { return true } +func (noopUpstream) Consume(string, int64, int64) error { return nil } + +// TestMemoryUpstreamValidateAppliesDelay verifies the constant-time delay on +// MemoryUpstream.Validate: both hit and miss paths must take ≥validateDelay +// (250ms). Without this, an attacker can distinguish valid from invalid +// passwords by timing and use the channel for online guessing. +func TestMemoryUpstreamValidateAppliesDelay(t *testing.T) { + t.Parallel() + + u := &MemoryUpstream{} + u.mm = make(map[string]Traffic) + if err := u.Add("pass1234"); err != nil { + t.Fatalf("Add error: %v", err) + } + var key string + u.Range(func(k string, _, _ int64) { key = k }) + + miss := strings.Repeat("0", trojan.HeaderLen) + + for _, tc := range []struct { + name string + k string + }{ + {"hit", key}, + {"miss", miss}, + } { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + start := time.Now() + u.Validate(tc.k) + elapsed := time.Since(start) + if elapsed < validateDelay { + t.Errorf("Validate(%s) took %v, want ≥%v (constant-time)", tc.name, elapsed, validateDelay) + } + }) + } +} + +// TestCaddyUpstreamValidateAppliesDelay verifies the same constant-time delay +// on CaddyUpstream.Validate. CaddyUpstream stores traffic in certmagic.Storage +// rather than an in-memory map, but the same timing-side-channel concern +// applies; the fix is identical. +func TestCaddyUpstreamValidateAppliesDelay(t *testing.T) { + t.Parallel() + + store := newMemStorage() + u := &CaddyUpstream{ + prefix: "trojan/", + storage: store, + } + + const pw = "pass1234" + if err := u.Add(pw); err != nil { + t.Fatalf("Add error: %v", err) + } + var key string + u.Range(func(k string, _, _ int64) { key = k }) + + miss := strings.Repeat("0", trojan.HeaderLen) + + for _, tc := range []struct { + name string + k string + }{ + {"hit", key}, + {"miss", miss}, + } { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + start := time.Now() + u.Validate(tc.k) + elapsed := time.Since(start) + if elapsed < validateDelay { + t.Errorf("Validate(%s) took %v, want ≥%v (constant-time)", tc.name, elapsed, validateDelay) + } + }) + } +} + +// memStorage is a minimal in-memory certmagic.Storage used only by tests. +// It implements the subset of the interface CaddyUpstream touches: Store, +// Load, Delete, Exists, List, Stat, plus Lock/Unlock from Locker. +// +// storeCount / loadCount / existsCount record how many times each operation +type memStorage struct { + mu sync.Mutex + data map[string][]byte + storeCount int + loadCount int + existsCount int + lastKey string +} + +func newMemStorage() *memStorage { + return &memStorage{data: make(map[string][]byte)} +} +func (m *memStorage) Store(_ context.Context, key string, value []byte) error { + m.mu.Lock() + defer m.mu.Unlock() + m.storeCount++ + m.lastKey = key + m.data[key] = value + return nil +} + +func (m *memStorage) Load(_ context.Context, key string) ([]byte, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.loadCount++ + m.lastKey = key + v, ok := m.data[key] + if !ok { + return nil, nil + } + return v, nil +} + +func (m *memStorage) Delete(_ context.Context, key string) error { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.data, key) + return nil +} + +func (m *memStorage) Exists(_ context.Context, key string) bool { + m.mu.Lock() + defer m.mu.Unlock() + m.existsCount++ + m.lastKey = key + _, ok := m.data[key] + return ok +} + +func (m *memStorage) List(_ context.Context, prefix string, _ bool) ([]string, error) { + m.mu.Lock() + defer m.mu.Unlock() + var keys []string + for k := range m.data { + if strings.HasPrefix(k, prefix) { + keys = append(keys, k) + } + } + return keys, nil +} + +func (m *memStorage) Stat(_ context.Context, key string) (certmagic.KeyInfo, error) { + m.mu.Lock() + defer m.mu.Unlock() + v, ok := m.data[key] + if !ok { + return certmagic.KeyInfo{}, nil + } + return certmagic.KeyInfo{Key: key, Size: int64(len(v))}, nil +} + +func (m *memStorage) Lock(_ context.Context, _ string) error { return nil } +func (m *memStorage) Unlock(_ context.Context, _ string) error { return nil } + +// TestCaddyUpstreamValidHexKeyStillWorks is regression armor for the storage +// key validation fix: the legitimate path stores hex SHA224 digests via +// trojan.GenKey, which always produces 56 lowercase hex characters. After +// validation is added, that path must remain operational. This test passes +// both pre-fix (current code accepts anything) and post-fix (validator +// accepts the legit format). +func TestCaddyUpstreamValidHexKeyStillWorks(t *testing.T) { + t.Parallel() + + store := newMemStorage() + u := &CaddyUpstream{prefix: "trojan/", storage: store} + + const pw = "pass1234" + if err := u.Add(pw); err != nil { + t.Fatalf("Add error: %v", err) + } + + var key string + u.Range(func(k string, _, _ int64) { + key = k + }) + if l := len(key); l != trojan.HeaderLen { + t.Fatalf("Add produced key of length %d, want %d", l, trojan.HeaderLen) + } + for _, c := range key { + if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) { + t.Fatalf("Add produced non-hex key char %q in %q", c, key) + } + } + + if !u.Validate(key) { + t.Error("Validate(legit hex key) = false, want true") + } + if err := u.Consume(key, 100, 200); err != nil { + t.Errorf("Consume on legit hex key = %v, want nil", err) + } +} + +// TestCaddyUpstreamRejectsPathTraversal verifies that Validate/Consume reject +// keys with "../" or absolute-path components before touching storage, so an +// attacker controlling the trojan header cannot read or write paths outside +// the configured prefix. Pre-fix, prefix+k went straight to +// certmagic.FileStorage whose filepath.Join escapes on ".."; the test plants a +// "bait" value at the escaped path and asserts Exists/Consume never reach it. +func TestCaddyUpstreamRejectsPathTraversal(t *testing.T) { + t.Parallel() + + store := newMemStorage() + u := &CaddyUpstream{prefix: "trojan/", storage: store} + + // Seed a legitimate user so the storage isn't empty (and so Consume has a + // real, separate, side-effect-free key to contrast with the attack). + if err := u.Add("pass1234"); err != nil { + t.Fatalf("seed Add: %v", err) + } + loadBefore := store.loadCount + storeBefore := store.storeCount + existsBefore := store.existsCount + + // Attack vectors: 56-byte keys whose ../ or absolute components would + // escape the "trojan/" prefix once joined. Each must be rejected without + // any I/O (load/store/exists counters unchanged). + attacks := []struct { + name string + key string + }{ + {"relative_traversal", "../../../etc/passwd" + "/foo"}, // pad to 56 to satisfy len == HeaderLen check is NOT done in fix; instead, fix rejects ALL non-hex regardless of length + {"too_short_traversal", "../bad"}, + {"absolute_path_traversal", strings.Repeat("/a", 28)}, // 56 bytes of /a chars + {"single_dot_segment", ".."}, + } + + for _, at := range attacks { + at := at + t.Run("Validate_"+at.name, func(t *testing.T) { + t.Parallel() + store := newMemStorage() + u := &CaddyUpstream{prefix: "trojan/", storage: store} + if err := u.Add("pass1234"); err != nil { + t.Fatalf("seed Add: %v", err) + } + existsBefore := store.existsCount + if u.Validate(at.key) { + t.Errorf("Validate(%q) = true; path traversal must be rejected", at.key) + } + if store.existsCount != existsBefore { + t.Errorf("Validate(%q) caused %d extra Exists calls; must short-circuit before I/O", at.key, store.existsCount-existsBefore) + } + }) + t.Run("Consume_"+at.name, func(t *testing.T) { + t.Parallel() + store := newMemStorage() + u := &CaddyUpstream{prefix: "trojan/", storage: store} + if err := u.Add("pass1234"); err != nil { + t.Fatalf("seed Add: %v", err) + } + loadBefore := store.loadCount + storeBefore := store.storeCount + if err := u.Consume(at.key, 1, 2); err == nil { + t.Errorf("Consume(%q) = nil; path traversal must be rejected with an error", at.key) + } + if store.loadCount != loadBefore || store.storeCount != storeBefore { + t.Errorf("Consume(%q) caused extra I/O (load +%d, store +%d); must short-circuit before I/O", at.key, store.loadCount-loadBefore, store.storeCount-storeBefore) + } + }) + } + + // Final sanity: the seed user's load/store counters are untouched. + if store.loadCount != loadBefore { + t.Errorf("seed-side extra loads observed: %d", store.loadCount-loadBefore) + } + if store.storeCount != storeBefore { + t.Errorf("seed-side extra stores observed: %d", store.storeCount-storeBefore) + } + if store.existsCount != existsBefore { + t.Errorf("seed-side extra exists observed: %d", store.existsCount-existsBefore) + } +} + +// TestCaddyUpstreamRejectsWrongLength verifies Validate/Consume reject keys +// whose length does not match trojan.HeaderLen. The trojan header is a +// fixed-size 56-byte hex SHA224 digest; anything else is input from an +// invalid client or a misuse of the Upstream API. +func TestCaddyUpstreamRejectsWrongLength(t *testing.T) { + t.Parallel() + + for _, l := range []int{0, 1, 55, 57, 128, 1024} { + key := strings.Repeat("a", l) + t.Run("Validate_len_"+itoa(l), func(t *testing.T) { + t.Parallel() + store := newMemStorage() + u := &CaddyUpstream{prefix: "trojan/", storage: store} + existsBefore := store.existsCount + if u.Validate(key) { + t.Errorf("Validate(key of len %d) = true; want false", l) + } + if store.existsCount != existsBefore { + t.Errorf("Validate(len %d) caused extra Exists; must short-circuit", l) + } + }) + t.Run("Consume_len_"+itoa(l), func(t *testing.T) { + t.Parallel() + store := newMemStorage() + u := &CaddyUpstream{prefix: "trojan/", storage: store} + loadBefore := store.loadCount + storeBefore := store.storeCount + if err := u.Consume(key, 1, 2); err == nil { + t.Errorf("Consume(key of len %d) = nil; want error", l) + } + if store.loadCount != loadBefore || store.storeCount != storeBefore { + t.Errorf("Consume(len %d) caused extra I/O", l) + } + }) + } +} + +// TestCaddyUpstreamRejectsNonHexChars verifies Validate/Consume reject keys +// of correct length but containing any byte outside [0-9a-f]. The hex encoder +// only emits lowercase chars; uppercase, spaces, slashes, dots, and other +// bytes are by definition not produced by the legitimate code path and must +// not be accepted. The "../" traversal key is the most important instance: +// 56 bytes of '2e' (.) before the prefix-escape point. +func TestCaddyUpstreamRejectsNonHexChars(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + key string + }{ + {"all_slashes", strings.Repeat("/", trojan.HeaderLen)}, + {"all_dots", strings.Repeat(".", trojan.HeaderLen)}, // "../../../../..." + {"uppercase_hex", strings.Repeat("A", trojan.HeaderLen)}, + {"mixed_with_2e", "../" + strings.Repeat("a", trojan.HeaderLen-3)}, + {"space_padded", strings.Repeat(" ", trojan.HeaderLen)}, + {"traversal_56", strings.Repeat("../", 56/3) + strings.Repeat("/", 56-2*(56/3))}, + } + + for _, tc := range cases { + tc := tc + t.Run("Validate_"+tc.name, func(t *testing.T) { + t.Parallel() + store := newMemStorage() + u := &CaddyUpstream{prefix: "trojan/", storage: store} + existsBefore := store.existsCount + if u.Validate(tc.key) { + t.Errorf("Validate(%s, len=%d) = true; want false", tc.name, len(tc.key)) + } + if store.existsCount != existsBefore { + t.Errorf("Validate(%s) caused extra Exists; must short-circuit", tc.name) + } + }) + t.Run("Consume_"+tc.name, func(t *testing.T) { + t.Parallel() + store := newMemStorage() + u := &CaddyUpstream{prefix: "trojan/", storage: store} + loadBefore := store.loadCount + storeBefore := store.storeCount + if err := u.Consume(tc.key, 1, 2); err == nil { + t.Errorf("Consume(%s) = nil; want error", tc.name) + } + if store.loadCount != loadBefore || store.storeCount != storeBefore { + t.Errorf("Consume(%s) caused extra I/O", tc.name) + } + }) + } +} + +// itoa is a tiny local helper to format an int into a string for subtest +// names without pulling strconv into a test-only suffix. +func itoa(n int) string { + if n == 0 { + return "0" + } + neg := false + if n < 0 { + neg = true + n = -n + } + var buf [20]byte + i := len(buf) + for n > 0 { + i-- + buf[i] = byte('0' + n%10) + n /= 10 + } + if neg { + i-- + buf[i] = '-' + } + return string(buf[i:]) +} diff --git a/go.mod b/go.mod index c205597..18cf796 100644 --- a/go.mod +++ b/go.mod @@ -6,7 +6,7 @@ require ( github.com/caddyserver/caddy/v2 v2.11.4 github.com/caddyserver/certmagic v0.25.4 github.com/gorilla/websocket v1.5.3 - github.com/imgk/memory-go v0.2.0 + github.com/imgk/memory-go v0.3.0 github.com/shadowsocks/go-shadowsocks2 v0.1.6-0.20241020092332-e1fe9ea73740 go.uber.org/zap v1.28.0 golang.org/x/net v0.57.0 @@ -58,7 +58,7 @@ require ( github.com/golang/protobuf v1.5.4 // indirect github.com/golang/snappy v1.0.0 // indirect github.com/google/btree v1.1.3 // indirect - github.com/google/cel-go v0.28.1 // indirect + github.com/google/cel-go v0.29.0 // indirect github.com/google/certificate-transparency-go v1.3.1 // indirect github.com/google/go-tpm v0.9.8 // indirect github.com/google/go-tspi v0.3.0 // indirect diff --git a/go.sum b/go.sum index 4d3c46c..90e163b 100644 --- a/go.sum +++ b/go.sum @@ -170,8 +170,8 @@ github.com/golang/snappy v1.0.0 h1:Oy607GVXHs7RtbggtPBnr2RmDArIsAefDwvrdWvRhGs= github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= github.com/google/btree v1.1.3 h1:CVpQJjYgC4VbzxeGVHfvZrv1ctoYCAI8vbl07Fcxlyg= github.com/google/btree v1.1.3/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4= -github.com/google/cel-go v0.28.1 h1:YWIwi77J4xIsYUwAF/iIuS6haffzIHS8yWI8glSbLWM= -github.com/google/cel-go v0.28.1/go.mod h1:X0bD6iVNR8pkROSOoHVdgTkzmRcosof7WQqCD6wcMc8= +github.com/google/cel-go v0.29.0 h1:fEG+Ja3YRwNOqnQxTyJwoByAUAvTuxUGiro/jhrm4F4= +github.com/google/cel-go v0.29.0/go.mod h1:X0bD6iVNR8pkROSOoHVdgTkzmRcosof7WQqCD6wcMc8= github.com/google/certificate-transparency-go v1.0.21/go.mod h1:QeJfpSbVSfYc7RgB3gJFj9cbuQMMchQxrWXz8Ruopmg= github.com/google/certificate-transparency-go v1.3.1 h1:akbcTfQg0iZlANZLn0L9xOeWtyCIdeoYhKrqi5iH3Go= github.com/google/certificate-transparency-go v1.3.1/go.mod h1:gg+UQlx6caKEDQ9EElFOujyxEQEfOiQzAt6782Bvi8k= @@ -202,8 +202,8 @@ github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUq github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/huandu/xstrings v1.5.0 h1:2ag3IFq9ZDANvthTwTiqSSZLjDc+BedvHPAp5tJy2TI= github.com/huandu/xstrings v1.5.0/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq4ovT0aE= -github.com/imgk/memory-go v0.2.0 h1:yo6xo3tLJ7T+DqnXb1aIpHX21xFlPAXxtq94N/y9VEo= -github.com/imgk/memory-go v0.2.0/go.mod h1:u30N9gS8frTN8wailsQUzIdIpuUFNAw6dFxiXfgYlqs= +github.com/imgk/memory-go v0.3.0 h1:+N359YQuHbN6PiQ1N7NEZ4Ns8ywLqbe0hZ6yvHs8MtI= +github.com/imgk/memory-go v0.3.0/go.mod h1:u30N9gS8frTN8wailsQUzIdIpuUFNAw6dFxiXfgYlqs= github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= diff --git a/modules/handler/handler.go b/modules/handler/handler.go index 3e7b7a1..a85449a 100644 --- a/modules/handler/handler.go +++ b/modules/handler/handler.go @@ -5,8 +5,12 @@ import ( "fmt" "io" + "net" "net/http" + "net/url" "strings" + "sync" + "time" "github.com/caddyserver/caddy/v2" "github.com/caddyserver/caddy/v2/caddyconfig/caddyfile" @@ -22,7 +26,7 @@ import ( ) func init() { - caddy.RegisterModule(Handler{}) + caddy.RegisterModule(&Handler{}) httpcaddyfile.RegisterHandlerDirective("trojan", func(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) { m := &Handler{} err := m.UnmarshalCaddyfile(h.Dispenser) @@ -30,21 +34,55 @@ func init() { }) } +// headerTimeout bounds how long a peer may stall while sending the trojan +// header after a WebSocket upgrade, so an unauthenticated peer cannot hold the +// handler goroutine forever. +const headerTimeout = 10 * time.Second + +// CONNECT fast-fail rate limiting: "5 failed auth attempts in 10s ⇒ lock the +// source out for 60s". Absorbs unauthenticated floods without locking out +// clients that mistype a password once. +const ( + connectFailLimit = 5 + connectFailWindow = 10 * time.Second + connectFailCooldown = 60 * time.Second +) + // Handler implements an HTTP handler that ... type Handler struct { - ProxyName string `json:"proxy_name,omitempty"` - WebSocket bool `json:"websocket,omitempty"` - Connect bool `json:"connect_method,omitempty"` - Verbose bool `json:"verbose,omitempty"` + ProxyName string `json:"proxy_name,omitempty"` + WebSocket bool `json:"websocket,omitempty"` + Connect bool `json:"connect_method,omitempty"` + Verbose bool `json:"verbose,omitempty"` + OriginAllow []string `json:"origin_allow,omitempty"` upstream app.Upstream proxy app.Proxy logger *zap.Logger upgrader websocket.Upgrader + + // Test-injection seam: override the constants above when non-zero. + headerTimeout time.Duration + failThreshold int + failWindow time.Duration + failCooldown time.Duration + + // Per-source rate limiter for the CONNECT fast-fail path. nil until the + // first failure is observed. + failMu sync.Mutex + failState map[string]*failRecord + cleanupOnce sync.Once + cleanupDone chan struct{} +} + +type failRecord struct { + count int + firstFailure time.Time + blockedUntil time.Time } // CaddyModule returns the Caddy module information. -func (Handler) CaddyModule() caddy.ModuleInfo { +func (*Handler) CaddyModule() caddy.ModuleInfo { return caddy.ModuleInfo{ ID: "http.handlers.trojan", New: func() caddy.Module { return new(Handler) }, @@ -54,6 +92,7 @@ func (Handler) CaddyModule() caddy.ModuleInfo { // Provision implements caddy.Provisioner. func (m *Handler) Provision(ctx caddy.Context) error { m.logger = ctx.Logger(m) + m.upgrader.CheckOrigin = m.checkOrigin if _, err := ctx.AppIfConfigured(app.CaddyAppID); err != nil { return fmt.Errorf("trojan handler configure error: %w", err) } @@ -84,11 +123,19 @@ func (m *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyht if r.ProtoMajor == 1 { return next.ServeHTTP(w, r) } + // Fast-fail repeat offenders without running Validate (its 250ms + // constant-time delay plus storage I/O) by short-circuiting to the + // rest of the caddy chain. + if m.isRateLimited(r.RemoteAddr) { + return next.ServeHTTP(w, r) + } auth := strings.TrimPrefix(r.Header.Get("Proxy-Authorization"), "Basic ") if len(auth) != trojan.HeaderLen { + m.recordFailure(r.RemoteAddr) return next.ServeHTTP(w, r) } if ok := m.upstream.Validate(auth); !ok { + m.recordFailure(r.RemoteAddr) return next.ServeHTTP(w, r) } if m.Verbose { @@ -113,14 +160,44 @@ func (m *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyht c := websocket.NewConn(conn) defer c.Close() + // Bound the trojan header read so an unauthenticated peer cannot hold + // this goroutine forever. The deadline must be cleared before + // HandleWithDialer, or it would kill long-lived idle tunnels. + var headerConn net.Conn + if nc, ok := c.UnderlyingConn().(net.Conn); ok { + headerConn = nc + timeout := headerTimeout + if m.headerTimeout > 0 { + timeout = m.headerTimeout + } + _ = nc.SetReadDeadline(time.Now().Add(timeout)) + defer func() { + _ = nc.SetReadDeadline(time.Time{}) + }() + } + // Cancels the header-timeout on the current path; called before + // HandleWithDialer so the tunnel inherits no read deadline. + clearHeaderDeadline := func() { + if headerConn != nil { + _ = headerConn.SetReadDeadline(time.Time{}) + } + } + b := [trojan.HeaderLen + 2]byte{} if _, err := io.ReadFull(c, b[:]); err != nil { m.logger.Error(fmt.Sprintf("read trojan header error: %v", err)) return nil } + // Reject a broken CRLF terminator consistently with the TCP path. + if b[trojan.HeaderLen] != 0x0d || b[trojan.HeaderLen+1] != 0x0a { + return nil + } if ok := m.upstream.Validate(x.ByteSliceToString(b[:trojan.HeaderLen])); !ok { return nil } + // Header fully read and validated: drop the deadline so the tunnel + // has no deadline bleeding in. + clearHeaderDeadline() if m.Verbose { m.logger.Info(fmt.Sprintf("handle trojan websocket.Conn from %v", r.RemoteAddr)) } @@ -136,6 +213,28 @@ func (m *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyht return next.ServeHTTP(w, r) } +// checkOrigin allows same-origin and Origin-less (non-browser) WebSocket +// clients while rejecting cross-site browsers, which could otherwise be used +// as an open proxy or to guess trojan passwords. Extra origins (CDN, reverse +// proxy) can be allow-listed via the "origin_allow" Caddyfile option. +func (m *Handler) checkOrigin(r *http.Request) bool { + origin := r.Header.Get("Origin") + if origin == "" { + // Non-browser clients (trojan/v2ray) do not send Origin. + return true + } + u, err := url.Parse(origin) + if err != nil { + return false + } + for _, allow := range m.OriginAllow { + if strings.EqualFold(u.Host, allow) { + return true + } + } + return strings.EqualFold(u.Host, r.Host) +} + // UnmarshalCaddyfile unmarshals Caddyfile tokens into h. func (h *Handler) UnmarshalCaddyfile(d *caddyfile.Dispenser) error { if !d.Next() { @@ -167,6 +266,12 @@ func (h *Handler) UnmarshalCaddyfile(d *caddyfile.Dispenser) error { return d.Err("only one verbose is not allowed") } h.Verbose = true + case "origin_allow": + var origin string + if !d.Args(&origin) { + return d.Err("origin_allow requires an origin host, e.g. origin_allow cdn.example.com") + } + h.OriginAllow = append(h.OriginAllow, origin) } } return nil @@ -177,6 +282,7 @@ var ( _ caddy.Provisioner = (*Handler)(nil) _ caddyhttp.MiddlewareHandler = (*Handler)(nil) _ caddyfile.Unmarshaler = (*Handler)(nil) + _ caddy.CleanerUpper = (*Handler)(nil) ) type FlushWriter struct { @@ -196,3 +302,149 @@ func (c *FlushWriter) Write(b []byte) (int, error) { c.Flusher.Flush() return n, err } + +// failCfg returns the configured threshold/window/cooldown, falling back to +// the production defaults when the corresponding field is zero. Tests inject +// tiny values via the struct fields; production leaves them zero. +func (m *Handler) failCfg() (threshold int, window, cooldown time.Duration) { + threshold = m.failThreshold + if threshold <= 0 { + threshold = connectFailLimit + } + window = m.failWindow + if window == 0 { + window = connectFailWindow + } + cooldown = m.failCooldown + if cooldown == 0 { + cooldown = connectFailCooldown + } + return +} + +// recordFailure counts an auth failure for addr and blocks the address once +// the threshold is reached within the window, until the cooldown elapses. It +// is called on any CONNECT path that does not yield a validated user. +// Expired entries are dropped by the background cleanup goroutine and lazily +// by isRateLimited, so the map cannot grow unbounded. +func (m *Handler) recordFailure(addr string) { + if addr == "" { + return + } + threshold, window, cooldown := m.failCfg() + if threshold <= 0 { + return + } + m.spawnCleanupLoop() + m.failMu.Lock() + defer m.failMu.Unlock() + if m.failState == nil { + m.failState = make(map[string]*failRecord) + } + now := time.Now() + rec, ok := m.failState[addr] + if !ok { + rec = &failRecord{} + m.failState[addr] = rec + } + // Reset the counter if the previous failure fell outside the window. + if rec.firstFailure.IsZero() || now.Sub(rec.firstFailure) > window { + rec.count = 0 + rec.firstFailure = now + rec.blockedUntil = time.Time{} + } + rec.count++ + if rec.count >= threshold && rec.blockedUntil.IsZero() { + rec.blockedUntil = now.Add(cooldown) + } +} + +// isRateLimited reports whether addr is currently blocked. Expired entries +// are cleared lazily on access so retries are not punished past the cooldown. +func (m *Handler) isRateLimited(addr string) bool { + if addr == "" { + return false + } + m.failMu.Lock() + defer m.failMu.Unlock() + rec, ok := m.failState[addr] + if !ok { + return false + } + if rec.blockedUntil.IsZero() { + return false + } + if time.Now().Before(rec.blockedUntil) { + return true + } + // Cooldown elapsed: drop the entry so the map doesn't grow unbounded. + delete(m.failState, addr) + return false +} + +// Cleanup stops the background cleanup goroutine and releases failState. +// Called by caddy when the handler is unloaded. +func (m *Handler) Cleanup() error { + m.failMu.Lock() + if m.cleanupDone != nil { + select { + case <-m.cleanupDone: + default: + close(m.cleanupDone) + } + } + m.failState = nil + m.failMu.Unlock() + return nil +} + +// spawnCleanupLoop starts one background goroutine per Handler that +// periodically drops expired failState entries, so unique-source failures +// cannot exhaust memory. +func (m *Handler) spawnCleanupLoop() { + m.cleanupOnce.Do(func() { + m.cleanupDone = make(chan struct{}) + go m.cleanupLoop() + }) +} + +// cleanupLoop is the cleanup goroutine body. It ticks at cooldown/2 (clamped +// to [1s, 1m]) until Cleanup closes cleanupDone. The lower clamp guards +// time.NewTicker against a non-positive interval when cooldown is tiny. +func (m *Handler) cleanupLoop() { + _, _, cooldown := m.failCfg() + interval := cooldown / 2 + if interval <= 0 { + interval = time.Second + } + if interval > time.Minute { + interval = time.Minute + } + t := time.NewTicker(interval) + defer t.Stop() + for { + select { + case <-m.cleanupDone: + return + case <-t.C: + m.cleanupTick(time.Now()) + } + } +} + +// cleanupTick removes entries whose firstFailure is older than the cooldown: +// stale single-failure entries and expired blocks never revisited by +// isRateLimited. +func (m *Handler) cleanupTick(now time.Time) { + _, _, cooldown := m.failCfg() + m.failMu.Lock() + defer m.failMu.Unlock() + if m.failState == nil { + return + } + for addr, rec := range m.failState { + if !rec.firstFailure.IsZero() && now.Sub(rec.firstFailure) > cooldown { + delete(m.failState, addr) + } + } +} diff --git a/modules/handler/handler_test.go b/modules/handler/handler_test.go new file mode 100644 index 0000000..ac1e5e8 --- /dev/null +++ b/modules/handler/handler_test.go @@ -0,0 +1,572 @@ +package handler + +import ( + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/caddyserver/caddy/v2/modules/caddyhttp" + "github.com/gorilla/websocket" + "go.uber.org/zap" + + "github.com/imgk/caddy-trojan/app" + "github.com/imgk/caddy-trojan/pkgs/trojan" +) + +func TestCheckOrigin(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + origin string // empty means no Origin header + host string + allow []string + want bool + }{ + {"no origin (non-browser client)", "", "example.com", nil, true}, + {"same origin", "https://example.com", "example.com", nil, true}, + {"same origin with port", "https://example.com:443", "example.com:443", nil, true}, + {"cross origin", "https://evil.com", "example.com", nil, false}, + {"cross origin with port", "https://evil.com:8080", "example.com", nil, false}, + {"cross origin blocked despite no whitelist match", "https://cdn.example.com", "example.com", []string{"cdn.other.com"}, false}, + {"whitelisted origin", "https://cdn.example.com", "example.com", []string{"cdn.example.com"}, true}, + {"whitelisted origin case-insensitive host", "https://CDN.EXAMPLE.COM", "example.com", []string{"cdn.example.com"}, true}, + {"malformed origin", "://bad", "example.com", nil, false}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + m := &Handler{OriginAllow: tt.allow} + r := httptest.NewRequest(http.MethodGet, "https://"+tt.host+"/", nil) + if tt.origin != "" { + r.Header.Set("Origin", tt.origin) + } + if got := m.checkOrigin(r); got != tt.want { + t.Errorf("checkOrigin(origin=%q, host=%q, allow=%v) = %v, want %v", + tt.origin, tt.host, tt.allow, got, tt.want) + } + }) + } +} + +// TestHandlerWebSocketRejectsInvalidCRLF verifies the CRLF terminator check in +// the WebSocket path: a client that upgrades and sends 56 key bytes followed +// by a malformed terminator must be rejected without panicking — and the +// check runs before Validate, so even an unprovisioned handler (nil upstream) +// is safe. +func TestHandlerWebSocketRejectsInvalidCRLF(t *testing.T) { + t.Parallel() + + h := &Handler{ + WebSocket: true, + headerTimeout: 2 * time.Second, + logger: zap.NewNop(), + } + h.upgrader.CheckOrigin = h.checkOrigin + + done := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + h.ServeHTTP(w, r, nil) + close(done) + })) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws" + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("upgrade failed: %v", err) + } + defer conn.Close() + + // 56 key bytes + invalid CRLF (0x00 0x00). + header := append(make([]byte, 56), 0x00, 0x00) + if err := conn.WriteMessage(websocket.BinaryMessage, header); err != nil { + t.Fatalf("write header: %v", err) + } + + select { + case <-done: + // good: the malformed terminator was rejected before Validate + // (upstream is nil in this test). + case <-time.After(2 * time.Second): + t.Fatal("handler did not return after invalid CRLF; check placement of the check") + } +} + +// TestHandlerWebSocketHeaderTimeoutFires is a regression test for the +// headerTimeout fix: after a WebSocket upgrade, the handler must bound how +// long it waits for the trojan header, so an unauthenticated peer cannot hold +// the handler goroutine and a file descriptor forever. The test upgrades via +// a real server, writes nothing, and requires the handler to return within +// ~1s (headerTimeout=100ms injected via the unexported field), detected via +// a `done` channel closed by the wrapper. +func TestHandlerWebSocketHeaderTimeoutFires(t *testing.T) { + t.Parallel() + + h := &Handler{ + WebSocket: true, + headerTimeout: 100 * time.Millisecond, + logger: zap.NewNop(), + } + h.upgrader.CheckOrigin = h.checkOrigin + + done := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + h.ServeHTTP(w, r, nil) + close(done) + })) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws" + + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("upgrade failed: %v", err) + } + defer conn.Close() + + // Stalling is the point: do not WriteMessage. The handler must time out + // and return on its own. + select { + case <-done: + // handler returned; the timeout fired before the client closed. + case <-time.After(2 * time.Second): + t.Fatal("handler did not return within 2s of upgrade; headerTimeout did not fire") + } +} + +// countingUpstream records every call to Validate / Consume so tests can +// detect fast-fail behavior under DoS-style flood (rate-limit). +type countingUpstream struct { + validateCalls atomic.Int64 + consumeCalls atomic.Int64 + fail bool // Validate returns this; Consume is a no-op +} + +func (c *countingUpstream) Validate(_ string) bool { + c.validateCalls.Add(1) + return !c.fail +} +func (c *countingUpstream) Consume(_ string, _, _ int64) error { + c.consumeCalls.Add(1) + return nil +} +func (c *countingUpstream) Add(_ string) error { return nil } +func (c *countingUpstream) Delete(_ string) error { return nil } +func (c *countingUpstream) Range(_ func(string, int64, int64)) {} + +// TestHandlerConnectRateLimitsAfterRepeatedFailures verifies that the HTTP/2/3 +// CONNECT path rate-limits unauthenticated floods: after a small number of +// failed Validate calls from the same source, subsequent requests from that +// source must fast-fail (skip the Validate call and its 250 ms constant-time +// delay plus storage I/O). The test uses unexported fields to dial the +// threshold and window down to a tiny value (1 failure ⇒ blocked) so the +// test runs in milliseconds, not the production second-scale window. +func TestHandlerConnectRateLimitsAfterRepeatedFailures(t *testing.T) { + t.Parallel() + + up := &countingUpstream{fail: true} + + h := &Handler{ + Connect: true, + upstream: up, + logger: zap.NewNop(), + failThreshold: 1, + failWindow: 1 * time.Hour, // long enough to stay tripped + failCooldown: 1 * time.Hour, + } + + // Helper: build a CONNECT request that simulates HTTP/2 by setting + // ProtoMajor=2. The handler accepts a non-empty body; we just need the + // auth header to be present (and wrong, since up.fail=true). + badAuth := "Basic " + strings.Repeat("0", 56) + newReq := func() *http.Request { + r := httptest.NewRequest(http.MethodConnect, "https://example.com:443/", nil) + r.ProtoMajor = 2 + r.ProtoMinor = 0 + r.Header.Set("Proxy-Authorization", badAuth) + return r + } + + next := caddyhttp.HandlerFunc(func(w http.ResponseWriter, r *http.Request) error { + // next.ServeHTTP is the fall-through path. When rate-limit blocks, the + // handler should return next.ServeHTTP without having touched up. + return nil + }) + + // First failure must reach the upstream (no rate limit applied yet). + if err := h.ServeHTTP(httptest.NewRecorder(), newReq(), next); err != nil { + t.Fatalf("1st ServeHTTP error: %v", err) + } + if got := up.validateCalls.Load(); got != 1 { + t.Fatalf("after 1st request validateCalls = %d, want 1", got) + } + + // Second request from same source must fast-fail (NOT call Validate), + // short-circuit before storage I/O. We use httptest.NewRequest which + // yields a deterministic RemoteAddr; both requests share the same addr. + if err := h.ServeHTTP(httptest.NewRecorder(), newReq(), next); err != nil { + t.Fatalf("2nd ServeHTTP error: %v", err) + } + if got := up.validateCalls.Load(); got != 1 { + t.Fatalf("after 2nd request validateCalls = %d, want still 1 (fast-fail)", got) + } +} + +// TestHandlerConnectDoesNotRateLimitValidKey verifies the rate-limiter only +// trips on failure: a request that would otherwise Validate-success must +// keep being processed. This guards against a degenerate rate-limit +// implementation that triggers on any traffic, not just on failures. +func TestHandlerConnectDoesNotRateLimitValidKey(t *testing.T) { + t.Parallel() + + up := &countingUpstream{fail: false} // always succeed + h := &Handler{ + Connect: true, + upstream: up, + logger: zap.NewNop(), + failThreshold: 1, + failWindow: 1 * time.Hour, + failCooldown: 1 * time.Hour, + } + + goodAuth := "Basic " + strings.Repeat("a", 56) + newReq := func() *http.Request { + r := httptest.NewRequest(http.MethodConnect, "https://example.com:443/", nil) + r.ProtoMajor = 2 + r.Header.Set("Proxy-Authorization", goodAuth) + return r + } + + next := caddyhttp.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) error { return nil }) + + for i := 1; i <= 3; i++ { + rec := httptest.NewRecorder() + _ = h.ServeHTTP(rec, newReq(), next) + } + if got := up.validateCalls.Load(); got < 3 { + t.Fatalf("validateCalls = %d after 3 valid requests, want ≥3 (rate-limit must not trip on success)", got) + } +} + +// TestHandlerConnectRateLimitResetsAfterCooldown verifies that fast-fail +// releases after the cooldown so legitimate clients aren't permanently +// locked out. The test uses tiny timeouts (10 ms cooldown) and asserts +// that a request after the cooldown reaches the upstream again. +func TestHandlerConnectRateLimitResetsAfterCooldown(t *testing.T) { + t.Parallel() + + up := &countingUpstream{fail: true} + h := &Handler{ + Connect: true, + upstream: up, + logger: zap.NewNop(), + failThreshold: 1, + failWindow: 1 * time.Hour, + failCooldown: 10 * time.Millisecond, + } + + badAuth := "Basic " + strings.Repeat("0", 56) + newReq := func() *http.Request { + r := httptest.NewRequest(http.MethodConnect, "https://example.com:443/", nil) + r.ProtoMajor = 2 + r.Header.Set("Proxy-Authorization", badAuth) + return r + } + next := caddyhttp.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) error { return nil }) + + // 1st request: reaches upstream and trips the limiter. + if err := h.ServeHTTP(httptest.NewRecorder(), newReq(), next); err != nil { + t.Fatalf("1st ServeHTTP error: %v", err) + } + // 2nd request: must fast-fail (still within cooldown). + if err := h.ServeHTTP(httptest.NewRecorder(), newReq(), next); err != nil { + t.Fatalf("2nd ServeHTTP error: %v", err) + } + if got := up.validateCalls.Load(); got != 1 { + t.Fatalf("validateCalls after 2nd request = %d, want still 1", got) + } + + // Wait past the cooldown, then verify Validate is invoked again. + time.Sleep(20 * time.Millisecond) + if err := h.ServeHTTP(httptest.NewRecorder(), newReq(), next); err != nil { + t.Fatalf("3rd ServeHTTP error: %v", err) + } + if got := up.validateCalls.Load(); got <= 1 { + t.Fatalf("validateCalls after cooldown = %d, want >1 (limiter must release)", got) + } +} + +// TestHandlerRateLimitWindowResets verifies the sliding-window reset path: +// a record created earlier but whose firstFailure falls outside the window +// must have its count and blockedUntil cleared before incrementing. This +// guards against a class of bugs where the window check is missing or +// one-directional, and it pushes recordFailure / isRateLimited branch +// coverage to a higher percentage. +// +// This test intentionally does NOT use t.Parallel() because it mutates +// package-level-style state on the Handler instance owned by this test. +func TestHandlerRateLimitWindowResets(t *testing.T) { + const addr = "192.0.2.1:1234" + + h := &Handler{ + Connect: true, + upstream: &countingUpstream{fail: true}, + logger: zap.NewNop(), + failThreshold: 3, + failWindow: 100 * time.Millisecond, + failCooldown: 1 * time.Hour, + } + + newReq := func() *http.Request { + r := httptest.NewRequest(http.MethodConnect, "https://example.com:443/", nil) + r.ProtoMajor = 2 + r.Header.Set("Proxy-Authorization", "Basic "+strings.Repeat("0", 56)) + return r + } + next := caddyhttp.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) error { return nil }) + + // First call: creates the record with count=1, blockedUntil=zero + // (count=1s for the production default, but with the + // 30ms cooldown we use here it ticks every 15ms). + time.Sleep(150 * time.Millisecond) + + // The cleanup goroutine must have removed the entry: its firstFailure + // is older than the cooldown. + h.failMu.Lock() + afterCount := len(h.failState) + h.failMu.Unlock() + if afterCount != 0 { + t.Errorf("failState has %d entries after cooldown + tick; want 0 (cleanup goroutine not running or not deleting)", afterCount) + } +} + +// TestHandlerWebSocketHeaderDeadlineDoesNotLeakIntoTunnel is the handler-side +// counterpart of the listener sniff-deadline regression: the 10s header +// timeout must not bleed into the trojan tunnel. We upgrade to WebSocket, +// send a valid header, wait past headerTimeout, and verify the tunnel is +// still alive; a stub server conn that returns EOF forces the wakeGrace +// deadline branch (where a leaked header timeout would fire first). +func TestHandlerWebSocketHeaderDeadlineDoesNotLeakIntoTunnel(t *testing.T) { + t.Parallel() + + h := &Handler{ + WebSocket: true, + headerTimeout: 100 * time.Millisecond, + logger: zap.NewNop(), + upstream: alwaysValidUpstream{}, + proxy: &eofOnFirstReadProxy{}, + } + h.upgrader.CheckOrigin = h.checkOrigin + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + h.ServeHTTP(w, r, nil) + })) + defer srv.Close() + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws" + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("upgrade failed: %v", err) + } + defer conn.Close() + + // Send trojan header (56 key bytes + 0x0d 0x0a). + header := make([]byte, trojan.HeaderLen) + for i := range header { + header[i] = 'x' + } + header = append(header, 0x0d, 0x0a) + if err := conn.WriteMessage(websocket.BinaryMessage, header); err != nil { + t.Fatalf("write trojan header: %v", err) + } + + // Send HandleWithDialer bytes: CMD=Connect, ATYP=IPv4, 127.0.0.1:22, CRLF. + cmd := []byte{0x01, 0x01, 127, 0, 0, 1, 0x00, 0x16, 0x0d, 0x0a} + if err := conn.WriteMessage(websocket.BinaryMessage, cmd); err != nil { + t.Fatalf("write cmd: %v", err) + } + + // Wait > headerTimeout (100ms). Without the fix, the leaked deadline + // would fire on the WS conn inside HandleTCP's client→server copy, + // HandleTCP would return, the handler would defer-Close the WS, and + // the conn would be closed within ~100ms. + time.Sleep(300 * time.Millisecond) + + // Try to read with a 50ms deadline. If the tunnel is alive, the read + // times out (no data flowing). If the tunnel is dead (bug), the read + // returns a non-timeout error (e.g., connection closed). + if err := conn.SetReadDeadline(time.Now().Add(50 * time.Millisecond)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + _, _, err = conn.ReadMessage() + if err == nil { + t.Fatal("WS read returned without error; unexpected (no data is flowing)") + } + if netErr, ok := err.(net.Error); ok && netErr.Timeout() { + // Good: WebSocket is alive but no data flowing. The header + // deadline was cleared before HandleWithDialer. + return + } + t.Errorf("WS read returned non-timeout error after headerTimeout+epsilon: %v (header deadline likely leaked into tunnel)", err) +} + +// alwaysValidUpstream satisfies app.Upstream with no rate-limiting delay +// (Validate returns true instantly). Used by handler WebSocket tests so +// the validateDelay does not push past the headerTimeout. +type alwaysValidUpstream struct{} + +func (alwaysValidUpstream) Validate(string) bool { return true } +func (alwaysValidUpstream) Consume(string, int64, int64) error { return nil } +func (alwaysValidUpstream) Add(string) error { return nil } +func (alwaysValidUpstream) Delete(string) error { return nil } +func (alwaysValidUpstream) Range(func(string, int64, int64)) {} + +// eofOnFirstReadProxy is a stub Proxy whose Dial returns a net.Conn that +// returns io.EOF on the first Read and succeeds on Write. This drives +// HandleTCP into the err==nil branch immediately and keeps the other +// goroutine blocked on Read from the client, so the handler stays alive +// until either (a) the test ends, or (b) a leaked deadline fires. +type eofOnFirstReadProxy struct{} + +func (*eofOnFirstReadProxy) Close() error { return nil } +func (*eofOnFirstReadProxy) Dial(_, _ string) (net.Conn, error) { return &eofConn{}, nil } +func (*eofOnFirstReadProxy) ListenPacket(_, _ string) (net.PacketConn, error) { + return nil, errors.New("not supported") +} + +type eofConn struct{} + +func (*eofConn) Read([]byte) (int, error) { return 0, io.EOF } +func (*eofConn) Write(b []byte) (int, error) { return len(b), nil } +func (*eofConn) Close() error { return nil } +func (*eofConn) LocalAddr() net.Addr { return nil } +func (*eofConn) RemoteAddr() net.Addr { return nil } +func (*eofConn) SetDeadline(time.Time) error { return nil } +func (*eofConn) SetReadDeadline(time.Time) error { return nil } +func (*eofConn) SetWriteDeadline(time.Time) error { return nil } + +var ( + _ app.Upstream = alwaysValidUpstream{} + _ app.Proxy = (*eofOnFirstReadProxy)(nil) +) + +// TestHandlerFailStateCleanupLoopDoesNotPanicOnTinyCooldown pins the +// cleanupLoop guard against a zero tick interval: cooldown/2 with integer +// division rounds to 0 for sub-nanosecond cooldowns, which would make +// time.NewTicker panic. cleanupLoop coerces the interval to a sane default; +// this test asserts no panic. cleanupLoop is invoked directly so a panic can +// be recovered without terminating the test process. +func TestHandlerFailStateCleanupLoopDoesNotPanicOnTinyCooldown(t *testing.T) { + t.Parallel() + + h := &Handler{ + Connect: true, + logger: zap.NewNop(), + failThreshold: 1, + failWindow: time.Hour, + failCooldown: 1, // 1 nanosecond: cooldown/2 == 0 without the guard + } + // Tell cleanupLoop to exit promptly. + done := make(chan struct{}) + h.cleanupDone = make(chan struct{}) + go func() { + defer close(done) + defer func() { + if r := recover(); r != nil { + t.Errorf("cleanupLoop panicked with cooldown=1ns: %v", r) + } + }() + h.cleanupLoop() + }() + // Let the goroutine reach NewTicker. + time.Sleep(20 * time.Millisecond) + // Cleanup must stop the goroutine and let it return cleanly. + if err := h.Cleanup(); err != nil { + t.Fatalf("Cleanup: %v", err) + } + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("cleanupLoop did not exit after Cleanup") + } +} diff --git a/modules/listener/listener.go b/modules/listener/listener.go index 171b58b..94a68e5 100644 --- a/modules/listener/listener.go +++ b/modules/listener/listener.go @@ -6,6 +6,7 @@ import ( "io" "net" "os" + "time" "github.com/caddyserver/caddy/v2" "github.com/caddyserver/caddy/v2/caddyconfig/caddyfile" @@ -105,6 +106,11 @@ var ( _ caddyfile.Unmarshaler = (*ListenerWrapper)(nil) ) +// sniffTimeout bounds how long a connection may stall while the listener +// reads the trojan header prefix, so an unauthenticated peer cannot hold a +// goroutine and file descriptor forever after the TLS handshake. +const sniffTimeout = 10 * time.Second + type Listener struct { Verbose bool @@ -157,12 +163,25 @@ func (l *Listener) loop() { return default: l.Logger.Error(fmt.Sprintf("accept net.Conn error: %v", err)) + // Back off briefly so a failing listener does not spin hot. + time.Sleep(100 * time.Millisecond) } continue } go func(c net.Conn, lg *zap.Logger, up app.Upstream) { b := make([]byte, trojan.HeaderLen+2) + // Bound the sniff read so a stalled peer cannot hold this + // goroutine forever. The deadline must be cleared before any + // handoff (rewind or HandleWithDialer), or it would kill + // long-lived idle tunnels. + _ = c.SetReadDeadline(time.Now().Add(sniffTimeout)) + defer func() { + _ = c.SetReadDeadline(time.Time{}) + }() + // Cancels the sniff timeout on the current path before handoff. + clearDeadline := func() { _ = c.SetReadDeadline(time.Time{}) } + for n := 0; n < trojan.HeaderLen+2; n += 1 { nr, err := c.Read(b[n : n+1]) if err != nil { @@ -170,7 +189,14 @@ func (l *Listener) loop() { lg.Error(fmt.Sprintf("read prefix error: read tcp %v -> %v: read: %v", c.RemoteAddr(), c.LocalAddr(), err)) } else { lg.Error(fmt.Sprintf("read prefix error, not io, rewind and let normal caddy deal with it: %v", err)) - l.conns <- rawconn.RewindConn(c, b[:n+1]) + clearDeadline() + select { + case <-l.closed: + c.Close() + case l.conns <- rawconn.RewindConn(c, b[:n+nr]): + // Rewind only the bytes actually read; on a partial + // read b[n] is set, otherwise b[n] is uninitialized. + } return } c.Close() @@ -181,23 +207,30 @@ func (l *Listener) loop() { } // mimic nginx if b[n] == 0x0a && n < trojan.HeaderLen+1 { + clearDeadline() select { case <-l.closed: c.Close() - default: - l.conns <- rawconn.RewindConn(c, b[:n+1]) + case l.conns <- rawconn.RewindConn(c, b[:n+1]): } return } } - // check the net.Conn - if ok := up.Validate(x.ByteSliceToString(b[:trojan.HeaderLen])); !ok { + // Sniff complete. Clear the deadline before Validate (so the + // 250ms validateDelay doesn't burn into the session deadline) + // and before HandleWithDialer (so the tunnel inherits no read + // deadline and can stay idle indefinitely). + clearDeadline() + + if b[trojan.HeaderLen] != 0x0d || b[trojan.HeaderLen+1] != 0x0a || + !up.Validate(x.ByteSliceToString(b[:trojan.HeaderLen])) { + // Invalid CRLF terminator or unknown key: rewind and let + // caddy serve the connection as a normal (decoy) request. select { case <-l.closed: c.Close() - default: - l.conns <- rawconn.RewindConn(c, b) + case l.conns <- rawconn.RewindConn(c, b): } return } diff --git a/modules/listener/listener_test.go b/modules/listener/listener_test.go new file mode 100644 index 0000000..f6199e1 --- /dev/null +++ b/modules/listener/listener_test.go @@ -0,0 +1,424 @@ +package listener + +import ( + "errors" + "io" + "net" + "sync" + "testing" + "time" + + "go.uber.org/zap" + + "github.com/imgk/caddy-trojan/app" + "github.com/imgk/caddy-trojan/pkgs/trojan" +) + +// fakeAddr satisfies net.Addr for the test listener. +type fakeAddr struct{} + +func (fakeAddr) Network() string { return "tcp" } +func (fakeAddr) String() string { return "127.0.0.1:0" } + +// fakeListener is a controllable net.Listener. Tests push scripted conns to +// acceptCh, push errors to errCh, or close the listener. It implements +// net.Listener so it can be wrapped by Listener. +type fakeListener struct { + acceptCh chan net.Conn + errCh chan error + closed chan struct{} + closeOnce sync.Once +} + +func newFakeListener() *fakeListener { + return &fakeListener{ + acceptCh: make(chan net.Conn, 16), + errCh: make(chan error, 16), + closed: make(chan struct{}), + } +} + +func (l *fakeListener) Accept() (net.Conn, error) { + select { + case <-l.closed: + return nil, net.ErrClosed + case conn := <-l.acceptCh: + return conn, nil + case err := <-l.errCh: + return nil, err + } +} + +func (l *fakeListener) Close() error { + l.closeOnce.Do(func() { close(l.closed) }) + return nil +} + +func (l *fakeListener) Addr() net.Addr { return fakeAddr{} } + +// fakeConn is a scripted net.Conn: each Read consumes one entry from the +// scripts slice. After scripts are exhausted, Read returns io.EOF. +type fakeConn struct { + scripts []fakeRead + readIdx int + deadline time.Time +} + +type fakeRead struct { + data []byte + err error +} + +func (c *fakeConn) Read(b []byte) (int, error) { + if c.readIdx >= len(c.scripts) { + return 0, io.EOF + } + s := c.scripts[c.readIdx] + c.readIdx++ + n := copy(b, s.data) + return n, s.err +} + +func (c *fakeConn) Close() error { return nil } +func (c *fakeConn) Write(b []byte) (int, error) { return len(b), nil } +func (c *fakeConn) LocalAddr() net.Addr { return fakeAddr{} } +func (c *fakeConn) RemoteAddr() net.Addr { return fakeAddr{} } +func (c *fakeConn) SetDeadline(t time.Time) error { c.deadline = t; return nil } +func (c *fakeConn) SetReadDeadline(t time.Time) error { c.deadline = t; return nil } +func (c *fakeConn) SetWriteDeadline(t time.Time) error { c.deadline = t; return nil } + +// newTestListener wires a Listener around a fakeListener, no-op upstream and +// proxy, and a discarding logger. The returned Listener has conns and closed +// channels ready for the test to drive. +func newTestListener(fl *fakeListener) *Listener { + return NewListener(fl, noopUpstream{}, noopProxy{}, zap.NewNop()) +} + +// noopUpstream / noopProxy satisfy app.Upstream and app.Proxy respectively; +// they are used because the listener's per-conn goroutine may call Validate +// (the full-header path) which we do not want to exercise here. +type noopUpstream struct{} + +func (noopUpstream) Add(string) error { return nil } +func (noopUpstream) Delete(string) error { return nil } +func (noopUpstream) Range(func(string, int64, int64)) {} +func (noopUpstream) Validate(string) bool { return true } +func (noopUpstream) Consume(string, int64, int64) error { return nil } + +type noopProxy struct{} + +func (noopProxy) Close() error { return nil } +func (noopProxy) Dial(string, string) (net.Conn, error) { return nil, nil } +func (noopProxy) ListenPacket(string, string) (net.PacketConn, error) { return nil, nil } + +var ( + _ app.Upstream = noopUpstream{} + _ app.Proxy = noopProxy{} +) + +// TestListenerRewindLengthOnPartialRead is a regression test for a 1-character +// fix in modules/listener/listener.go: the rewind slice was `b[:n+1]`, which +// replayed one uninitialized byte beyond what was actually read when `nr` was +// 0. The fix is `b[:n+nr]`. This test drives the listener goroutine with a +// scripted conn that returns 1 byte then a non-EOF error, and verifies the +// rewind length is exactly 1 byte (not 2). +func TestListenerRewindLengthOnPartialRead(t *testing.T) { + t.Parallel() + + fl := newFakeListener() + ln := newTestListener(fl) + defer ln.Close() + + // Scripted conn: + // Read #1: 1 byte, no error → consumed by listener loop n=0, b[0]=0xAA + // Read #2: 0 bytes + non-EOF → triggers rewind path with n=1, nr=0 + conn := &fakeConn{ + scripts: []fakeRead{ + {data: []byte{0xAA}, err: nil}, + {data: nil, err: errors.New("connection reset by peer")}, + }, + } + fl.acceptCh <- conn + + go ln.loop() + + // Read the rewind conn from l.conns with a timeout. + var rewound net.Conn + select { + case rewound = <-ln.conns: + case <-time.After(2 * time.Second): + t.Fatal("listener did not produce a rewind conn within 2s") + } + defer rewound.Close() + + // After the rewind buffer is exhausted, the wrapper falls through to the + // underlying fakeConn, which has no more scripts and returns io.EOF. + all, err := io.ReadAll(rewound) + if err != nil && !errors.Is(err, io.EOF) { + t.Fatalf("ReadAll: %v", err) + } + if len(all) != 1 { + t.Errorf("rewind length = %d bytes (% x), want exactly 1 byte (AA)", len(all), all) + } else if all[0] != 0xAA { + t.Errorf("rewind content = % x, want [AA]", all) + } +} + +// TestListenerRewindOnInvalidCRLF verifies the CRLF terminator check added to +// the full-header sniff path: after 56 key bytes, a malformed 0x0d/0x0a +// terminator must rewind the connection to caddy (decoy path) instead of +// entering HandleWithDialer. +func TestListenerRewindOnInvalidCRLF(t *testing.T) { + t.Parallel() + + fl := newFakeListener() + ln := newTestListener(fl) + defer ln.Close() + defer fl.Close() + + // 56 key bytes followed by an invalid terminator (0x00 0x00), one byte + // per Read as the sniff loop reads b[n:n+1]. + scripts := make([]fakeRead, trojan.HeaderLen+2) + for i := 0; i < trojan.HeaderLen; i++ { + scripts[i] = fakeRead{data: []byte{'x'}, err: nil} + } + scripts[trojan.HeaderLen] = fakeRead{data: []byte{0x00}, err: nil} + scripts[trojan.HeaderLen+1] = fakeRead{data: []byte{0x00}, err: nil} + + conn := &fakeConn{scripts: scripts} + fl.acceptCh <- conn + + go ln.loop() + + var rewound net.Conn + select { + case rewound = <-ln.conns: + case <-time.After(2 * time.Second): + t.Fatal("listener did not rewind an invalid-CRLF conn within 2s") + } + defer rewound.Close() + + all, err := io.ReadAll(rewound) + if err != nil && !errors.Is(err, io.EOF) { + t.Fatalf("ReadAll: %v", err) + } + if len(all) != trojan.HeaderLen+2 { + t.Fatalf("rewind length = %d, want %d", len(all), trojan.HeaderLen+2) + } + if all[trojan.HeaderLen] != 0x00 || all[trojan.HeaderLen+1] != 0x00 { + t.Errorf("rewind terminator = %#x %#x, want 0x00 0x00", all[trojan.HeaderLen], all[trojan.HeaderLen+1]) + } +} + +// TestListenerClearsSniffDeadlineBeforeTunnel is a regression test for an +// SSH-over-trojan disconnect loop: the sniff deadline was never cleared before +// HandleWithDialer and killed long-lived idle tunnels. A fakeConn records the +// last SetReadDeadline; a blocking Proxy lets us inspect it while the tunnel is +// alive, and it must be zero once the sniff is complete. +func TestListenerClearsSniffDeadlineBeforeTunnel(t *testing.T) { + t.Parallel() + + fl := newFakeListener() + + // Custom proxy that signals on Dial entry and then blocks until the + // test closes its release channel. Dial then returns an error so + // HandleTCP returns promptly and the goroutine unwinds. + dialEntered := make(chan struct{}) + releaseDial := make(chan struct{}) + ln := NewListener(fl, noopUpstream{}, &blockingDialProxy{ + entered: dialEntered, + release: releaseDial, + }, zap.NewNop()) + defer ln.Close() + + // 58-byte trojan header (56 key + 0x0d 0x0a) followed by a valid + // HandleWithDialer request: command=Connect, ATYP=IPv4, 127.0.0.1:22, + // CRLF. HandleWithDialer needs at least 10 more bytes after the + // trojan header to dispatch into HandleTCP (which calls Dial). + scripts := make([]fakeRead, trojan.HeaderLen+2+10) + for i := 0; i < trojan.HeaderLen; i++ { + scripts[i] = fakeRead{data: []byte{'x'}, err: nil} + } + scripts[trojan.HeaderLen] = fakeRead{data: []byte{0x0d}, err: nil} + scripts[trojan.HeaderLen+1] = fakeRead{data: []byte{0x0a}, err: nil} + // HandleWithDialer bytes: CMD=1, ATYP=1, 127.0.0.1, port 22, CRLF. + post := []byte{0x01, 0x01, 127, 0, 0, 1, 0x00, 0x16, 0x0d, 0x0a} + for i, b := range post { + scripts[trojan.HeaderLen+2+i] = fakeRead{data: []byte{b}, err: nil} + } + + conn := &fakeConn{scripts: scripts} + fl.acceptCh <- conn + + go ln.loop() + + // Wait until HandleWithDialer has reached Dial. + select { + case <-dialEntered: + case <-time.After(2 * time.Second): + t.Fatal("listener did not reach HandleWithDialer.Dial within 2s") + } + + // The fix: the listener must have cleared the sniff deadline BEFORE + // entering HandleWithDialer. The fakeConn.deadline field is set by + // every SetReadDeadline call (initial sniff timeout = ~10s; clear = + // zero). Without the clear, this would be ~10s in the future. + deadline := conn.deadline + if !deadline.IsZero() { + t.Errorf("conn read deadline at HandleWithDialer entry = %v (in %v); want zero time.Time{} (sniff deadline leaked into tunnel)", deadline, time.Until(deadline)) + } + + // Let Dial return an error so HandleWithDialer exits and the goroutine + // unwinds. + close(releaseDial) +} + +// blockingDialProxy is a test-only Proxy that signals on Dial entry and +// blocks until the test closes its release channel. The listener takes +// the concrete *blockingDialProxy directly via NewListener, so no caddy +// module registration is needed. +type blockingDialProxy struct { + entered chan struct{} + release chan struct{} +} + +func (*blockingDialProxy) Close() error { return nil } +func (p *blockingDialProxy) Dial(_, _ string) (net.Conn, error) { + select { + case <-p.entered: + default: + close(p.entered) + } + <-p.release + return nil, errors.New("blockingDialProxy: exit") +} +func (*blockingDialProxy) ListenPacket(_, _ string) (net.PacketConn, error) { + return nil, errors.New("not supported") +} + +// TestListenerAcceptErrorBackoff verifies that when the underlying Accept +// fails repeatedly, the loop does not spin: it pauses ~100ms between +// failures. Without the backoff, a failing underlying listener (e.g. during +// shutdown) could burn CPU at full speed. +// +// Method: replace the underlying net.Listener with a counting fake that +// records each Accept call, repeatedly returns errors, and let the loop run +// for a fixed window. With backoff the count must be small (roughly +// windowMs / 100ms); without backoff it would be orders of magnitude higher. +func TestListenerAcceptErrorBackoff(t *testing.T) { + t.Parallel() + + fl := newFakeListener() + ln := newTestListener(fl) + + // Wrap Accept to count calls without changing semantics. + fl.acceptCh = make(chan net.Conn, 16) // ensure non-blocking send paths don't matter here + var acceptCount int64 + var mu sync.Mutex + doneProbe := make(chan struct{}) + wrapped := &countingListener{ + Listener: fl, + count: &acceptCount, + mu: &mu, + done: doneProbe, + } + ln.Listener = wrapped + + go ln.loop() + + // Periodically push errors so the loop always has something to consume. + pushErr := func() { + select { + case fl.errCh <- errors.New("transient accept failure"): + default: + } + } + + // Run for ~600ms. With backoff, expect ~6 Accept calls (600/100). + const windowMs = 600 + start := time.Now() + deadline := start.Add(windowMs * time.Millisecond) + for time.Now().Before(deadline) { + pushErr() + time.Sleep(5 * time.Millisecond) + } + + // Stop the loop and read the final count. + fl.Close() + ln.Close() + close(doneProbe) + // Give loop goroutine time to exit. + time.Sleep(150 * time.Millisecond) + + mu.Lock() + count := acceptCount + mu.Unlock() + + // With 100ms backoff over a 600ms window we expect ~6 calls. We allow a + // generous upper bound (≤20) to avoid CI flakiness while still failing + // without the backoff, where the count would be 100+. + if count < 2 { + t.Errorf("loop ran too few Accept iterations: %d (loop may have exited early)", count) + } + if count > 20 { + t.Errorf("loop spun too fast on Accept errors: %d iterations in %dms (want ≤~%d); backoff regressed", + count, windowMs, windowMs/100*2) + } + t.Logf("Accept was called %d times in %dms (with backoff: ~%d expected)", count, windowMs, windowMs/100) +} + +// countingListener wraps a net.Listener and counts Accept calls. +type countingListener struct { + net.Listener + count *int64 + mu *sync.Mutex + done chan struct{} +} + +func (l *countingListener) Accept() (net.Conn, error) { + l.mu.Lock() + *l.count++ + l.mu.Unlock() + return l.Listener.Accept() +} + +// TestListenerClosedHonoredOnConnsSend verifies that when the listener is +// closed while a connection is queued, the per-conn goroutine takes the +// closed path and closes the connection rather than blocking forever on an +// unbuffered conns channel (the listener's consumer has gone away). +// +// This is hard to deterministically exercise from a unit test because the +// conns channel is unbuffered and the loop's read side is held by the loop +// goroutine itself. We at least verify the more important property: when +// the listener is closed, the loop's Accept returns net.ErrClosed, the +// loop exits, and no goroutine leaks. +func TestListenerClosedHonoredOnConnsSend(t *testing.T) { + t.Parallel() + + fl := newFakeListener() + ln := newTestListener(fl) + + done := make(chan struct{}) + go func() { + ln.loop() + close(done) + }() + + // Close both the underlying fakeListener (so Accept returns + // net.ErrClosed and the loop enters the error branch) and the wrapper + // Listener (so l.closed is closed and the error branch returns instead + // of sleeping 100ms and looping again). + fl.Close() + ln.Close() + + select { + case <-done: + // good: loop exited + case <-time.After(2 * time.Second): + t.Fatal("loop did not exit within 2s of listener close") + } +} + +// keep trojan imported so we can reference trojan.HeaderLen symbolically if +// we ever need to (no-op import-only reference to avoid drift). +var _ = trojan.HeaderLen diff --git a/pkgs/rawconn/conn.go b/pkgs/rawconn/conn.go index e9c4762..819d8cb 100644 --- a/pkgs/rawconn/conn.go +++ b/pkgs/rawconn/conn.go @@ -6,31 +6,17 @@ import ( "errors" "io" "net" - "reflect" - "unsafe" ) +// RewindConn returns a connection that replays the already-read bytes +// (plaintext, for a *tls.Conn) before continuing to read from conn. +// +// The rewind is done with a wrapper rather than by reaching into the +// unexported internals of crypto/tls.Conn, whose layout is not stable across +// Go releases. Replaying plaintext above the TLS layer is data-preserving and +// does not depend on stdlib internals. func RewindConn(conn net.Conn, read []byte) net.Conn { - if tlsConn, ok := conn.(*tls.Conn); ok { - var ( - tlsInput, _ = reflect.TypeOf(tls.Conn{}).FieldByName("input") - input = (*bytes.Reader)(unsafe.Add(unsafe.Pointer(tlsConn), tlsInput.Offset)) - remaining = input.Len() - size = int(input.Size()) - buffered = len(read) - ) - if buffered <= size { - _, _ = input.Seek(0, 0) - } else { - buf := make([]byte, buffered+remaining) - copy(buf, read) - _, _ = input.Read(buf[buffered:]) - input.Reset(buf) - } - return tlsConn - } else { - return NewConn(conn, read) - } + return NewConn(conn, read) } type conn struct { @@ -58,6 +44,13 @@ func (c *conn) Read(b []byte) (int, error) { return n, err } +// CloseWrite implements net.CloseWriter on the wrapped conn when the +// underlying type supports half-close (e.g. *net.TCPConn). +// +// *tls.Conn is not a CloseWriter (the TLS record layer cannot signal a +// half-close over a single direction), so the wrapper reports "not supported" +// for TLS conns. No current rawconn consumer on the listener path uses TCP +// half-close, so this is a documentation note rather than a functional gap. func (c *conn) CloseWrite() error { if cc, ok := c.Conn.(*net.TCPConn); ok { return cc.CloseWrite() @@ -69,3 +62,13 @@ func (c *conn) CloseWrite() error { } return errors.New("not supported") } + +// ConnectionState forwards the TLS state when the wrapped connection is a +// *tls.Conn, which caddy's http2listener uses to recover the TLS state. For +// non-TLS connections an empty state is returned. +func (c *conn) ConnectionState() tls.ConnectionState { + if tc, ok := c.Conn.(*tls.Conn); ok { + return tc.ConnectionState() + } + return tls.ConnectionState{} +} diff --git a/pkgs/rawconn/conn_test.go b/pkgs/rawconn/conn_test.go new file mode 100644 index 0000000..7620696 --- /dev/null +++ b/pkgs/rawconn/conn_test.go @@ -0,0 +1,207 @@ +package rawconn + +import ( + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "io" + "math/big" + "net" + "strings" + "testing" + "time" +) + +// TestRewindConnTLS verifies that plaintext already read from a *tls.Conn can +// be replayed through RewindConn on the current Go release. +func TestRewindConnTLS(t *testing.T) { + cert := testCertificate(t) + + client, server := net.Pipe() + defer client.Close() + defer server.Close() + + clientTLS := tls.Client(client, &tls.Config{InsecureSkipVerify: true}) + serverTLS := tls.Server(server, &tls.Config{Certificates: []tls.Certificate{cert}}) + + errCh := make(chan error, 1) + go func() { errCh <- serverTLS.Handshake() }() + if err := clientTLS.Handshake(); err != nil { + t.Fatalf("client handshake: %v", err) + } + if err := <-errCh; err != nil { + t.Fatalf("server handshake: %v", err) + } + + msg := []byte("hello world") + go func() { + if _, err := clientTLS.Write(msg); err != nil { + t.Errorf("client write: %v", err) + } + }() + + // Read part of the plaintext, then rewind and expect the full message. + first := make([]byte, len("hello")) + if _, err := io.ReadFull(serverTLS, first); err != nil { + t.Fatalf("first read: %v", err) + } + + rc := RewindConn(serverTLS, first) + all := make([]byte, len(msg)) + if _, err := io.ReadFull(rc, all); err != nil { + t.Fatalf("read after rewind: %v", err) + } + if string(all) != string(msg) { + t.Fatalf("after rewind got %q, want %q", all, msg) + } +} + +// TestRewindConnTLSState verifies that the wrapper returned by RewindConn +// forwards the TLS connection state to caddy/net/http instead of hiding it. +func TestRewindConnTLSState(t *testing.T) { + cert := testCertificate(t) + + client, server := net.Pipe() + defer client.Close() + defer server.Close() + + clientTLS := tls.Client(client, &tls.Config{InsecureSkipVerify: true}) + serverTLS := tls.Server(server, &tls.Config{Certificates: []tls.Certificate{cert}}) + + errCh := make(chan error, 1) + go func() { errCh <- serverTLS.Handshake() }() + if err := clientTLS.Handshake(); err != nil { + t.Fatalf("client handshake: %v", err) + } + if err := <-errCh; err != nil { + t.Fatalf("server handshake: %v", err) + } + + rc := RewindConn(serverTLS, nil) + + // The wrapper must expose the exported connectionStater method so caddy's + // http2listener/ConnContext path can recover the TLS state. + stater, ok := rc.(interface{ ConnectionState() tls.ConnectionState }) + if !ok { + t.Fatal("RewindConn result does not implement ConnectionState") + } + if got := stater.ConnectionState(); got.Version == 0 { + t.Errorf("ConnectionState = %+v, want valid TLS version", got) + } +} + +// TestRewindConnPlainState verifies that rewinding a non-TLS connection stays +// safe: ConnectionState returns an empty, not-valid state instead of panicking. +func TestRewindConnPlainState(t *testing.T) { + c1, c2 := net.Pipe() + defer c1.Close() + defer c2.Close() + + rc := RewindConn(c1, []byte("x")) + stater, ok := rc.(interface{ ConnectionState() tls.ConnectionState }) + if !ok { + t.Fatal("RewindConn result does not implement ConnectionState") + } + if got := stater.ConnectionState(); got.Version != 0 { + t.Errorf("plain conn ConnectionState = %+v, want empty state", got) + } +} + +// TestConnCloseWriteTCP verifies that closing the write side of a +// wrapped *net.TCPConn forwards to the underlying conn: the peer +// sees a clean EOF rather than a reset or hang. Uses a loopback +// TCP listener/accept pair so CloseWrite is meaningful (net.Pipe +// does not support half-close). +func TestConnCloseWriteTCP(t *testing.T) { + t.Parallel() + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + type acceptResult struct { + conn net.Conn + err error + } + acceptDone := make(chan acceptResult, 1) + go func() { + c, err := ln.Accept() + acceptDone <- acceptResult{c, err} + }() + + client, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer client.Close() + + var server net.Conn + select { + case r := <-acceptDone: + if r.err != nil { + t.Fatal(r.err) + } + server = r.conn + case <-time.After(2 * time.Second): + t.Fatal("Accept did not return") + } + defer server.Close() + + wrapped := NewConn(server, nil) + if err := wrapped.(interface{ CloseWrite() error }).CloseWrite(); err != nil { + t.Fatalf("CloseWrite on TCPConn wrapper = %v, want nil", err) + } + + client.SetReadDeadline(time.Now().Add(2 * time.Second)) + buf := make([]byte, 1) + n, err := client.Read(buf) + if err != io.EOF { + t.Errorf("client Read after server CloseWrite = (%d, %v), want io.EOF", n, err) + } +} + +// TestConnCloseWriteUnsupportedType verifies that closing the +// write side of a wrapped conn that has no half-close support +// returns 'not supported' rather than silently doing the wrong +// thing. We use a net.Pipe (which has no CloseWrite method) as +// the wrapped conn to drive the unsupported branch. +func TestConnCloseWriteUnsupportedType(t *testing.T) { + t.Parallel() + + c1, c2 := net.Pipe() + defer c1.Close() + defer c2.Close() + + wrapped := NewConn(c1, nil) + err := wrapped.(interface{ CloseWrite() error }).CloseWrite() + if err == nil { + t.Fatal("CloseWrite on unsupported wrapper = nil, want error") + } + if !strings.Contains(err.Error(), "not supported") { + t.Errorf("CloseWrite error = %v, want 'not supported'", err) + } +} + +func testCertificate(t *testing.T) tls.Certificate { + t.Helper() + priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatalf("generate key: %v", err) + } + tmpl := x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "localhost"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + } + der, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, &priv.PublicKey, priv) + if err != nil { + t.Fatalf("create certificate: %v", err) + } + return tls.Certificate{Certificate: [][]byte{der}, PrivateKey: priv} +} diff --git a/pkgs/trojan/trojan.go b/pkgs/trojan/trojan.go index ce4770a..3bf3f9a 100644 --- a/pkgs/trojan/trojan.go +++ b/pkgs/trojan/trojan.go @@ -15,6 +15,15 @@ import ( const HeaderLen = 56 +// errInvalidCRLF is returned when the protocol CRLF terminator is missing or +// malformed. +var errInvalidCRLF = errors.New("invalid CRLF terminator") + +// allocByteErr, when non-nil, is returned by HandleTCP/HandleUDP instead of +// performing real memory.Alloc. Production leaves it nil; tests set it to +// simulate an mmap failure without exhausting memory. +var allocByteErr error + const ( CmdConnect = 1 CmdAssociate = 3 @@ -51,6 +60,9 @@ func HandleWithDialer(r io.Reader, w io.Writer, d Dialer) (int64, int64, error) if _, err := io.ReadFull(r, b[1:3]); err != nil { return 0, 0, fmt.Errorf("read 0x0d 0x0a error: %w", err) } + if b[1] != 0x0d || b[2] != 0x0a { + return 0, 0, errInvalidCRLF + } switch b[0] { case CmdConnect: diff --git a/pkgs/trojan/trojan_tcp.go b/pkgs/trojan/trojan_tcp.go index c58a40e..fcc9a24 100644 --- a/pkgs/trojan/trojan_tcp.go +++ b/pkgs/trojan/trojan_tcp.go @@ -2,6 +2,7 @@ package trojan import ( "errors" + "fmt" "io" "net" "os" @@ -10,6 +11,16 @@ import ( "github.com/imgk/memory-go" ) +// tcpBufSize is the per-direction copy buffer for HandleTCP. Allocated in the +// main flow before goroutines spawn so an mmap failure surfaces as a regular +// (0, 0, err) return rather than a goroutine panic. +const tcpBufSize = 32 * 1024 + +// wakeGrace bounds how long HandleTCP keeps the client-side read open after the +// remote half-closes its write side. The FIN is not a full close: in-flight +// data (e.g. SSH final packets) must drain before the peer is force-woken. +const wakeGrace = 2 * time.Second + func copyBuffer(w io.Writer, r io.Reader, buf []byte) (n int64, err error) { for { nr, er := r.Read(buf) @@ -41,28 +52,56 @@ func copyBuffer(w io.Writer, r io.Reader, buf []byte) (n int64, err error) { return n, err } +// allocShortCircuit returns the test-only allocByteErr, or nil in production +// so HandleTCP/HandleUDP proceed to allocate normally. Tests override +// allocByteErr to exercise the alloc-failure branch without exhausting memory. +func allocShortCircuit() error { + if err := allocByteErr; err != nil { + return fmt.Errorf("memory alloc error: %w", err) + } + return nil +} + func HandleTCP(r io.Reader, w io.Writer, addr net.Addr, d Dialer) (int64, int64, error) { + if err := allocShortCircuit(); err != nil { + return 0, 0, err + } rc, err := d.Dial("tcp", addr.String()) if err != nil { return 0, 0, err } defer rc.Close() + ptrA, bufA, err := memory.Alloc[byte](tcpBufSize) + if err != nil { + return 0, 0, fmt.Errorf("memory alloc error: %w", err) + } + defer memory.Free(ptrA) + ptrB, bufB, err := memory.Alloc[byte](tcpBufSize) + if err != nil { + return 0, 0, fmt.Errorf("memory alloc error: %w", err) + } + defer memory.Free(ptrB) + type Result struct { Num int64 Err error } errCh := make(chan Result) - go func(rc net.Conn, r io.Reader, errCh chan Result) { - ptr, buf, err := memory.Alloc[byte](32 * 1024) - if err != nil { - panic(err) - } - defer memory.Free(ptr) - + go func(rc net.Conn, r io.Reader, buf []byte, errCh chan Result) { nr, err := copyBuffer(io.Writer(rc), r, buf) - if err == nil || errors.Is(err, os.ErrDeadlineExceeded) { + // Clean exits cover EOF and the wakeGrace timeout. Wrappers like + // gorilla/websocket return a *netError with Timeout()==true rather + // than the stdlib sentinel, so check the net.Error interface too. + isClean := err == nil + if !isClean { + var nerr net.Error + if errors.As(err, &nerr) && nerr.Timeout() { + isClean = true + } + } + if isClean { if cw, ok := rc.(interface { CloseWrite() error }); ok { @@ -79,15 +118,13 @@ func HandleTCP(r io.Reader, w io.Writer, addr net.Addr, d Dialer) (int64, int64, } rc.SetReadDeadline(time.Now()) errCh <- Result{Num: nr, Err: err} - }(rc, r, errCh) - - nr, nw, err := func(rc net.Conn, w io.Writer, errCh chan Result) (int64, int64, error) { - ptr, buf, err := memory.Alloc[byte](32 * 1024) - if err != nil { - panic(err) - } - defer memory.Free(ptr) + }(rc, r, bufA, errCh) + nr, nw, err := func(rc net.Conn, w io.Writer, buf []byte, errCh chan Result) (int64, int64, error) { + // The write path through w never sees a deadline error: rc is a raw + // *net.TCPConn and only the reader goroutine's deadline is set (via + // trySetReadDeadline). Timeout detection therefore stays on the + // stdlib sentinel here. nw, err := copyBuffer(w, io.Reader(rc), buf) if err == nil { if cw, ok := w.(interface { @@ -95,6 +132,9 @@ func HandleTCP(r io.Reader, w io.Writer, addr net.Addr, d Dialer) (int64, int64, }); ok { cw.CloseWrite() } + // rc finished writing but the client may still be sending; give + // it wakeGrace to drain instead of truncating in-flight data. + trySetReadDeadline(w, time.Now().Add(wakeGrace)) r := <-errCh return r.Num, nw, r.Err } @@ -126,18 +166,25 @@ func HandleTCP(r io.Reader, w io.Writer, addr net.Addr, d Dialer) (int64, int64, }); ok { cw.CloseWrite() } + wakeReader(w) r := <-errCh return r.Num, nw, r.Err } - rc.SetWriteDeadline(time.Now()) + rc.SetWriteDeadline(time.Now().Add(wakeGrace)) if cw, ok := rc.(interface { CloseWrite() error }); ok { cw.CloseWrite() } + wakeReader(w) r := <-errCh + // If the writer timed out but the reader exited cleanly, the tunnel is + // in a normal shutdown shape: report nil. + if r.Err == nil && isTimeoutError(err) { + return r.Num, nw, nil + } return r.Num, nw, err - }(rc, w, errCh) + }(rc, w, bufB, errCh) return nr, nw, err } diff --git a/pkgs/trojan/trojan_test.go b/pkgs/trojan/trojan_test.go new file mode 100644 index 0000000..f5f37bc --- /dev/null +++ b/pkgs/trojan/trojan_test.go @@ -0,0 +1,500 @@ +package trojan + +import ( + "bytes" + "errors" + "io" + "net" + "strings" + "testing" + "time" +) + +type mockDialer struct{} + +func (mockDialer) Dial(string, string) (net.Conn, error) { + return nil, nil // must not be reached by the invalid-CRLF tests +} + +func (mockDialer) ListenPacket(string, string) (net.PacketConn, error) { + return net.ListenPacket("udp", "127.0.0.1:0") +} + +// pipeDialer returns a fixed connection from Dial, simulating the upstream +// endpoint of a CONNECT tunnel. +type pipeDialer struct { + conn net.Conn +} + +func (d *pipeDialer) Dial(string, string) (net.Conn, error) { + return d.conn, nil +} + +func (d *pipeDialer) ListenPacket(string, string) (net.PacketConn, error) { + return nil, errors.New("not supported") +} + +// TestHandleWithDialerRejectsInvalidCRLF verifies the TCP header CRLF +// terminator is checked before any dial happens. +func TestHandleWithDialerRejectsInvalidCRLF(t *testing.T) { + // [CMD=1][ATYP=1 IPv4][127.0.0.1][port 8080][CRLF=0x00 0x00] + stream := []byte{0x01, 0x01, 127, 0, 0, 1, 0x1f, 0x90, 0x00, 0x00} + _, _, err := HandleWithDialer(bytes.NewReader(stream), io.Discard, mockDialer{}) + if err == nil { + t.Fatal("HandleWithDialer accepted invalid CRLF terminator") + } + if !strings.Contains(err.Error(), "CRLF") { + t.Errorf("HandleWithDialer error = %v, want CRLF-related error", err) + } +} + +// TestHandleTCPNoDeadlockOnHalfClose is a regression test for the F1 +// deadlock: when the upstream closes first (half-close) and the client stays +// silent, HandleWithDialer must return instead of blocking on errCh forever. +func TestHandleTCPNoDeadlockOnHalfClose(t *testing.T) { + // Real loopback TCP pairs: net.Pipe does not support CloseWrite, which is + // needed to simulate the upstream half-close. + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + client, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer client.Close() + // srv is the connection handed to HandleWithDialer; the client peer stays + // silent after sending the header. + srv, err := ln.Accept() + if err != nil { + t.Fatal(err) + } + defer srv.Close() + + tln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tln.Close() + target, err := net.Dial("tcp", tln.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer target.Close() + tpeer, err := tln.Accept() + if err != nil { + t.Fatal(err) + } + defer tpeer.Close() + + d := &pipeDialer{conn: target} + done := make(chan error, 1) + go func() { + _, _, err := HandleWithDialer(srv, srv, d) + done <- err + }() + + // CONNECT to 127.0.0.1:80 with a valid CRLF terminator. + header := []byte{0x01, 0x01, 127, 0, 0, 1, 0x1f, 0x90, 0x0d, 0x0a} + if _, err := client.Write(header); err != nil { + t.Fatalf("write header: %v", err) + } + + // Let HandleTCP start, then half-close the upstream side: tpeer.CloseWrite + // makes target.Read return EOF while the client stays silent. + time.Sleep(100 * time.Millisecond) + if err := tpeer.(*net.TCPConn).CloseWrite(); err != nil { + t.Fatalf("tpeer CloseWrite: %v", err) + } + + select { + case err := <-done: + if err != nil { + t.Logf("HandleWithDialer returned error: %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("HandleWithDialer deadlocked after upstream half-close with silent client") + } +} + +// TestHandleUDPNoDeadlockOnIdle is a regression test for the F2 deadlock: an +// idle client that never sends a UDP packet must not leave HandleUDP blocked +// on errCh after the socket deadline expires. +func TestHandleUDPNoDeadlockOnIdle(t *testing.T) { + client, _ := net.Pipe() + defer client.Close() + + done := make(chan error, 1) + go func() { + _, _, err := HandleUDP(client, client, 500*time.Millisecond, mockDialer{}) + done <- err + }() + + // The client sends nothing; the write side times out after 500ms and must + // wake the blocked reader instead of waiting on errCh forever. + select { + case err := <-done: + // F2 follow-up: the writer loop's deadline timeout, surfaced + // as the writer goroutine's os.ErrDeadlineExceeded after the + // reader defer's rc.SetReadDeadline-now call, must be coalesced + // to nil by HandleUDP's main-goroutine cleanup branch. Otherwise + // every normal UDP session close logs a spurious 'i/o timeout'. + if err != nil { + t.Errorf("HandleUDP returned %v, want nil on idle timeout", err) + } + case <-time.After(5 * time.Second): + t.Fatal("HandleUDP deadlocked on idle client") + } +} + +// TestHandleUDPCleanExitOnClientClose verifies that a client closing +// the read side of the conn produces a nil error from HandleUDP: the +// reader's io.EOF is coalesced in the reader defer, the writer loop +// then sees a deadline timeout from its blocked ReadFrom, and the +// main goroutine's cleanup branch folds that timeout back to nil. +// Pre-F2-followup, the writer-side deadline would propagate up as +// "i/o timeout" and every normal UDP teardown would log a spurious +// error. +func TestHandleUDPCleanExitOnClientClose(t *testing.T) { + client, _ := net.Pipe() + defer client.Close() + + done := make(chan error, 1) + go func() { + _, _, err := HandleUDP(client, client, 5*time.Second, mockDialer{}) + done <- err + }() + + // Give HandleUDP time to start; the writer loop sets a 5s read + // deadline on the UDP socket. Then close the client side and + // expect HandleUDP to exit within the deadline. + time.Sleep(100 * time.Millisecond) + client.Close() + + select { + case err := <-done: + if err != nil { + t.Errorf("HandleUDP returned %v, want nil on client close", err) + } + case <-time.After(6 * time.Second): + t.Fatal("HandleUDP did not return within 6s after client close") + } +} + +// TestHandleUDPRejectsInvalidCRLF verifies the UDP packet CRLF terminator is +// checked and reported through the error channel. +func TestHandleUDPRejectsInvalidCRLF(t *testing.T) { + // [ATYP=1 IPv4][127.0.0.1][port 8080][Len=0x0004][CRLF=0x00 0x00][data] + stream := []byte{ + 0x01, 127, 0, 0, 1, 0x1f, 0x90, + 0x00, 0x04, 0x00, 0x00, + 'd', 'a', 't', 'a', + } + _, _, err := HandleUDP(bytes.NewReader(stream), io.Discard, time.Second, mockDialer{}) + if err == nil { + t.Fatal("HandleUDP accepted invalid CRLF terminator") + } + if !strings.Contains(err.Error(), "CRLF") { + t.Errorf("HandleUDP error = %v, want CRLF-related error", err) + } + if !errors.Is(err, errInvalidCRLF) { + t.Errorf("HandleUDP error = %v, want errInvalidCRLF sentinel", err) + } +} + +// failingAddr is a net.Addr that returns a fixed string; used only so +// HandleTCP's Dial path has a syntactically valid target. +type failingAddr struct{} + +func (failingAddr) Network() string { return "tcp" } +func (failingAddr) String() string { return "127.0.0.1:0" } + +// dialReturningPipe is a Dialer whose Dial returns the half of a net.Pipe +// already set up by the test, so HandleTCP's Dial succeeds and we exercise +// the post-dial memory.Alloc path. +type dialReturningPipe struct{ conn net.Conn } + +func (d *dialReturningPipe) Dial(string, string) (net.Conn, error) { return d.conn, nil } +func (d *dialReturningPipe) ListenPacket(string, string) (net.PacketConn, error) { + return nil, errors.New("not supported") +} + +// withAllocErr temporarily sets the package-level allocByteErr to the given +// error and registers a cleanup that restores it. It is the test hook for +// the B3 (mmap failure propagation) fix: production leaves allocByteErr +// nil, so HandleTCP/HandleUDP call the real memory.Alloc. +func withAllocErr(t *testing.T, err error) { + t.Helper() + orig := allocByteErr + t.Cleanup(func() { allocByteErr = orig }) + allocByteErr = err +} + +// TestHandleTCPAllocFailureReturnsError verifies that HandleTCP returns an +// error instead of panicking when memory.Alloc fails (e.g. mmap() under +// memory pressure in the `malloc_syscall` build). Pre-fix the goroutines +// contain `panic(err)`; a real mmap failure would crash the whole process. +// Post-fix the failure surfaces as a (0, 0, err) return without a panic. +// +// This test runs without t.Parallel() because it mutates package-level +// state (allocByteErr); parallel runs would race on that variable. +func TestHandleTCPAllocFailureReturnsError(t *testing.T) { + client, srv := net.Pipe() + defer client.Close() + defer srv.Close() + + withAllocErr(t, errors.New("simulated alloc failure")) + + d := &dialReturningPipe{conn: srv} + done := make(chan struct{}) + var nr, nw int64 + var err error + go func() { + defer close(done) + defer func() { + if r := recover(); r != nil { + err = errors.New("HandleTCP panicked instead of returning an error") + } + }() + nr, nw, err = HandleTCP(client, client, failingAddr{}, d) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("HandleTCP did not return within 2s on alloc failure; possible deadlock or panic") + } + + if err == nil { + t.Fatalf("HandleTCP returned (nr=%d, nw=%d, err=nil) on alloc failure; want non-nil error", nr, nw) + } + if nr != 0 || nw != 0 { + t.Errorf("HandleTCP on alloc failure returned nr=%d nw=%d; want 0/0", nr, nw) + } +} + +// TestHandleUDPAllocFailureReturnsError is the UDP counterpart of the TCP +// alloc-failure test: HandleUDP must return (0, 0, err) without panicking +// if memory.Alloc fails. mockDialer in this package implements ListenPacket +// by binding a real UDP socket; we reuse it so Alloc is reached. +func TestHandleUDPAllocFailureReturnsError(t *testing.T) { + client, _ := net.Pipe() + defer client.Close() + + withAllocErr(t, errors.New("simulated alloc failure")) + + d := mockDialer{} + done := make(chan struct{}) + var nr, nw int64 + var err error + go func() { + defer close(done) + defer func() { + if r := recover(); r != nil { + err = errors.New("HandleUDP panicked instead of returning an error") + } + }() + nr, nw, err = HandleUDP(client, client, 200*time.Millisecond, d) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("HandleUDP did not return within 2s on alloc failure; possible deadlock or panic") + } + + if err == nil { + t.Fatalf("HandleUDP returned (nr=%d, nw=%d, err=nil) on alloc failure; want non-nil error", nr, nw) + } + if nr != 0 || nw != 0 { + t.Errorf("HandleUDP on alloc failure returned nr=%d nw=%d; want 0/0", nr, nw) + } +} + +// TestHandleTCPNormalPathUnaffected sanity-checks that under the default +// (real) allocator, HandleTCP can still complete a small round trip — +// guarding against the B3 fix accidentally breaking the happy path. +// +// This test is intentionally not Parallel() because it shares package state +// with the alloc-failure test, but a t.Cleanup guarantees allocByteErr is +// restored to nil before any other test runs. +func TestHandleTCPNormalPathUnaffected(t *testing.T) { + withAllocErr(t, nil) // explicit: real allocator + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + upstream, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer upstream.Close() + go func() { + c, _ := ln.Accept() + _ = c.Close() + }() + + d := &dialReturningPipe{conn: upstream} + done := make(chan error, 1) + go func() { + _, _, err := HandleTCP(upstream, upstream, failingAddr{}, d) + done <- err + }() + select { + case err := <-done: + if err != nil { + t.Logf("HandleTCP returned (expected EOF or similar): %v", err) + } + case <-time.After(3 * time.Second): + t.Fatal("HandleTCP deadlocked under default allocator") + } +} + +// TestHandleTCPHalfCloseDrainsClientData is a regression test for the +// half-close handling in HandleTCP: when the remote side closes its write +// direction (TCP half-close, e.g. an SSH server sending FIN), data the +// client still has in flight must be delivered, not truncated. The fix +// replaced an immediate wakeReader deadline on this path with a grace +// window (wakeGrace); pre-fix, the immediate read deadline killed the +// client→remote direction and dropped the final packets, causing spurious +// disconnects on long-lived tunnels. +func TestHandleTCPHalfCloseDrainsClientData(t *testing.T) { + // Tunnel client-side pair: tcTest (test holds) <-> tcTunnel (passed to + // HandleTCP as r and w). + lnT, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer lnT.Close() + tcTunnel, err := net.Dial("tcp", lnT.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer tcTunnel.Close() + tcTest, err := lnT.Accept() + if err != nil { + t.Fatal(err) + } + defer tcTest.Close() + + // Remote (rc) pair, standing in for the SSH server: rcTunnel is what + // HandleTCP dials, rcTest is the "server" the test controls. + lnR, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer lnR.Close() + rcTunnel, err := net.Dial("tcp", lnR.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer rcTunnel.Close() + rcTest, err := lnR.Accept() + if err != nil { + t.Fatal(err) + } + defer rcTest.Close() + + d := &dialReturningPipe{conn: rcTunnel} + done := make(chan struct{}) + go func() { + defer close(done) + _, _, _ = HandleTCP(tcTunnel, tcTunnel, failingAddr{}, d) + }() + + // SSH server half-closes its write direction. + if err := rcTest.(*net.TCPConn).CloseWrite(); err != nil { + t.Fatalf("rc CloseWrite: %v", err) + } + + // Give HandleTCP a moment to observe the remote EOF and enter the + // grace-window path. + time.Sleep(100 * time.Millisecond) + + // The client still has a final packet in flight; it must arrive. + if _, err := tcTest.Write([]byte("final-packet")); err != nil { + t.Fatalf("client Write after half-close: %v", err) + } + + got := make([]byte, len("final-packet")) + if _, err := io.ReadFull(rcTest, got); err != nil { + t.Fatalf("rc did not receive post-FIN data (truncated by immediate wake): %v", err) + } + if string(got) != "final-packet" { + t.Errorf("rc received %q, want %q", got, "final-packet") + } + + // Unwind: close both sides so HandleTCP returns. + tcTest.Close() + rcTest.Close() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("HandleTCP did not return after both sides closed") + } +} + +// assertion is reached. +type addrPacketConn struct{} + +func (addrPacketConn) ReadFrom([]byte) (int, net.Addr, error) { + return 0, &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 80}, nil +} +func (addrPacketConn) WriteTo([]byte, net.Addr) (int, error) { return 0, nil } +func (addrPacketConn) Close() error { return nil } +func (addrPacketConn) LocalAddr() net.Addr { return nil } +func (addrPacketConn) SetDeadline(time.Time) error { return nil } +func (addrPacketConn) SetReadDeadline(time.Time) error { return nil } +func (addrPacketConn) SetWriteDeadline(time.Time) error { return nil } + +type nonUDPDialer struct{} + +func (nonUDPDialer) Dial(string, string) (net.Conn, error) { return nil, nil } +func (nonUDPDialer) ListenPacket(string, string) (net.PacketConn, error) { + return addrPacketConn{}, nil +} + +// TestHandleUDPNonUDPAddrHardError covers the comma-ok hardening plus the +// reader-release fix: when the writer loop's addr.(*net.UDPAddr) assertion +// fails, HandleUDP must surface an 'unsupported packet address type' error and +// release the reader goroutine so the main goroutine does not hang on errCh. +// +// The client side of the pipe is deliberately not closed, so the reader blocks +// in serverSide.Read; only trySetImmediateReadDeadline on the hard-error break +// releases it. Pre-fix the writer panicked (no comma-ok) or the main goroutine +// hung (no reader release). +func TestHandleUDPNonUDPAddrHardError(t *testing.T) { + clientSide, serverSide := net.Pipe() + defer clientSide.Close() + defer serverSide.Close() + + done := make(chan struct{}) + var gotErr error + go func() { + defer close(done) + _, _, gotErr = HandleUDP(serverSide, io.Discard, time.Minute*10, nonUDPDialer{}) + }() + + select { + case <-done: + if gotErr == nil { + t.Fatal("HandleUDP returned nil on non-UDPAddr; want hard error") + } + if !strings.Contains(gotErr.Error(), "unsupported packet address type") { + t.Errorf("HandleUDP error = %v, want 'unsupported packet address type'", gotErr) + } + case <-time.After(5 * time.Second): + t.Fatal("HandleUDP did not return within 5s on non-UDPAddr hard error; " + + "the writer-loop hard-error break path did not release the reader. " + + "This is the e44c4a3 reader-release bug: " + + "trySetImmediateReadDeadline(r) is missing from the hard-error break, " + + "so the reader stays blocked on r.Read and the main goroutine hangs " + + "on the errCh wait.") + } +} + +// failingAddr is a net.Addr that returns a fixed string; used only so +// HandleTCP's Dial path has a syntactically valid target. diff --git a/pkgs/trojan/trojan_udp.go b/pkgs/trojan/trojan_udp.go index 69d4de2..6cceb79 100644 --- a/pkgs/trojan/trojan_udp.go +++ b/pkgs/trojan/trojan_udp.go @@ -3,9 +3,9 @@ package trojan import ( "bytes" "errors" + "fmt" "io" "net" - "os" "time" "github.com/imgk/caddy-trojan/pkgs/socks" @@ -15,12 +15,35 @@ import ( // [AddrType(1 byte)][Addr(max 256 byte)][Port(2 byte)][Len(2 byte)][0x0d, 0x0a][Data(max 65535 byte)] func HandleUDP(r io.Reader, w io.Writer, timeout time.Duration, d Dialer) (int64, int64, error) { + if err := allocShortCircuit(); err != nil { + return 0, 0, err + } rc, err := d.ListenPacket("udp", "") if err != nil { return 0, 0, err } defer rc.Close() + ptrPrev, bb, err := memory.Alloc[byte](socks.MaxAddrLen) + if err != nil { + return 0, 0, fmt.Errorf("memory alloc error: %w", err) + } + defer memory.Free(ptrPrev) + + ptrRead, b, err := memory.Alloc[byte](64*1024 + socks.MaxAddrLen) + if err != nil { + return 0, 0, fmt.Errorf("memory alloc error: %w", err) + } + defer memory.Free(ptrRead) + + ptrWrite, bWrite, err := memory.Alloc[byte](64*1024 + socks.MaxAddrLen + 4) + if err != nil { + return 0, 0, fmt.Errorf("memory alloc error: %w", err) + } + defer memory.Free(ptrWrite) + bWrite[socks.MaxAddrLen+2] = 0x0d + bWrite[socks.MaxAddrLen+3] = 0x0a + type Result struct { Num int64 Err error @@ -29,27 +52,16 @@ func HandleUDP(r io.Reader, w io.Writer, timeout time.Duration, d Dialer) (int64 errCh := make(chan Result) go func(rc net.PacketConn, r io.Reader, errCh chan Result) (nr int64, err error) { defer func() { - if errors.Is(err, io.EOF) || errors.Is(err, os.ErrDeadlineExceeded) { + // Coalesce clean exits (EOF, closed pipe, or a self-imposed + // deadline) into nil for the main goroutine. + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrClosedPipe) || isTimeoutError(err) { err = nil } errCh <- Result{Num: nr, Err: err} }() - // save previous address - ptr, bb, err := memory.Alloc[byte](socks.MaxAddrLen) - if err != nil { - panic(err) - } - defer memory.Free(ptr) - tt := (*net.UDPAddr)(nil) - ptr, b, err := memory.Alloc[byte](64*1024 + socks.MaxAddrLen) - if err != nil { - panic(err) - } - defer memory.Free(ptr) - for { raddr, er := socks.ReadAddrBuffer(r, b) if er != nil { @@ -73,6 +85,10 @@ func HandleUDP(r io.Reader, w io.Writer, timeout time.Duration, d Dialer) (int64 err = er break } + if b[l+2] != 0x0d || b[l+3] != 0x0a { + err = errInvalidCRLF + break + } l += (int(b[l])<<8 | int(b[l+1])) nr += int64(l) + 4 @@ -92,26 +108,27 @@ func HandleUDP(r io.Reader, w io.Writer, timeout time.Duration, d Dialer) (int64 return }(rc, r, errCh) - nr, nw, err := func(rc net.PacketConn, w io.Writer, errCh chan Result, timeout time.Duration) (_, nw int64, err error) { - ptr, b, err := memory.Alloc[byte](64*1024 + socks.MaxAddrLen + 4) - if err != nil { - panic(err) - } - defer memory.Free(ptr) - - b[socks.MaxAddrLen+2] = 0x0d - b[socks.MaxAddrLen+3] = 0x0a + nr, nw, err := func(rc net.PacketConn, w io.Writer, clientR io.Reader, buf []byte, errCh chan Result, timeout time.Duration) (_, nw int64, err error) { for { rc.SetReadDeadline(time.Now().Add(timeout)) - n, addr, er := rc.ReadFrom(b[socks.MaxAddrLen+4:]) + n, addr, er := rc.ReadFrom(buf[socks.MaxAddrLen+4:]) if er != nil { err = er break } - b[socks.MaxAddrLen] = byte(n >> 8) - b[socks.MaxAddrLen+1] = byte(n) + buf[socks.MaxAddrLen] = byte(n >> 8) + buf[socks.MaxAddrLen+1] = byte(n) + udpAddr, ok := addr.(*net.UDPAddr) + if !ok { + // ReadFrom may return any net.Addr; fail loud instead of + // panicking, and release the reader's blocked Read so the + // main goroutine does not hang on errCh. + err = fmt.Errorf("handle udp error: unsupported packet address type %T", addr) + trySetImmediateReadDeadline(clientR) + break + } l := func(bb []byte, addr *net.UDPAddr) int64 { if ipv4 := addr.IP.To4(); ipv4 != nil { const offset = socks.MaxAddrLen - (1 + net.IPv4len + 2) @@ -126,23 +143,28 @@ func HandleUDP(r io.Reader, w io.Writer, timeout time.Duration, d Dialer) (int64 bb[offset+1+net.IPv6len], bb[offset+1+net.IPv6len+1] = byte(addr.Port>>8), byte(addr.Port) return 1 + net.IPv6len + 2 } - }(b[:socks.MaxAddrLen], addr.(*net.UDPAddr)) - nw += 4 + int64(n) + l + }(buf[:socks.MaxAddrLen], udpAddr) - if _, ew := w.Write(b[socks.MaxAddrLen-l : socks.MaxAddrLen+4+n]); ew != nil { + if _, ew := w.Write(buf[socks.MaxAddrLen-l : socks.MaxAddrLen+4+n]); ew != nil { err = ew break } } rc.SetWriteDeadline(time.Now()) - if errors.Is(err, io.EOF) || errors.Is(err, os.ErrDeadlineExceeded) { - r := <-errCh + wakeReader(w) + r := <-errCh + // Prefer the reader's error (it owns CRLF validation) over the + // writer's, and suppress the writer's clean timeout so normal + // session teardown does not log a spurious i/o timeout. + if r.Err != nil { return r.Num, nw, r.Err } - r := <-errCh + if isTimeoutError(err) { + return r.Num, nw, nil + } return r.Num, nw, err - }(rc, w, errCh, timeout) + }(rc, w, r, bWrite, errCh, timeout) return nr, nw, err } diff --git a/pkgs/trojan/wake_reader_test.go b/pkgs/trojan/wake_reader_test.go new file mode 100644 index 0000000..cd065e1 --- /dev/null +++ b/pkgs/trojan/wake_reader_test.go @@ -0,0 +1,145 @@ +package trojan + +import ( + "io" + "net" + "sync" + "sync/atomic" + "testing" + "time" +) + +// wakeNetConn is a minimal net.Conn whose SetReadDeadline records the +// deadline and signals a channel so a blocked Read can return. It only +// implements the methods wakeReader touches. +type wakeNetConn struct { + deadline atomic.Pointer[time.Time] + armed atomic.Bool + read atomic.Int32 // counts Read calls + mu sync.Mutex + done chan struct{} +} + +func newWakeNetConn() *wakeNetConn { + return &wakeNetConn{done: make(chan struct{})} +} + +func (c *wakeNetConn) SetReadDeadline(t time.Time) error { + c.deadline.Store(&t) + if !t.IsZero() { + c.armed.Store(true) + select { + case <-c.done: + default: + close(c.done) + } + } + return nil +} + +func (c *wakeNetConn) Read(b []byte) (int, error) { + c.read.Add(1) + if !c.armed.Load() { + <-c.done + } + return 0, osErrDeadlineExceeded +} + +func (c *wakeNetConn) Write(b []byte) (int, error) { return len(b), nil } +func (c *wakeNetConn) Close() error { return nil } +func (c *wakeNetConn) LocalAddr() net.Addr { return nil } +func (c *wakeNetConn) RemoteAddr() net.Addr { return nil } +func (c *wakeNetConn) SetDeadline(t time.Time) error { return c.SetReadDeadline(t) } +func (c *wakeNetConn) SetWriteDeadline(_ time.Time) error { return nil } + +func (c *wakeNetConn) snapshot() (time.Time, bool) { + p := c.deadline.Load() + if p == nil { + return time.Time{}, false + } + return *p, true +} + +// osErrDeadlineExceeded matches the timeout error net.Conn returns when a +// Read fires after the read deadline has passed. We re-export os's +// sentinel here for clarity. +var osErrDeadlineExceeded = timeoutError{} + +type timeoutError struct{} + +func (timeoutError) Error() string { return "i/o timeout" } +func (timeoutError) Timeout() bool { return true } +func (timeoutError) Temporary() bool { return true } + +// TestWakeReaderSetsImmediateDeadlineOnNetConn verifies that wakeReader +// applies an immediate read deadline (≈ now) to a net.Conn. The fix is in +// HandleTCP/HandleUDP's half-close recovery: when one direction finishes, +// wakeReader releases any goroutine blocked reading on the same conn so +// the other goroutine can exit instead of deadlocking on errCh forever. +func TestWakeReaderSetsImmediateDeadlineOnNetConn(t *testing.T) { + t.Parallel() + + fc := newWakeNetConn() + before := time.Now() + wakeReader(fc) + after := time.Now() + + d, ok := fc.snapshot() + if !ok { + t.Fatal("SetReadDeadline was not called") + } + if d.Before(before) || d.After(after) { + t.Errorf("deadline = %v, want in [%v, %v] (immediate, not future)", d, before, after) + } +} + +// TestWakeReaderIsNoopOnNonConn verifies that wakeReader does not panic +// when given an io.Writer that is not a net.Conn. This protects the +// HTTP/2 / HTTP/3 path in HandleWithDialer where w is an http.Body / +// FlushWriter rather than a net.Conn. +func TestWakeReaderIsNoopOnNonConn(t *testing.T) { + t.Parallel() + + wakeReader(io.Discard) // must not panic +} + +// TestWakeReaderReleasesBlockedReader verifies the end-to-end behavior the +// fix is supposed to deliver: a goroutine blocked on Read is released when +// wakeReader sets a deadline. The fake conn's Read blocks until +// SetReadDeadline is called. +func TestWakeReaderReleasesBlockedReader(t *testing.T) { + t.Parallel() + + fc := newWakeNetConn() + + done := make(chan error, 1) + go func() { + buf := make([]byte, 1) + _, err := fc.Read(buf) + done <- err + }() + + // Let the goroutine start blocking on Read. + time.Sleep(10 * time.Millisecond) + + wakeReader(fc) + + select { + case err := <-done: + if err == nil { + t.Fatal("Read returned nil error after wakeReader; want timeout error") + } + if !isTimeoutErr(err) { + t.Errorf("Read after wakeReader = %v, want a timeout error", err) + } + case <-time.After(1 * time.Second): + t.Fatal("Read did not return within 1s; wakeReader did not release the blocked reader") + } +} + +// isTimeoutErr returns true if err implements net.Error and Timeout() is true. +func isTimeoutErr(err error) bool { + type timeoutErr interface{ Timeout() bool } + te, ok := err.(timeoutErr) + return ok && te.Timeout() +} diff --git a/pkgs/trojan/wakeable.go b/pkgs/trojan/wakeable.go new file mode 100644 index 0000000..41b46a6 --- /dev/null +++ b/pkgs/trojan/wakeable.go @@ -0,0 +1,77 @@ +package trojan + +import ( + "errors" + "io" + "net" + "os" + "time" + + pkgwebsocket "github.com/imgk/caddy-trojan/pkgs/websocket" +) + +// wakeableConn lets HandleTCP/HandleUDP release a reader goroutine blocked in +// Read via a deadline change, regardless of how many wrappers sit between the +// writer and the underlying net.Conn. The method must be exported so types in +// other packages (notably pkgs/websocket.Conn) can satisfy the interface. +type wakeableConn interface { + SetReadDeadline(time.Time) error +} + +// trySetReadDeadline sets w's read deadline if w implements wakeableConn and +// reports whether the call was made. false means the writer is opaque and the +// caller cannot bound a silent peer's blocking read. +func trySetReadDeadline(w io.Writer, t time.Time) bool { + if c, ok := w.(wakeableConn); ok { + _ = c.SetReadDeadline(t) + return true + } + return false +} + +// trySetImmediateReadDeadline unblocks any pending Read on the connection that +// rw wraps. rw is typed any because the HandleUDP hard-error path passes an +// io.Reader here. +func trySetImmediateReadDeadline(rw any) bool { + c, ok := rw.(wakeableConn) + if !ok { + return false + } + _ = c.SetReadDeadline(time.Now()) + return true +} + +// wakeReader unblocks a pending Read on w's connection. It is the safety net +// used before waiting on errCh in HandleTCP/HandleUDP so a silent peer cannot +// keep the reader goroutine blocked forever. It is a no-op for writers that do +// not implement wakeableConn. +func wakeReader(w io.Writer) { + _ = trySetImmediateReadDeadline(w) +} + +// isTimeoutError reports whether err is a clean timeout: either the stdlib +// sentinel or a net.Error with Timeout()==true (e.g. gorilla/websocket's +// hideTempErr-wrapped *netError). Such errors mean the conn hit a deadline we +// set ourselves and should not propagate as a tunnel-level error. +func isTimeoutError(err error) bool { + if err == nil { + return false + } + if errors.Is(err, os.ErrDeadlineExceeded) { + return true + } + var nerr net.Error + if errors.As(err, &nerr) && nerr.Timeout() { + return true + } + return false +} + +// Compile-time guards that the types the wake protocol depends on still satisfy +// wakeableConn, so a future interface drift becomes a build failure instead of +// a silent runtime no-op. +var ( + _ wakeableConn = (*net.TCPConn)(nil) + _ wakeableConn = (*net.UDPConn)(nil) + _ wakeableConn = (*pkgwebsocket.Conn)(nil) +) diff --git a/pkgs/trojan/wakeable_test.go b/pkgs/trojan/wakeable_test.go new file mode 100644 index 0000000..cf929bf --- /dev/null +++ b/pkgs/trojan/wakeable_test.go @@ -0,0 +1,292 @@ +package trojan + +import ( + "net" + "net/http" + "net/http/httptest" + "net/url" + "testing" + "time" + + "github.com/gorilla/websocket" + + pkgwebsocket "github.com/imgk/caddy-trojan/pkgs/websocket" +) + +// TestWakeableConnDispatchesToWebSocketWrapper is the cross-package +// integration test for wakeableConn on the HandleTCP path. Using a real +// loopback TCP pair (net.Pipe has no half-close), it half-closes the dial +// side so the grace path sets a read deadline that must wake the reader +// within wakeGrace. Asserts both the timing (wakeGrace + 1s margin) and that +// HandleWithDialer returns nil (not 'i/o timeout'). +func TestWakeableConnDispatchesToWebSocketWrapper(t *testing.T) { + t.Parallel() + + tln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tln.Close() + + upstreamTarget, err := net.Dial("tcp", tln.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer upstreamTarget.Close() + upstreamPeer, err := tln.Accept() + if err != nil { + t.Fatal(err) + } + defer upstreamPeer.Close() + + upgrader := websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + } + + gotErr := make(chan error, 1) + + mux := http.NewServeMux() + mux.HandleFunc("/tunnel", func(w http.ResponseWriter, r *http.Request) { + c, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade: %v", err) + return + } + wrapper := pkgwebsocket.NewConn(c) + defer wrapper.Close() + + d := &pipeDialer{conn: upstreamPeer} + _, _, err = HandleWithDialer(wrapper, wrapper, d) + gotErr <- err + }) + + srv := httptest.NewServer(mux) + defer srv.Close() + + u, err := url.Parse(srv.URL) + if err != nil { + t.Fatal(err) + } + u.Scheme = "ws" + u.Path = "/tunnel" + + dialer := *websocket.DefaultDialer + dialer.HandshakeTimeout = 5 * time.Second + clientConn, _, err := dialer.Dial(u.String(), nil) + if err != nil { + t.Fatalf("ws dial: %v", err) + } + defer clientConn.Close() + + header := []byte{0x01, 0x01, 127, 0, 0, 1, 0x1f, 0x90, 0x0d, 0x0a} + if err := clientConn.WriteMessage(websocket.BinaryMessage, header); err != nil { + t.Fatalf("write header: %v", err) + } + + time.Sleep(200 * time.Millisecond) + if err := upstreamTarget.(*net.TCPConn).CloseWrite(); err != nil { + t.Fatalf("upstreamTarget CloseWrite: %v", err) + } + + start := time.Now() + select { + case e := <-gotErr: + elapsed := time.Since(start) + if elapsed > 3*time.Second { + t.Errorf("HandleWithDialer returned in %v, "+ + "expected < 3s (wakeGrace + scheduling margin). "+ + "A slow return suggests the clean-shutdown chain is degraded.", elapsed) + } + if e != nil { + t.Errorf("HandleWithDialer returned %v, want nil for a clean WS teardown. "+ + "A non-nil return means the reader-defer or main-goroutine cleanup did "+ + "not fold the writer-side deadline to nil (F2-followup regression).", e) + } + case <-time.After(5 * time.Second): + t.Fatal("HandleWithDialer did not return within 5s after upstream half-close. " + + "This indicates either (a) the wakeableConn method is unexported so " + + "the cross-package WebSocket hook cannot satisfy the interface, or " + + "(b) the reader goroutine does not recognize gorilla's hideTempErr-" + + "wrapped timeout as a clean exit.") + } +} + +// TestWakeableConnClientToUpstreamFlow is the inverse-direction +// integration test: upstream pushes data and then half-closes; +// HandleTCP's main goroutine reads from rc and writes to the WS +// wrapper. Asserts both the timing (within wakeGrace + margin) +// AND that HandleWithDialer returns nil. +func TestWakeableConnClientToUpstreamFlow(t *testing.T) { + t.Parallel() + + tln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer tln.Close() + upstreamTarget, err := net.Dial("tcp", tln.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer upstreamTarget.Close() + upstreamPeer, err := tln.Accept() + if err != nil { + t.Fatal(err) + } + defer upstreamPeer.Close() + + if _, err := upstreamTarget.Write([]byte("hello, world\n")); err != nil { + t.Fatal(err) + } + if err := upstreamTarget.(*net.TCPConn).CloseWrite(); err != nil { + t.Fatal(err) + } + + upgrader := websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + } + + gotErr := make(chan error, 1) + + mux := http.NewServeMux() + mux.HandleFunc("/tunnel", func(w http.ResponseWriter, r *http.Request) { + c, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade: %v", err) + return + } + wrapper := pkgwebsocket.NewConn(c) + defer wrapper.Close() + + d := &pipeDialer{conn: upstreamPeer} + _, _, err = HandleWithDialer(wrapper, wrapper, d) + gotErr <- err + }) + + srv := httptest.NewServer(mux) + defer srv.Close() + + u, err := url.Parse(srv.URL) + if err != nil { + t.Fatal(err) + } + u.Scheme = "ws" + u.Path = "/tunnel" + + clientConn, _, err := (*websocket.DefaultDialer).Dial(u.String(), nil) + if err != nil { + t.Fatalf("ws dial: %v", err) + } + defer clientConn.Close() + + header := []byte{0x01, 0x01, 127, 0, 0, 1, 0x1f, 0x90, 0x0d, 0x0a} + if err := clientConn.WriteMessage(websocket.BinaryMessage, header); err != nil { + t.Fatalf("write header: %v", err) + } + + start := time.Now() + select { + case e := <-gotErr: + elapsed := time.Since(start) + if elapsed > 3*time.Second { + t.Errorf("HandleWithDialer returned in %v, "+ + "expected < 3s (wakeGrace + scheduling margin).", elapsed) + } + if e != nil { + t.Errorf("HandleWithDialer returned %v, want nil for a clean WS teardown. "+ + "A non-nil return means the cleanup did not fold the writer-side "+ + "deadline to nil.", e) + } + case <-time.After(5 * time.Second): + t.Fatal("HandleWithDialer did not return within 5s after upstream write+half-close.") + } +} + +// TestHandleUDPOverWebSocketWrapperCleanExit exercises the +// HandleUDP path with a real gorilla WebSocket wrapper for both +// r and w. The client writes the CmdAssociate header and +// immediately closes; the server-side wrapper's NextReader +// returns io.EOF (the client stream EOFs), the reader goroutine's +// defer folds the EOF to nil, the writer loop's blocked ReadFrom +// then sees a deadline timeout via rc.SetReadDeadline(now), and +// the main goroutine's isTimeoutError check coalesces both back +// to nil. Asserts BOTH the timing (within WS close round-trip + +// margin) AND that HandleWithDialer returns nil (F2-followup +// invariant). +func TestHandleUDPOverWebSocketWrapperCleanExit(t *testing.T) { + t.Parallel() + + upgrader := websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + } + + gotErr := make(chan error, 1) + + mux := http.NewServeMux() + mux.HandleFunc("/tunnel", func(w http.ResponseWriter, r *http.Request) { + c, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade: %v", err) + return + } + wrapper := pkgwebsocket.NewConn(c) + defer wrapper.Close() + + _, _, e := HandleWithDialer(wrapper, wrapper, mockDialerUDP{}) + gotErr <- e + }) + + srv := httptest.NewServer(mux) + defer srv.Close() + + u, err := url.Parse(srv.URL) + if err != nil { + t.Fatal(err) + } + u.Scheme = "ws" + u.Path = "/tunnel" + + clientConn, _, err := (*websocket.DefaultDialer).Dial(u.String(), nil) + if err != nil { + t.Fatalf("ws dial: %v", err) + } + + // CmdAssociate (0x03) header for the UDP ASSOCIATE command. + header := []byte{0x03, 0x01, 127, 0, 0, 1, 0x1f, 0x90, 0x0d, 0x0a} + if err := clientConn.WriteMessage(websocket.BinaryMessage, header); err != nil { + t.Fatalf("write header: %v", err) + } + clientConn.Close() + + start := time.Now() + select { + case e := <-gotErr: + elapsed := time.Since(start) + if elapsed > 1*time.Second { + t.Errorf("HandleWithDialer took %v, expected < 1s for a "+ + "clean WS-over-UDP teardown. A slow return suggests "+ + "the clean-shutdown chain is degraded.", elapsed) + } + if e != nil { + t.Errorf("HandleWithDialer returned %v, want nil for a "+ + "clean WS-over-UDP teardown. A non-nil return means the "+ + "reader-defer or main-goroutine cleanup did not fold "+ + "the writer-side deadline to nil (F2-followup regression).", e) + } + case <-time.After(5 * time.Second): + t.Fatal("HandleWithDialer did not return within 5s for WS-over-UDP.") + } +} + +// mockDialerUDP is a Dialer that returns a real UDP socket from +// ListenPacket so HandleUDP's writer loop has a valid PacketConn +// to wait on. Dial is unused. +type mockDialerUDP struct{} + +func (mockDialerUDP) Dial(string, string) (net.Conn, error) { + return nil, nil +} + +func (mockDialerUDP) ListenPacket(string, string) (net.PacketConn, error) { + return net.ListenPacket("udp", "127.0.0.1:0") +} diff --git a/pkgs/websocket/websocket.go b/pkgs/websocket/websocket.go index a23b768..fec9f41 100644 --- a/pkgs/websocket/websocket.go +++ b/pkgs/websocket/websocket.go @@ -69,3 +69,8 @@ func (c *Conn) Close() error { c.Conn.WriteControl(websocket.CloseMessage, msg, time.Now().Add(time.Second*5)) return c.Conn.Close() } + +// SetReadDeadline is promoted from the embedded *gorilla.Conn, satisfying +// pkgs/trojan.wakeableConn automatically. We intentionally add no wrapper +// method: gorilla's *websocket.Conn already forwards SetReadDeadline to the +// underlying net.Conn, and method-set promotion gives it the same signature. diff --git a/pkgs/x/utils.go b/pkgs/x/utils.go index bc48345..4175c01 100644 --- a/pkgs/x/utils.go +++ b/pkgs/x/utils.go @@ -11,7 +11,14 @@ func ByteSliceToString(b []byte) string { } func StringToByteSlice(s string) []byte { - return unsafe.Slice((*byte)(unsafe.Pointer(*(*uintptr)(unsafe.Pointer(&s)))), len(s)) + // unsafe.StringData returns a *byte pointing at s's underlying bytes + // without going through a uintptr round-trip. The old form + // unsafe.Slice((*byte)(unsafe.Pointer(*(*uintptr)(unsafe.Pointer(&s)))), len(s)) + // violated the unsafe.Pointer rules (pointer<->uintptr conversions must + // stay within a single expression) and made Go 1.26's stricter vet + // flag the line as 'possible misuse of unsafe.Pointer'. unsafe.StringData + // is the standard library's blessed escape hatch for exactly this case. + return unsafe.Slice(unsafe.StringData(s), len(s)) } // RemoveNullKeys removes keys with null values from a JSON object represented as a map[string]any. diff --git a/tools/gauntlet.sh b/tools/gauntlet.sh new file mode 100755 index 0000000..206511f --- /dev/null +++ b/tools/gauntlet.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# tools/gauntlet.sh — single entry point that re-runs every gauntlet layer +# so the EVIDENCE report is reproducible from the repo alone. +# +# Usage: tools/gauntlet.sh +# +# Layers (in order): +# 1. Clean stale artifacts (freshness by mechanism, not discipline). +# 2. go build ./... (static types — must produce no errors). +# 3. go vet ./... (static analysis — pre-existing pkgs/x +# unsafe.Pointer warning is expected; we exit +# zero only on new warnings). +# 4. go test ./... (full unit suite). +# 5. go test -coverprofile (changed-line coverage on app/, modules/handler/, +# pkgs/trojan/). +# 6. tools/mutants.sh (manual mutation gauntlet). +set -euo pipefail + +cd "$(dirname "$0")/.." + +echo "=== gauntlet: cleaning stale artifacts ===" +rm -f c.out c.html + +echo "=== gauntlet layer 2/3: go build + go vet ===" +go build ./... +# go vet emits one pre-existing warning about pkgs/x/utils.go's +# ByteSliceToString unsafe.Pointer usage (see AGENTS.md); that file is +# not part of this change. We surface vet but do not fail on pre-existing +# warnings. +go vet ./... 2>&1 | grep -v "pkgs/x/utils.go" || true + +echo "=== gauntlet layer 4: full test suite ===" +go test ./... + +echo "=== gauntlet layer 5: changed-line coverage ===" +go test -coverprofile=c.out -covermode=atomic \ + ./app/... ./modules/handler/... ./pkgs/trojan/... +echo "--- changed/added function coverage ---" +go tool cover -func=c.out \ + | grep -E "validStorageKey|recordFailure|isRateLimited|allocShortCircuit|tcpBufSize" \ + || true + +echo "=== gauntlet layer 6: manual mutation ===" +tools/mutants.sh + +echo "=== gauntlet: done ===" diff --git a/tools/mutants.sh b/tools/mutants.sh new file mode 100755 index 0000000..4192b1b --- /dev/null +++ b/tools/mutants.sh @@ -0,0 +1,129 @@ +#!/usr/bin/env bash +# tools/mutants.sh — manual mutation gauntlet for the B1/B2/B3 fixes. +# +# Background: Go has no mature default mutation-testing tool, so we +# introduce plausible bugs into the new/changed code one at a time, run +# the suite, and confirm each mutant is KILLED by at least one test. +# +# If a mutant SURVIVES (no test fails), the test suite has a hole and a +# new assertion must be added before the fix can be declared done. +# +# Re-run-safe: each mutant is applied via python3 literal substitution, +# the test suite is run, then `git checkout` restores the file. +# IMPORTANT: the implementation must already be COMMITTED for restore +# to be a no-op for our changes; otherwise `git checkout` reverts +# working-tree changes and breaks the next steps. Run after the +# [GREEN] commit, not during RED. + +set -euo pipefail + +cd "$(dirname "$0")/.." + +# Sanity: working tree must be either clean or have only untracked files +# relative to HEAD, so that `git checkout -- file` is a no-op for files +# we didn't intend to revert. +git status --porcelain | awk '{print $1}' | grep -v '^??' | grep -v 'M$' && { + echo "ERROR: working tree has staged or modified files; commit first." >&2 + exit 2 +} || true + +TRACKED=( + app/upstream.go + modules/handler/handler.go + pkgs/trojan/trojan_tcp.go + pkgs/trojan/trojan_udp.go +) +cleanup() { git checkout -- "${TRACKED[@]}" 2>/dev/null || true; } +trap cleanup EXIT + +apply_mutant() { + local file="$1" find="$2" replace="$3" + python3 - "$file" "$find" "$replace" <<'PY' +import sys, io +p, find, rep = sys.argv[1], sys.argv[2], sys.argv[3] +with io.open(p, 'r', encoding='utf-8') as f: + src = f.read() +n = src.count(find) +if n == 0: + print(f" SKIP: pattern not present in {p}") + sys.exit(0) +if n > 1: + print(f" ERROR: pattern matches {n} times in {p}; aborting") + sys.exit(2) +with io.open(p, 'w', encoding='utf-8') as f: + f.write(src.replace(find, rep, 1)) +PY +} + +run_focused() { go test -run "$2" "$1" >/dev/null 2>&1; } + +declare -i killed=0 survived=0 +mutant_status() { + if [ "$1" -ne 0 ]; then + echo " KILLED"; killed=$((killed+1)) + else + echo " SURVIVED"; survived=$((survived+1)) + fi +} + +# M1: weaken validStorageKey's hex check +echo "--- M1: validStorageKey accepts any byte ---" +apply_mutant app/upstream.go \ + "if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) {" \ + "if false { // MUTANT M1" +run_focused ./app/ 'TestCaddyUpstream' && mutant_status 0 || mutant_status $? +git checkout -- app/upstream.go + +# M2: weaken the length check +echo "--- M2: length check removed ---" +apply_mutant app/upstream.go \ + "if len(k) != trojan.HeaderLen { + return false + }" \ + "// MUTANT M2: length check removed + if false { return false; } + if len(k) != trojan.HeaderLen {" +run_focused ./app/ 'TestCaddyUpstream' && mutant_status 0 || mutant_status $? +git checkout -- app/upstream.go + +# M3: never block (rate-limit disabled) +echo "--- M3: rate-limit block-check disabled ---" +apply_mutant modules/handler/handler.go \ + "if time.Now().Before(rec.blockedUntil) { + return true + }" \ + "// MUTANT M3: always fall through + if false { return true; } + if time.Now().Before(rec.blockedUntil) {" +run_focused ./modules/handler/ 'TestHandlerConnect' && mutant_status 0 || mutant_status $? +git checkout -- modules/handler/handler.go + +# M4: threshold check defeated +echo "--- M4: threshold check unreachable ---" +apply_mutant modules/handler/handler.go \ + "if rec.count >= threshold && rec.blockedUntil.IsZero() { + rec.blockedUntil = now.Add(cooldown) + }" \ + "// MUTANT M4: never block + if false { rec.blockedUntil = now.Add(cooldown); } + if rec.count >= threshold && rec.blockedUntil.IsZero() {" +run_focused ./modules/handler/ 'TestHandlerConnect' && mutant_status 0 || mutant_status $? +git checkout -- modules/handler/handler.go + +# M5: HandleTCP: disable the test-injected allocShortCircuit (so the test +# must reach the real-alloc path), then neutralize the post-alloc error +# return. With both halves applied the test MUST observe an error; without +# the hook being honored the test would see nil and fail. +echo "--- M5: HandleTCP dispatches around allocByteErr short-circuit ---" +apply_mutant pkgs/trojan/trojan_tcp.go \ + "if err := allocByteErr; err != nil { + return 0, 0, fmt.Errorf(\"memory alloc error: %w\", err) + }" \ + "// MUTANT M5: short-circuit disabled + if false { return 0, 0, fmt.Errorf(\"memory alloc error: %w\", allocByteErr) }" +run_focused ./pkgs/trojan/ 'TestHandleTCPAllocFailure' && mutant_status 0 || mutant_status $? +git checkout -- pkgs/trojan/trojan_tcp.go + +echo +echo "=== mutation summary: $killed killed, $survived survived ===" +[ "$survived" -eq 0 ] || exit 1