refactor(handler/sni): extract helpers, add SnifferBuilder, comprehensive tests
- Extract newRecorderObject, checkRateLimit, sniffingDial, and handleSniffedProtocol from the Handle method into dedicated files (sniffing.go, util.go) to reduce Handle from ~100 to ~65 lines. - Introduce SnifferBuilder: populated once during Init, reused via Build() per connection. Avoids reconstructing a 9-field sniffing.Sniffer literal on every inbound connection. - Hoist router nil check from the lazy dial closure into Handle so misconfigured handlers fail immediately instead of silently dropping unrecognised traffic with no error. - Add nil-addr guard in checkRateLimit (missing from the old inline version — would panic if the remote address were somehow nil). - Add comprehensive unit tests (44 tests across 5 test files) covering handler creation, Init, metadata parsing (defaults, custom, MITM, error paths), protocol dispatch (HTTP, TLS, unknown, empty), rate limiting, recorder objects, dial plumbing, and SnifferBuilder. - Add package-level godoc explaining the handler scope, limitations vs tcp/forward, and the connection processing flow. - Fix metadata comment: s/SnifferBuilder/sniffing.Sniffer/.
This commit is contained in:
@@ -0,0 +1,148 @@
|
||||
package sni
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-gost/core/handler"
|
||||
"github.com/go-gost/core/limiter/rate"
|
||||
xctx "github.com/go-gost/x/ctx"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// newRecorderObject
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestNewRecorderObject_TCP(t *testing.T) {
|
||||
h := newInitdHandler()
|
||||
conn := newStringConn(nil)
|
||||
start := time.Now()
|
||||
|
||||
ro := h.newRecorderObject(context.Background(), conn, start)
|
||||
|
||||
if ro.Network != "tcp" {
|
||||
t.Errorf("expected network tcp, got %s", ro.Network)
|
||||
}
|
||||
if ro.RemoteAddr != "10.0.0.1:12345" {
|
||||
t.Errorf("expected remote 10.0.0.1:12345, got %s", ro.RemoteAddr)
|
||||
}
|
||||
if ro.LocalAddr != "127.0.0.1:8080" {
|
||||
t.Errorf("expected local 127.0.0.1:8080, got %s", ro.LocalAddr)
|
||||
}
|
||||
if !ro.Time.Equal(start) {
|
||||
t.Errorf("expected time %v, got %v", start, ro.Time)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewRecorderObject_SID(t *testing.T) {
|
||||
h := newInitdHandler()
|
||||
conn := newStringConn(nil)
|
||||
ctx := xctx.ContextWithSid(context.Background(), xctx.Sid("test-sid"))
|
||||
|
||||
ro := h.newRecorderObject(ctx, conn, time.Now())
|
||||
|
||||
if ro.SID != "test-sid" {
|
||||
t.Errorf("expected SID 'test-sid', got '%s'", ro.SID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewRecorderObject_NilSrcAddr(t *testing.T) {
|
||||
h := newInitdHandler()
|
||||
conn := newStringConn(nil)
|
||||
|
||||
ro := h.newRecorderObject(context.Background(), conn, time.Now())
|
||||
|
||||
if ro.ClientAddr != "" {
|
||||
t.Errorf("expected empty ClientAddr, got %s", ro.ClientAddr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewRecorderObject_SrcAddr(t *testing.T) {
|
||||
h := newInitdHandler()
|
||||
conn := newStringConn(nil)
|
||||
srcAddr := &net.TCPAddr{IP: net.IPv4(192, 168, 1, 1), Port: 54321}
|
||||
ctx := xctx.ContextWithSrcAddr(context.Background(), srcAddr)
|
||||
|
||||
ro := h.newRecorderObject(ctx, conn, time.Now())
|
||||
|
||||
if ro.ClientAddr != "192.168.1.1:54321" {
|
||||
t.Errorf("expected ClientAddr '192.168.1.1:54321', got '%s'", ro.ClientAddr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewRecorderObject_Service(t *testing.T) {
|
||||
h := newTestHandler(func(o *handler.Options) {
|
||||
o.Service = "my-service"
|
||||
})
|
||||
conn := newStringConn(nil)
|
||||
|
||||
ro := h.newRecorderObject(context.Background(), conn, time.Now())
|
||||
|
||||
if ro.Service != "my-service" {
|
||||
t.Errorf("expected Service 'my-service', got '%s'", ro.Service)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// checkRateLimit
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestCheckRateLimit_NilLimiter(t *testing.T) {
|
||||
h := newInitdHandler()
|
||||
conn := newStringConn(nil)
|
||||
|
||||
if !h.checkRateLimit(conn.RemoteAddr()) {
|
||||
t.Error("expected true when RateLimiter is nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckRateLimit_NilAddr(t *testing.T) {
|
||||
h := newInitdHandler(withRateLimiter(&stubRateLimiter{
|
||||
limiterFn: func(key string) rate.Limiter {
|
||||
return &stubLimiter{allowFn: func(n int) bool { return false }}
|
||||
},
|
||||
}))
|
||||
|
||||
if !h.checkRateLimit(nil) {
|
||||
t.Error("expected true when addr is nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckRateLimit_Allowed(t *testing.T) {
|
||||
h := newInitdHandler(withRateLimiter(&stubRateLimiter{
|
||||
limiterFn: func(key string) rate.Limiter {
|
||||
return &stubLimiter{allowFn: func(n int) bool { return true }}
|
||||
},
|
||||
}))
|
||||
conn := newStringConn(nil)
|
||||
|
||||
if !h.checkRateLimit(conn.RemoteAddr()) {
|
||||
t.Error("expected true when limiter allows")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckRateLimit_Blocked(t *testing.T) {
|
||||
h := newInitdHandler(withRateLimiter(&stubRateLimiter{
|
||||
limiterFn: func(key string) rate.Limiter {
|
||||
return &stubLimiter{allowFn: func(n int) bool { return false }}
|
||||
},
|
||||
}))
|
||||
conn := newStringConn(nil)
|
||||
|
||||
if h.checkRateLimit(conn.RemoteAddr()) {
|
||||
t.Error("expected false when limiter blocks")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckRateLimit_NoLimiterForKey(t *testing.T) {
|
||||
h := newInitdHandler(withRateLimiter(&stubRateLimiter{
|
||||
limiterFn: func(key string) rate.Limiter { return nil },
|
||||
}))
|
||||
conn := newStringConn(nil)
|
||||
|
||||
if !h.checkRateLimit(conn.RemoteAddr()) {
|
||||
t.Error("expected true when no limiter found for key")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user