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:
ginuerzh
@@ -0,0 +1,244 @@
|
||||
package sni
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"testing"
|
||||
|
||||
"github.com/go-gost/core/chain"
|
||||
"github.com/go-gost/core/handler"
|
||||
"github.com/go-gost/core/limiter/rate"
|
||||
"github.com/go-gost/core/recorder"
|
||||
xmd "github.com/go-gost/x/metadata"
|
||||
xrecorder "github.com/go-gost/x/recorder"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// NewHandler
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestNewHandler(t *testing.T) {
|
||||
h := NewHandler(func(o *handler.Options) {
|
||||
o.Service = "test-svc"
|
||||
})
|
||||
if h == nil {
|
||||
t.Fatal("expected non-nil handler")
|
||||
}
|
||||
sh, ok := h.(*sniHandler)
|
||||
if !ok {
|
||||
t.Fatal("expected *sniHandler")
|
||||
}
|
||||
if sh.options.Service != "test-svc" {
|
||||
t.Errorf("expected Service 'test-svc', got %s", sh.options.Service)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Init
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestInit_Defaults(t *testing.T) {
|
||||
h := newTestHandler()
|
||||
err := h.Init(xmd.NewMetadata(nil))
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if h.md.readTimeout <= 0 {
|
||||
t.Error("expected positive default readTimeout")
|
||||
}
|
||||
if h.sniffer == nil {
|
||||
t.Error("expected sniffer to be initialised")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInit_SelectsRecorder(t *testing.T) {
|
||||
h := newTestHandler(withRecorder(recorder.RecorderObject{
|
||||
Record: xrecorder.RecorderServiceHandler,
|
||||
Recorder: &mockRecorder{},
|
||||
}))
|
||||
err := h.Init(xmd.NewMetadata(nil))
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if h.recorder.Recorder == nil {
|
||||
t.Error("expected recorder to be set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInit_CertPool(t *testing.T) {
|
||||
cert, key := generateTestCertKey(t)
|
||||
h := newTestHandler()
|
||||
h.md.certificate = cert
|
||||
h.md.privateKey = key
|
||||
err := h.Init(xmd.NewMetadata(nil))
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if h.certPool == nil {
|
||||
t.Error("expected certPool to be created")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInit_ParseMetadataError(t *testing.T) {
|
||||
h := newTestHandler()
|
||||
md := xmd.NewMetadata(map[string]any{
|
||||
"mitm.certFile": "/nonexistent/cert.pem",
|
||||
"mitm.keyFile": "/nonexistent/key.pem",
|
||||
})
|
||||
err := h.Init(md)
|
||||
if err == nil {
|
||||
t.Fatal("expected error from Init with bad cert files")
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Handle — error paths
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestHandle_NoRouter(t *testing.T) {
|
||||
h := newTestHandler(withRouter(nil))
|
||||
conn := newStringConn(nil)
|
||||
|
||||
err := h.Handle(context.Background(), conn)
|
||||
if err == nil {
|
||||
t.Fatal("expected error when router is nil")
|
||||
}
|
||||
if !errors.Is(err, errRouterNotAvailable) {
|
||||
t.Errorf("expected errRouterNotAvailable, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandle_RateLimited(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)
|
||||
|
||||
err := h.Handle(context.Background(), conn)
|
||||
if err == nil {
|
||||
t.Fatal("expected rate limit error")
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Handle — sniffing (protocol detection)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestHandle_SniffError(t *testing.T) {
|
||||
// An empty connection causes sniffing.Sniff to return an error because
|
||||
// there are no bytes to peek. The handler should log the error and fall
|
||||
// through to the unrecognised path (silent drop).
|
||||
h := newInitdHandler()
|
||||
conn := newStringConn(nil)
|
||||
|
||||
err := h.Handle(context.Background(), conn)
|
||||
if err != nil {
|
||||
t.Fatalf("expected nil error for unrecognised traffic, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandle_UnrecognizedProto(t *testing.T) {
|
||||
h := newInitdHandler()
|
||||
// Non-HTTP/non-TLS bytes won't match any sniffer protocol.
|
||||
conn := newStringConn([]byte("ssh-2.0-OpenSSH"))
|
||||
|
||||
err := h.Handle(context.Background(), conn)
|
||||
if err != nil {
|
||||
t.Fatalf("expected nil error for unrecognised traffic, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandle_HTTP(t *testing.T) {
|
||||
server, client := net.Pipe()
|
||||
defer server.Close()
|
||||
defer client.Close()
|
||||
|
||||
go func() {
|
||||
buf := make([]byte, 4096)
|
||||
n, err := client.Read(buf)
|
||||
if err == nil && n > 0 {
|
||||
client.Write([]byte("HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n"))
|
||||
}
|
||||
}()
|
||||
|
||||
h := newInitdHandler(withRouter(&mockRouter{
|
||||
opts: &chain.RouterOptions{},
|
||||
dialFn: func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return server, nil
|
||||
},
|
||||
}))
|
||||
conn := newStringConn([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n"))
|
||||
|
||||
err := h.Handle(context.Background(), conn)
|
||||
if err != nil {
|
||||
t.Logf("sniffed HTTP returned: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandle_TLS(t *testing.T) {
|
||||
server, client := net.Pipe()
|
||||
defer server.Close()
|
||||
defer client.Close()
|
||||
|
||||
h := newInitdHandler(withRouter(&mockRouter{
|
||||
opts: &chain.RouterOptions{},
|
||||
dialFn: func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return server, nil
|
||||
},
|
||||
}))
|
||||
conn := newStringConn([]byte{
|
||||
0x16, 0x03, 0x01, 0x00, 0x06, 0x01, 0x00, 0x00, 0x02, 0x03, 0x01,
|
||||
})
|
||||
|
||||
err := h.Handle(context.Background(), conn)
|
||||
if err == nil {
|
||||
t.Log("sniffed TLS ClientHello without error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandle_RecorderError(t *testing.T) {
|
||||
// The deferred recorder error should not affect the return value.
|
||||
h := newTestHandler(withRouter(&mockRouter{
|
||||
opts: &chain.RouterOptions{},
|
||||
dialFn: func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return nil, nil
|
||||
},
|
||||
}))
|
||||
h.recorder = recorder.RecorderObject{Recorder: &errorRecorder{}}
|
||||
conn := newStringConn(nil)
|
||||
|
||||
err := h.Handle(context.Background(), conn)
|
||||
if err != nil {
|
||||
t.Fatalf("expected nil from unrecognised traffic, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandle_WithReadTimeout(t *testing.T) {
|
||||
h := newTestHandler(withRouter(&mockRouter{
|
||||
opts: &chain.RouterOptions{},
|
||||
}))
|
||||
h.md.readTimeout = 0 // zero disables the deadline
|
||||
h.recorder = recorder.RecorderObject{Recorder: &mockRecorder{}}
|
||||
conn := newStringConn([]byte("ssh"))
|
||||
|
||||
err := h.Handle(context.Background(), conn)
|
||||
if err != nil {
|
||||
t.Fatalf("expected nil error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandle_ConnClosed(t *testing.T) {
|
||||
h := newInitdHandler()
|
||||
conn := newStringConn(nil)
|
||||
|
||||
err := h.Handle(context.Background(), conn)
|
||||
if err != nil {
|
||||
t.Fatalf("expected nil error, got %v", err)
|
||||
}
|
||||
if !conn.closed {
|
||||
t.Error("expected connection to be closed after Handle returns")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user