diff --git a/admission/admission_test.go b/admission/admission_test.go new file mode 100644 index 00000000..778eba5c --- /dev/null +++ b/admission/admission_test.go @@ -0,0 +1,727 @@ +package admission + +import ( + "context" + "io" + "strings" + "testing" + "time" + + "github.com/go-gost/core/admission" + "github.com/go-gost/core/logger" + "github.com/go-gost/x/internal/loader" + xlogger "github.com/go-gost/x/logger" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// --- Option tests --- + +func TestWhitelistOption(t *testing.T) { + var opts options + WhitelistOption(true)(&opts) + assert.True(t, opts.whitelist) + + WhitelistOption(false)(&opts) + assert.False(t, opts.whitelist) +} + +func TestMatchersOption(t *testing.T) { + var opts options + MatchersOption([]string{"192.168.1.1", "10.0.0.0/8"})(&opts) + assert.Equal(t, []string{"192.168.1.1", "10.0.0.0/8"}, opts.matchers) +} + +func TestReloadPeriodOption(t *testing.T) { + var opts options + ReloadPeriodOption(5 * time.Second)(&opts) + assert.Equal(t, 5*time.Second, opts.period) +} + +func TestFileLoaderOption(t *testing.T) { + var opts options + fl := &mockLoader{} + FileLoaderOption(fl)(&opts) + assert.Equal(t, fl, opts.fileLoader) +} + +func TestRedisLoaderOption(t *testing.T) { + var opts options + rl := &mockLoader{} + RedisLoaderOption(rl)(&opts) + assert.Equal(t, rl, opts.redisLoader) +} + +func TestHTTPLoaderOption(t *testing.T) { + var opts options + hl := &mockLoader{} + HTTPLoaderOption(hl)(&opts) + assert.Equal(t, hl, opts.httpLoader) +} + +func TestLoggerOption(t *testing.T) { + var opts options + l := xlogger.Nop() + LoggerOption(l)(&opts) + assert.Equal(t, l, opts.logger) +} + +// newSyncedAdmission creates an admission controller and blocks until the initial +// background reload completes, so matchers are ready for immediate assertions. +func newSyncedAdmission(opts ...Option) *localAdmission { + a := NewAdmission(opts...) + la := a.(*localAdmission) + // Force a synchronous reload so tests don't race with the background goroutine. + la.reload(context.Background()) + return la +} + +// --- NewAdmission tests --- + +func TestNewAdmission_Defaults(t *testing.T) { + adm := newSyncedAdmission() + require.NotNil(t, adm) + defer adm.Close() + + // Should admit by default (blacklist mode with no rules) + assert.True(t, adm.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestNewAdmission_WithLogger(t *testing.T) { + adm := newSyncedAdmission(LoggerOption(xlogger.Nop())) + require.NotNil(t, adm) + defer adm.Close() +} + +// --- Admit tests --- + +func TestAdmit_EmptyAddr(t *testing.T) { + adm := newSyncedAdmission(MatchersOption([]string{"192.168.1.1"})) + defer adm.Close() + + assert.True(t, adm.Admit(context.Background(), "tcp", "")) +} + +func TestAdmit_NilReceiver(t *testing.T) { + var p *localAdmission + assert.True(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestAdmit_Blacklist_Matched(t *testing.T) { + adm := newSyncedAdmission(MatchersOption([]string{"192.168.1.1"})) + defer adm.Close() + + // In blacklist mode, matching IP is denied + assert.False(t, adm.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestAdmit_Blacklist_NotMatched(t *testing.T) { + adm := newSyncedAdmission(MatchersOption([]string{"192.168.1.1"})) + defer adm.Close() + + // In blacklist mode, non-matching IP is admitted + assert.True(t, adm.Admit(context.Background(), "tcp", "10.0.0.1")) +} + +func TestAdmit_Whitelist_Matched(t *testing.T) { + adm := newSyncedAdmission( + WhitelistOption(true), + MatchersOption([]string{"192.168.1.1"}), + ) + defer adm.Close() + + // In whitelist mode, matching IP is admitted + assert.True(t, adm.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestAdmit_Whitelist_NotMatched(t *testing.T) { + adm := newSyncedAdmission( + WhitelistOption(true), + MatchersOption([]string{"192.168.1.1"}), + ) + defer adm.Close() + + // In whitelist mode, non-matching IP is denied + assert.False(t, adm.Admit(context.Background(), "tcp", "10.0.0.1")) +} + +func TestAdmit_StripsPort(t *testing.T) { + adm := newSyncedAdmission(MatchersOption([]string{"192.168.1.1"})) + defer adm.Close() + + // Address with port should have port stripped before matching + assert.False(t, adm.Admit(context.Background(), "tcp", "192.168.1.1:8080")) +} + +func TestAdmit_CIDRMatch(t *testing.T) { + adm := newSyncedAdmission(MatchersOption([]string{"10.0.0.0/8"})) + defer adm.Close() + + assert.False(t, adm.Admit(context.Background(), "tcp", "10.0.0.1")) + assert.True(t, adm.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestAdmit_InvalidAddr(t *testing.T) { + adm := newSyncedAdmission(MatchersOption([]string{"192.168.1.1"})) + defer adm.Close() + + // Invalid IP should not match and be admitted in blacklist mode + assert.True(t, adm.Admit(context.Background(), "tcp", "invalid")) +} + +// --- parseLine tests --- + +func TestParseLine_Normal(t *testing.T) { + adm := &localAdmission{} + assert.Equal(t, "192.168.1.1", adm.parseLine("192.168.1.1")) +} + +func TestParseLine_WithComment(t *testing.T) { + adm := &localAdmission{} + assert.Equal(t, "192.168.1.1", adm.parseLine("192.168.1.1 # a comment")) +} + +func TestParseLine_CommentOnly(t *testing.T) { + adm := &localAdmission{} + assert.Equal(t, "", adm.parseLine("# comment only")) +} + +func TestParseLine_Empty(t *testing.T) { + adm := &localAdmission{} + assert.Equal(t, "", adm.parseLine("")) +} + +func TestParseLine_Whitespace(t *testing.T) { + adm := &localAdmission{} + assert.Equal(t, "192.168.1.1", adm.parseLine(" 192.168.1.1 ")) +} + +// --- parsePatterns tests --- + +func TestParsePatterns_NilReader(t *testing.T) { + adm := &localAdmission{} + patterns, err := adm.parsePatterns(nil) + assert.NoError(t, err) + assert.Nil(t, patterns) +} + +func TestParsePatterns_Normal(t *testing.T) { + adm := &localAdmission{} + r := strings.NewReader("192.168.1.1\n10.0.0.1\n# comment\n\n fd00::1 \n") + patterns, err := adm.parsePatterns(r) + assert.NoError(t, err) + assert.Equal(t, []string{"192.168.1.1", "10.0.0.1", "fd00::1"}, patterns) +} + +// --- matched tests --- + +func TestMatched_IP(t *testing.T) { + adm := newSyncedAdmission(MatchersOption([]string{"192.168.1.1"})) + defer adm.Close() + + assert.True(t, adm.matched("192.168.1.1")) + assert.False(t, adm.matched("10.0.0.1")) +} + +func TestMatched_CIDR(t *testing.T) { + adm := newSyncedAdmission(MatchersOption([]string{"10.0.0.0/8"})) + defer adm.Close() + + assert.True(t, adm.matched("10.0.0.1")) + assert.False(t, adm.matched("192.168.1.1")) +} + +func TestMatched_BothIPAndCIDR(t *testing.T) { + adm := newSyncedAdmission(MatchersOption([]string{"192.168.1.1", "10.0.0.0/8"})) + defer adm.Close() + + assert.True(t, adm.matched("192.168.1.1")) + assert.True(t, adm.matched("10.0.0.1")) + assert.False(t, adm.matched("172.16.0.1")) +} + +// --- reload tests --- + +func TestReload_WithHostnameResolution(t *testing.T) { + adm := newSyncedAdmission( + MatchersOption([]string{"localhost"}), + LoggerOption(xlogger.Nop()), + ) + defer adm.Close() + + // After reload, localhost should be resolved to 127.0.0.1 or ::1 + err := adm.reload(context.Background()) + assert.NoError(t, err) +} + +// --- load tests --- + +func TestLoad_AllNilLoaders(t *testing.T) { + adm := &localAdmission{ + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Nil(t, patterns) +} + +func TestLoad_FileLoaderWithList(t *testing.T) { + adm := &localAdmission{ + options: options{ + fileLoader: &mockListerLoader{list: []string{"192.168.1.1", "# comment", "10.0.0.1"}}, + }, + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Equal(t, []string{"192.168.1.1", "10.0.0.1"}, patterns) +} + +func TestLoad_FileLoaderWithLoad(t *testing.T) { + adm := &localAdmission{ + options: options{ + fileLoader: &mockLoader{data: "192.168.1.1\n10.0.0.1\n"}, + }, + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Equal(t, []string{"192.168.1.1", "10.0.0.1"}, patterns) +} + +func TestLoad_RedisLoaderWithList(t *testing.T) { + adm := &localAdmission{ + options: options{ + redisLoader: &mockListerLoader{list: []string{"redis-host-1", "redis-host-2"}}, + }, + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Equal(t, []string{"redis-host-1", "redis-host-2"}, patterns) +} + +func TestLoad_RedisLoaderWithLoad(t *testing.T) { + adm := &localAdmission{ + options: options{ + redisLoader: &mockLoader{data: "192.168.2.1\n"}, + }, + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Equal(t, []string{"192.168.2.1"}, patterns) +} + +func TestLoad_HTTPLoader(t *testing.T) { + adm := &localAdmission{ + options: options{ + httpLoader: &mockLoader{data: "10.10.10.1\n10.10.10.2\n"}, + }, + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Equal(t, []string{"10.10.10.1", "10.10.10.2"}, patterns) +} + +func TestLoad_FileLoaderWithListError(t *testing.T) { + adm := &localAdmission{ + options: options{ + fileLoader: &mockListerLoader{listErr: io.ErrUnexpectedEOF}, + }, + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Nil(t, patterns) +} + +func TestLoad_FileLoaderWithLoadError(t *testing.T) { + adm := &localAdmission{ + options: options{ + fileLoader: &mockLoader{loadErr: io.ErrUnexpectedEOF}, + }, + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Nil(t, patterns) +} + +func TestLoad_RedisLoaderWithListError(t *testing.T) { + adm := &localAdmission{ + options: options{ + redisLoader: &mockListerLoader{listErr: io.ErrUnexpectedEOF}, + }, + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Nil(t, patterns) +} + +func TestLoad_RedisLoaderWithLoadError(t *testing.T) { + adm := &localAdmission{ + options: options{ + redisLoader: &mockLoader{loadErr: io.ErrUnexpectedEOF}, + }, + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Nil(t, patterns) +} + +func TestLoad_HTTPLoaderError(t *testing.T) { + adm := &localAdmission{ + options: options{ + httpLoader: &mockLoader{loadErr: io.ErrUnexpectedEOF}, + }, + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Nil(t, patterns) +} + +func TestLoad_FileLoaderWithListParsing(t *testing.T) { + adm := &localAdmission{ + options: options{ + fileLoader: &mockListerLoader{list: []string{" # comment", "", " 192.168.1.1 "}}, + }, + logger: xlogger.Nop(), + } + patterns, err := adm.load(context.Background()) + assert.NoError(t, err) + assert.Equal(t, []string{"192.168.1.1"}, patterns) +} + +// --- periodReload tests --- + +func TestPeriodReload_ZeroPeriod(t *testing.T) { + adm := &localAdmission{ + logger: xlogger.Nop(), + options: options{period: 0}, + } + err := adm.periodReload(context.Background()) + assert.NoError(t, err) +} + +func TestPeriodReload_NegativePeriod(t *testing.T) { + adm := &localAdmission{ + logger: xlogger.Nop(), + options: options{period: -1}, + } + err := adm.periodReload(context.Background()) + assert.NoError(t, err) +} + +func TestPeriodReload_WithContextCancel(t *testing.T) { + // 500ms < 1s triggers the clamping branch (period = time.Second), + // then sleep > 1s to let at least one ticker cycle fire before cancel. + adm := &localAdmission{ + logger: xlogger.Nop(), + options: options{period: 500 * time.Millisecond}, + } + ctx, cancel := context.WithCancel(context.Background()) + + errCh := make(chan error, 1) + go func() { + errCh <- adm.periodReload(ctx) + }() + + time.Sleep(1100 * time.Millisecond) + cancel() + + select { + case err := <-errCh: + assert.Equal(t, context.Canceled, err) + case <-time.After(3 * time.Second): + t.Fatal("periodReload did not return after context cancel") + } +} + +// --- Close tests --- + +func TestClose_NoLoaders(t *testing.T) { + adm := newSyncedAdmission() + err := adm.Close() + assert.NoError(t, err) +} + +func TestClose_WithAllLoaders(t *testing.T) { + fl := &mockLoader{} + rl := &mockLoader{} + hl := &mockLoader{} + + adm := newSyncedAdmission( + FileLoaderOption(fl), + RedisLoaderOption(rl), + HTTPLoaderOption(hl), + ) + err := adm.Close() + assert.NoError(t, err) + + assert.True(t, fl.closed) + assert.True(t, rl.closed) + assert.True(t, hl.closed) +} + +// --- DNS resolution in reload --- + +func TestReload_DNSResolution(t *testing.T) { + adm := &localAdmission{ + ipMatcher: nil, + cidrMatcher: nil, + options: options{ + matchers: []string{"localhost"}, + }, + logger: xlogger.Nop(), + } + err := adm.reload(context.Background()) + assert.NoError(t, err) + + // localhost should resolve to 127.0.0.1 (or ::1) + assert.True(t, adm.matched("127.0.0.1") || adm.matched("::1")) +} + +// --- AdmissionGroup tests --- + +func TestAdmissionGroup_Empty(t *testing.T) { + g := AdmissionGroup() + assert.True(t, g.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestAdmissionGroup_AllAdmit(t *testing.T) { + g := AdmissionGroup( + alwaysAdmit{}, + alwaysAdmit{}, + ) + assert.True(t, g.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestAdmissionGroup_OneDenies(t *testing.T) { + g := AdmissionGroup( + alwaysAdmit{}, + alwaysDeny{}, + ) + assert.False(t, g.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestAdmissionGroup_NilMember(t *testing.T) { + g := AdmissionGroup( + nil, + alwaysAdmit{}, + nil, + ) + assert.True(t, g.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestAdmissionGroup_ShortCircuit(t *testing.T) { + // The first deny should short-circuit, never calling the second. + callCount := 0 + tracking := &trackingAdmission{admit: false, callCount: &callCount} + g := AdmissionGroup(tracking, alwaysAdmit{}) + assert.False(t, g.Admit(context.Background(), "tcp", "192.168.1.1")) + assert.Equal(t, 1, callCount) +} + +// --- Mock types --- + +type mockLoader struct { + data string + loadErr error + closed bool +} + +func (m *mockLoader) Load(ctx context.Context) (io.Reader, error) { + if m.loadErr != nil { + return nil, m.loadErr + } + return strings.NewReader(m.data), nil +} + +func (m *mockLoader) Close() error { + m.closed = true + return nil +} + +type mockListerLoader struct { + list []string + listErr error +} + +func (m *mockListerLoader) Load(ctx context.Context) (io.Reader, error) { + return strings.NewReader(strings.Join(m.list, "\n")), nil +} + +func (m *mockListerLoader) List(ctx context.Context) ([]string, error) { + if m.listErr != nil { + return nil, m.listErr + } + return m.list, nil +} + +func (m *mockListerLoader) Close() error { + return nil +} + +type alwaysAdmit struct{} + +func (alwaysAdmit) Admit(ctx context.Context, network, addr string, opts ...admission.Option) bool { + return true +} + +type alwaysDeny struct{} + +func (alwaysDeny) Admit(ctx context.Context, network, addr string, opts ...admission.Option) bool { + return false +} + +type trackingAdmission struct { + admit bool + callCount *int +} + +func (t *trackingAdmission) Admit(ctx context.Context, network, addr string, opts ...admission.Option) bool { + *t.callCount++ + return t.admit +} + +// Ensure mock types implement the interfaces. +var _ loader.Loader = (*mockLoader)(nil) +var _ loader.Lister = (*mockListerLoader)(nil) +var _ loader.Loader = (*mockListerLoader)(nil) +var _ admission.Admission = alwaysAdmit{} +var _ admission.Admission = alwaysDeny{} + +// TestAdmit_IPv6 tests that IPv6 addresses are handled correctly. +func TestAdmit_IPv6(t *testing.T) { + adm := newSyncedAdmission(MatchersOption([]string{"fd00::1"})) + defer adm.Close() + + assert.False(t, adm.Admit(context.Background(), "tcp", "fd00::1")) + assert.True(t, adm.Admit(context.Background(), "tcp", "fd00::2")) +} + +// TestAdmit_IPv6WithPort tests that IPv6 addresses with port brackets are handled. +func TestAdmit_IPv6WithPort(t *testing.T) { + adm := newSyncedAdmission(MatchersOption([]string{"::1"})) + defer adm.Close() + + // net.SplitHostPort handles [::1]:8080 format + assert.False(t, adm.Admit(context.Background(), "tcp", "[::1]:8080")) +} + +// TestNewAdmission_ReloadErrorInBackground verifies that the initial reload +// in the background goroutine doesn't panic, even if it fails. +func TestNewAdmission_ReloadError(t *testing.T) { + fl := &mockLoader{loadErr: io.ErrUnexpectedEOF} + adm := newSyncedAdmission( + FileLoaderOption(fl), + LoggerOption(xlogger.Nop()), + ) + defer adm.Close() + + // Should still work, just logged a warning + assert.True(t, adm.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +// TestAdmissionGroup_NilMemberFirst tests nil as first member. +func TestAdmissionGroup_NilMemberFirst(t *testing.T) { + g := AdmissionGroup(nil, alwaysDeny{}) + assert.False(t, g.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +// TestAdmissionGroup_AllNil tests all nil members. +func TestAdmissionGroup_AllNil(t *testing.T) { + g := AdmissionGroup(nil, nil) + assert.True(t, g.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +// TestLoggerDefault tests that default logger is set when nil is passed. +func TestLoggerOption_Nil(t *testing.T) { + var opts options + LoggerOption(nil)(&opts) + assert.Nil(t, opts.logger) +} + +// TestAdmit_IpMatcher test case where only IP matcher is valid. +func TestReload_OnlyIPs(t *testing.T) { + adm := &localAdmission{ + options: options{ + matchers: []string{"1.1.1.1", "8.8.8.8"}, + }, + logger: xlogger.Nop(), + } + err := adm.reload(context.Background()) + assert.NoError(t, err) + assert.True(t, adm.matched("1.1.1.1")) + assert.True(t, adm.matched("8.8.8.8")) + assert.False(t, adm.matched("9.9.9.9")) +} + +// TestReload_OnlyCIDRs tests reload with only CIDR patterns. +func TestReload_OnlyCIDRs(t *testing.T) { + adm := &localAdmission{ + options: options{ + matchers: []string{"192.168.0.0/16", "10.0.0.0/8"}, + }, + logger: xlogger.Nop(), + } + err := adm.reload(context.Background()) + assert.NoError(t, err) + assert.True(t, adm.matched("192.168.1.1")) + assert.True(t, adm.matched("10.0.0.1")) + assert.False(t, adm.matched("172.16.0.1")) +} + +// TestAdmit_LoggerDebugOutput tests that debug logging happens on deny. +func TestAdmit_DenyLogs(t *testing.T) { + adm := newSyncedAdmission( + MatchersOption([]string{"192.168.1.1"}), + LoggerOption(xlogger.Nop()), + ) + defer adm.Close() + + // Should not panic; debug log is written + assert.False(t, adm.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +// TestAdmit_WhitelistDenyLogs tests logging for whitelist deny. +func TestAdmit_WhitelistDenyLogs(t *testing.T) { + adm := newSyncedAdmission( + WhitelistOption(true), + MatchersOption([]string{"192.168.1.1"}), + LoggerOption(xlogger.Nop()), + ) + defer adm.Close() + + assert.False(t, adm.Admit(context.Background(), "tcp", "10.0.0.1")) +} + +// TestParsePatterns_WithCommentsAndEmptyLines tests edge cases. +func TestParsePatterns_EdgeCases(t *testing.T) { + adm := &localAdmission{} + r := strings.NewReader("# full comment\n\n # indented comment\n192.168.1.1 # inline") + patterns, err := adm.parsePatterns(r) + assert.NoError(t, err) + assert.Equal(t, []string{"192.168.1.1"}, patterns) +} + +func TestAdmit_NetworkIgnoredForMatching(t *testing.T) { + // Network parameter doesn't affect matching behavior + adm := newSyncedAdmission(MatchersOption([]string{"192.168.1.1"})) + defer adm.Close() + + assert.False(t, adm.Admit(context.Background(), "tcp", "192.168.1.1")) + assert.False(t, adm.Admit(context.Background(), "udp", "192.168.1.1")) + assert.False(t, adm.Admit(context.Background(), "ip4", "192.168.1.1")) +} + +// Compile-time interface check +var _ logger.Logger = xlogger.Nop() + +// Ensure admission.Admission is implemented +var _ admission.Admission = (*localAdmission)(nil) diff --git a/admission/plugin/grpc_test.go b/admission/plugin/grpc_test.go new file mode 100644 index 00000000..d9107ef2 --- /dev/null +++ b/admission/plugin/grpc_test.go @@ -0,0 +1,199 @@ +package admission + +import ( + "context" + "io" + "net" + "os" + "testing" + + "github.com/go-gost/core/admission" + "github.com/go-gost/plugin/admission/proto" + "github.com/go-gost/x/internal/plugin" + xlogger "github.com/go-gost/x/logger" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/status" + + corelogger "github.com/go-gost/core/logger" +) + +func TestMain(m *testing.M) { + corelogger.SetDefault(xlogger.Nop()) + os.Exit(m.Run()) +} + +// helper to set up a real gRPC server on a random port, returning the client conn and server stop func. +func newTestGRPCConn(t *testing.T, srv proto.AdmissionServer) (*grpc.ClientConn, func()) { + t.Helper() + + lis, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + gsrv := grpc.NewServer() + proto.RegisterAdmissionServer(gsrv, srv) + go gsrv.Serve(lis) + + conn, err := grpc.NewClient(lis.Addr().String(), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + require.NoError(t, err) + + return conn, func() { + conn.Close() + gsrv.Stop() + } +} + +func TestGRPCPlugin_FailOpen(t *testing.T) { + // Test fail-open: create directly with nil client to ensure nil path is covered. + p := &grpcPlugin{ + conn: nil, + client: nil, + log: xlogger.Nop(), + } + assert.True(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) + assert.NoError(t, p.Close()) +} + +func TestGRPCPlugin_Admit_Success(t *testing.T) { + conn, cleanup := newTestGRPCConn(t, &testAdmissionServer{admit: true}) + defer cleanup() + + p := &grpcPlugin{ + conn: conn, + client: proto.NewAdmissionClient(conn), + log: xlogger.Nop(), + } + assert.True(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestGRPCPlugin_Admit_Deny(t *testing.T) { + conn, cleanup := newTestGRPCConn(t, &testAdmissionServer{admit: false}) + defer cleanup() + + p := &grpcPlugin{ + conn: conn, + client: proto.NewAdmissionClient(conn), + log: xlogger.Nop(), + } + assert.False(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestGRPCPlugin_Admit_WithService(t *testing.T) { + server := &testAdmissionServer{admit: true} + conn, cleanup := newTestGRPCConn(t, server) + defer cleanup() + + p := &grpcPlugin{ + conn: conn, + client: proto.NewAdmissionClient(conn), + log: xlogger.Nop(), + } + assert.True(t, p.Admit(context.Background(), "tcp", "192.168.1.1", admission.WithService("myservice"))) + assert.Equal(t, "myservice", server.lastService) +} + +func TestGRPCPlugin_Admit_NilClient(t *testing.T) { + p := &grpcPlugin{ + conn: nil, + client: nil, + } + assert.True(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestGRPCPlugin_Close_NilConn(t *testing.T) { + p := &grpcPlugin{ + conn: nil, + client: nil, + } + assert.NoError(t, p.Close()) +} + +func TestGRPCPlugin_Close_NonCloserConn(t *testing.T) { + p := &grpcPlugin{ + conn: &nonCloserConn{}, + client: nil, + } + assert.NoError(t, p.Close()) +} + +func TestGRPCPlugin_Close_WithRealConn(t *testing.T) { + conn, cleanup := newTestGRPCConn(t, &testAdmissionServer{admit: true}) + defer cleanup() + + // conn is *grpc.ClientConn which implements io.Closer + p := &grpcPlugin{ + conn: conn, + client: nil, + log: xlogger.Nop(), + } + assert.NoError(t, p.Close()) +} + +func TestNewGRPCPlugin_RealConn(t *testing.T) { + lis, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + gsrv := grpc.NewServer() + proto.RegisterAdmissionServer(gsrv, &testAdmissionServer{admit: true}) + go gsrv.Serve(lis) + defer gsrv.Stop() + + // Call the real NewGRPCPlugin function. + p := NewGRPCPlugin("test", lis.Addr().String(), plugin.TimeoutOption(500)) + require.NotNil(t, p) + + assert.True(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) + assert.NoError(t, p.(io.Closer).Close()) +} + +func TestGRPCPlugin_Admit_ServerError(t *testing.T) { + conn, cleanup := newTestGRPCConn(t, &errorAdmissionServer{}) + defer cleanup() + + p := &grpcPlugin{ + conn: conn, + client: proto.NewAdmissionClient(conn), + log: xlogger.Nop(), + } + // Server returns an error; Admit should return false. + assert.False(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +// --- test helpers --- + +type testAdmissionServer struct { + proto.UnimplementedAdmissionServer + admit bool + lastService string +} + +func (s *testAdmissionServer) Admit(ctx context.Context, req *proto.AdmissionRequest) (*proto.AdmissionReply, error) { + s.lastService = req.Service + return &proto.AdmissionReply{Ok: s.admit}, nil +} + +type errorAdmissionServer struct { + proto.UnimplementedAdmissionServer +} + +func (s *errorAdmissionServer) Admit(ctx context.Context, req *proto.AdmissionRequest) (*proto.AdmissionReply, error) { + return nil, status.Error(codes.Internal, "internal error") +} + +type nonCloserConn struct{} + +func (n *nonCloserConn) Invoke(ctx context.Context, method string, args any, reply any, opts ...grpc.CallOption) error { + return nil +} + +func (n *nonCloserConn) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) { + return nil, nil +} + +var _ grpc.ClientConnInterface = (*nonCloserConn)(nil) +var _ io.Closer = (*grpcPlugin)(nil) diff --git a/admission/plugin/http_test.go b/admission/plugin/http_test.go new file mode 100644 index 00000000..87063b85 --- /dev/null +++ b/admission/plugin/http_test.go @@ -0,0 +1,161 @@ +package admission + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/go-gost/core/admission" + "github.com/go-gost/x/internal/plugin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNewHTTPPlugin(t *testing.T) { + p := NewHTTPPlugin("test", "http://localhost:9999") + require.NotNil(t, p) + + hp, ok := p.(*httpPlugin) + require.True(t, ok) + assert.Equal(t, "http://localhost:9999", hp.url) + assert.NotNil(t, hp.client) + assert.NotNil(t, hp.log) +} + +func TestHTTPPlugin_Admit_Success(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodPost, r.Method) + assert.Equal(t, "application/json", r.Header.Get("Content-Type")) + + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"ok": true}`)) + })) + defer srv.Close() + + p := NewHTTPPlugin("test", srv.URL) + assert.True(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestHTTPPlugin_Admit_Deny(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"ok": false}`)) + })) + defer srv.Close() + + p := NewHTTPPlugin("test", srv.URL) + assert.False(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestHTTPPlugin_Admit_Non200(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + })) + defer srv.Close() + + p := NewHTTPPlugin("test", srv.URL) + assert.False(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestHTTPPlugin_Admit_InvalidJSON(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`not json`)) + })) + defer srv.Close() + + p := NewHTTPPlugin("test", srv.URL) + assert.False(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestHTTPPlugin_Admit_ConnectionRefused(t *testing.T) { + p := NewHTTPPlugin("test", "http://127.0.0.1:0", plugin.TimeoutOption(0)) + assert.False(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestHTTPPlugin_Admit_NilClient(t *testing.T) { + hp := &httpPlugin{ + url: "http://localhost", + client: nil, + } + assert.False(t, hp.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestHTTPPlugin_Admit_WithServiceAndHeader(t *testing.T) { + var receivedService string + var receivedHeader string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedService = r.URL.Query().Get("check_service") + receivedHeader = r.Header.Get("X-Custom") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"ok": true}`)) + })) + defer srv.Close() + + p := NewHTTPPlugin("test", srv.URL, plugin.HeaderOption(http.Header{ + "X-Custom": []string{"test-value"}, + })) + assert.True(t, p.Admit(context.Background(), "tcp", "192.168.1.1", admission.WithService("myservice"))) + _ = receivedService + _ = receivedHeader +} + +func TestHTTPPlugin_Admit_WithNoHeader(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"ok": true}`)) + })) + defer srv.Close() + + // Create plugin without custom headers + hp := &httpPlugin{ + url: srv.URL, + client: plugin.NewHTTPClient(&plugin.Options{}), + header: nil, + } + assert.True(t, hp.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestHTTPPlugin_Close_NilClient(t *testing.T) { + hp := &httpPlugin{client: nil} + err := hp.Close() + assert.NoError(t, err) +} + +func TestHTTPPlugin_Close_WithClient(t *testing.T) { + hp := &httpPlugin{ + client: plugin.NewHTTPClient(&plugin.Options{}), + } + err := hp.Close() + assert.NoError(t, err) +} + +func TestHTTPPlugin_Admit_WithAllOptions(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "Bearer mytoken", r.Header.Get("Authorization")) + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"ok": true}`)) + })) + defer srv.Close() + + // Test with various plugin options + p := &httpPlugin{ + url: srv.URL, + client: plugin.NewHTTPClient(&plugin.Options{}), + header: http.Header{"Authorization": []string{"Bearer mytoken"}}, + } + assert.True(t, p.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +func TestHTTPPlugin_Admit_BadURL(t *testing.T) { + // A URL that causes http.NewRequestWithContext to fail. + hp := &httpPlugin{ + url: "http://invalid hostname", // space causes NewRequest to fail + client: plugin.NewHTTPClient(&plugin.Options{}), + } + assert.False(t, hp.Admit(context.Background(), "tcp", "192.168.1.1")) +} + +var _ io.Closer = (*httpPlugin)(nil) diff --git a/admission/wrapper/conn_test.go b/admission/wrapper/conn_test.go new file mode 100644 index 00000000..a5f55910 --- /dev/null +++ b/admission/wrapper/conn_test.go @@ -0,0 +1,1153 @@ +package wrapper + +import ( + "context" + "errors" + "io" + "net" + "syscall" + "testing" + "time" + + "github.com/go-gost/core/admission" + xctx "github.com/go-gost/x/ctx" + xio "github.com/go-gost/x/internal/io" + xnet "github.com/go-gost/x/internal/net" + "github.com/go-gost/x/internal/net/udp" + "github.com/stretchr/testify/assert" +) + +// --- WrapConn tests --- + +func TestWrapConn_NilAdmission(t *testing.T) { + c := &fakeConn{} + result := WrapConn(nil, c) + assert.Equal(t, c, result) // should return the original conn +} + +func TestWrapConn_WithAdmission(t *testing.T) { + c := &fakeConn{raddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}} + result := WrapConn(alwaysDenyAdmission{}, c) + assert.IsType(t, &serverConn{}, result) +} + +// --- serverConn.Read tests --- + +func TestServerConn_Read_Admit(t *testing.T) { + c := &fakeConn{ + raddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + } + sc := WrapConn(alwaysAdmitAdmission{}, c) + buf := make([]byte, 10) + n, err := sc.Read(buf) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.Equal(t, "hello", string(buf[:n])) +} + +func TestServerConn_Read_Deny(t *testing.T) { + c := &fakeConn{ + raddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + } + sc := WrapConn(alwaysDenyAdmission{}, c) + buf := make([]byte, 10) + n, err := sc.Read(buf) + assert.Equal(t, io.EOF, err) + assert.Equal(t, 0, n) +} + +func TestServerConn_Read_AdmissionNil(t *testing.T) { + // serverConn with nil admission should just delegate + c := &fakeConn{ + raddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + } + sc := &serverConn{Conn: c, admission: nil} + buf := make([]byte, 10) + n, err := sc.Read(buf) + assert.NoError(t, err) + assert.Equal(t, 5, n) +} + +// --- serverConn.SyscallConn tests --- + +func TestServerConn_SyscallConn_Supported(t *testing.T) { + c := &fakeSyscallConn{fakeConn: fakeConn{raddr: &net.TCPAddr{}}} + sc := WrapConn(alwaysAdmitAdmission{}, c) + rc, err := sc.(*serverConn).SyscallConn() + assert.NoError(t, err) + assert.NotNil(t, rc) +} + +func TestServerConn_SyscallConn_NotSupported(t *testing.T) { + c := &fakeConn{raddr: &net.TCPAddr{}} + sc := WrapConn(alwaysAdmitAdmission{}, c) + _, err := sc.(*serverConn).SyscallConn() + assert.Equal(t, errUnsupport, err) +} + +// --- serverConn.CloseRead tests --- + +func TestServerConn_CloseRead_Supported(t *testing.T) { + c := &fakeCloseReadConn{fakeConn: fakeConn{raddr: &net.TCPAddr{}}} + sc := WrapConn(alwaysAdmitAdmission{}, c) + err := sc.(*serverConn).CloseRead() + assert.NoError(t, err) +} + +func TestServerConn_CloseRead_NotSupported(t *testing.T) { + c := &fakeConn{raddr: &net.TCPAddr{}} + sc := WrapConn(alwaysAdmitAdmission{}, c) + err := sc.(*serverConn).CloseRead() + assert.Equal(t, xio.ErrUnsupported, err) +} + +// --- serverConn.CloseWrite tests --- + +func TestServerConn_CloseWrite_Supported(t *testing.T) { + c := &fakeCloseWriteConn{fakeConn: fakeConn{raddr: &net.TCPAddr{}}} + sc := WrapConn(alwaysAdmitAdmission{}, c) + err := sc.(*serverConn).CloseWrite() + assert.NoError(t, err) +} + +func TestServerConn_CloseWrite_NotSupported(t *testing.T) { + c := &fakeConn{raddr: &net.TCPAddr{}} + sc := WrapConn(alwaysAdmitAdmission{}, c) + err := sc.(*serverConn).CloseWrite() + assert.Equal(t, xio.ErrUnsupported, err) +} + +// --- serverConn.Context tests --- + +func TestServerConn_Context_Supported(t *testing.T) { + testCtx := context.WithValue(context.Background(), "key", "val") + c := &fakeContextConn{fakeConn: fakeConn{raddr: &net.TCPAddr{}}, ctx: testCtx} + sc := WrapConn(alwaysAdmitAdmission{}, c) + result := sc.(*serverConn).Context() + assert.Equal(t, testCtx, result) +} + +func TestServerConn_Context_NotSupported(t *testing.T) { + // Use plainNetConn which does NOT implement xctx.Context. + c := &plainNetConn{raddr: &net.TCPAddr{}} + sc := WrapConn(alwaysAdmitAdmission{}, c) + result := sc.(*serverConn).Context() + assert.Nil(t, result) +} + +// --- WrapPacketConn tests --- + +func TestWrapPacketConn_NilAdmission(t *testing.T) { + pc := &fakePacketConn{} + result := WrapPacketConn(nil, pc) + assert.Equal(t, pc, result) +} + +func TestWrapPacketConn_WithAdmission(t *testing.T) { + pc := &fakePacketConn{} + result := WrapPacketConn(alwaysAdmitAdmission{}, pc) + assert.IsType(t, &packetConn{}, result) +} + +// --- packetConn.ReadFrom tests --- + +func TestPacketConn_ReadFrom_Admit(t *testing.T) { + pc := &fakePacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + } + wrapped := WrapPacketConn(alwaysAdmitAdmission{}, pc) + buf := make([]byte, 10) + n, addr, err := wrapped.ReadFrom(buf) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.NotNil(t, addr) +} + +func TestPacketConn_ReadFrom_Deny(t *testing.T) { + // The first read is denied, the second is a different address that admits. + // But since the admission always denies, any read will loop forever. + // We use a conn that returns error on second call. + callCount := 0 + pc := &dynamicPacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + onRead: func() error { + callCount++ + if callCount >= 2 { + return io.ErrClosedPipe + } + return nil + }, + } + wrapped := WrapPacketConn(alwaysDenyAdmission{}, pc) + buf := make([]byte, 10) + _, _, err := wrapped.ReadFrom(buf) + assert.Equal(t, io.ErrClosedPipe, err) + assert.Equal(t, 2, callCount) +} + +func TestPacketConn_ReadFrom_Error(t *testing.T) { + pc := &fakePacketConn{readErr: io.ErrUnexpectedEOF} + wrapped := WrapPacketConn(alwaysAdmitAdmission{}, pc) + buf := make([]byte, 10) + _, _, err := wrapped.ReadFrom(buf) + assert.Equal(t, io.ErrUnexpectedEOF, err) +} + +// --- packetConn.Context tests --- + +func TestPacketConn_Context_Supported(t *testing.T) { + testCtx := context.WithValue(context.Background(), "key", "val") + pc := &fakeContextPacketConn{fakePacketConn: fakePacketConn{}, ctx: testCtx} + wrapped := WrapPacketConn(alwaysAdmitAdmission{}, pc) + result := wrapped.(*packetConn).Context() + assert.Equal(t, testCtx, result) +} + +func TestPacketConn_Context_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + wrapped := WrapPacketConn(alwaysAdmitAdmission{}, pc) + result := wrapped.(*packetConn).Context() + assert.Nil(t, result) +} + +// --- WrapUDPConn tests --- + +func TestWrapUDPConn(t *testing.T) { + pc := &fakePacketConn{} + result := WrapUDPConn(alwaysAdmitAdmission{}, pc) + assert.IsType(t, &udpConn{}, result) +} + +// --- udpConn.RemoteAddr tests --- + +func TestUDPConn_RemoteAddr_Supported(t *testing.T) { + ra := &net.UDPAddr{IP: net.ParseIP("10.0.0.1"), Port: 9999} + pc := &fakeRemoteAddrPacketConn{fakePacketConn: fakePacketConn{}, raddr: ra} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + result := uc.RemoteAddr() + assert.Equal(t, ra, result) +} + +func TestUDPConn_RemoteAddr_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + result := uc.RemoteAddr() + assert.Nil(t, result) +} + +// --- udpConn.SetReadBuffer tests --- + +func TestUDPConn_SetReadBuffer_Supported(t *testing.T) { + pc := &fakeSetBufferPacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + err := uc.SetReadBuffer(4096) + assert.NoError(t, err) +} + +func TestUDPConn_SetReadBuffer_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + err := uc.SetReadBuffer(4096) + assert.Equal(t, errUnsupport, err) +} + +// --- udpConn.SetWriteBuffer tests --- + +func TestUDPConn_SetWriteBuffer_Supported(t *testing.T) { + pc := &fakeSetBufferPacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + err := uc.SetWriteBuffer(4096) + assert.NoError(t, err) +} + +func TestUDPConn_SetWriteBuffer_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + err := uc.SetWriteBuffer(4096) + assert.Equal(t, errUnsupport, err) +} + +// --- udpConn.Read tests --- + +func TestUDPConn_Read_Supported(t *testing.T) { + pc := &fakeReaderPacketConn{data: []byte("hello")} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + buf := make([]byte, 10) + n, err := uc.Read(buf) + assert.NoError(t, err) + assert.Equal(t, 5, n) +} + +func TestUDPConn_Read_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + _, err := uc.Read(make([]byte, 10)) + assert.Equal(t, errUnsupport, err) +} + +// --- udpConn.ReadFrom tests --- + +func TestUDPConn_ReadFrom_Admit(t *testing.T) { + pc := &fakePacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + } + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + buf := make([]byte, 10) + n, addr, err := uc.ReadFrom(buf) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.NotNil(t, addr) +} + +func TestUDPConn_ReadFrom_Error(t *testing.T) { + pc := &fakePacketConn{readErr: io.ErrClosedPipe} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + _, _, err := uc.ReadFrom(make([]byte, 10)) + assert.Equal(t, io.ErrClosedPipe, err) +} + +// --- udpConn.ReadFromUDP tests --- + +func TestUDPConn_ReadFromUDP_Supported_Admit(t *testing.T) { + pc := &fakeReadUDPPacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + } + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + buf := make([]byte, 10) + n, addr, err := uc.ReadFromUDP(buf) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.NotNil(t, addr) +} + +func TestUDPConn_ReadFromUDP_Supported_Deny(t *testing.T) { + callCount := 0 + pc := &dynamicReadUDPPacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + onRead: func() error { + callCount++ + if callCount >= 2 { + return io.ErrClosedPipe + } + return nil + }, + } + uc := WrapUDPConn(alwaysDenyAdmission{}, pc) + buf := make([]byte, 10) + _, _, err := uc.ReadFromUDP(buf) + assert.Equal(t, io.ErrClosedPipe, err) +} + +func TestUDPConn_ReadFromUDP_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + _, _, err := uc.ReadFromUDP(make([]byte, 10)) + assert.Equal(t, errUnsupport, err) +} + +// --- udpConn.ReadMsgUDP tests --- + +func TestUDPConn_ReadMsgUDP_Supported_Admit(t *testing.T) { + pc := &fakeReadUDPPacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + } + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + buf := make([]byte, 10) + n, oobn, flags, addr, err := uc.ReadMsgUDP(buf, nil) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.Equal(t, 0, oobn) + assert.Equal(t, 0, flags) + assert.NotNil(t, addr) +} + +func TestUDPConn_ReadMsgUDP_Supported_Deny(t *testing.T) { + callCount := 0 + pc := &dynamicReadUDPPacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + onRead: func() error { + callCount++ + if callCount >= 2 { + return io.ErrClosedPipe + } + return nil + }, + } + uc := WrapUDPConn(alwaysDenyAdmission{}, pc) + buf := make([]byte, 10) + _, _, _, _, err := uc.ReadMsgUDP(buf, nil) + assert.Equal(t, io.ErrClosedPipe, err) +} + +func TestUDPConn_ReadMsgUDP_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + _, _, _, _, err := uc.ReadMsgUDP(make([]byte, 10), nil) + assert.Equal(t, errUnsupport, err) +} + +// --- udpConn.Write tests --- + +func TestUDPConn_Write_Supported(t *testing.T) { + pc := &fakeWriterPacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + n, err := uc.Write([]byte("hello")) + assert.NoError(t, err) + assert.Equal(t, 5, n) +} + +func TestUDPConn_Write_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + _, err := uc.Write([]byte("hello")) + assert.Equal(t, errUnsupport, err) +} + +// --- udpConn.WriteTo tests --- + +func TestUDPConn_WriteTo(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + addr := &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234} + n, err := uc.WriteTo([]byte("hello"), addr) + assert.NoError(t, err) + assert.Equal(t, 5, n) +} + +// --- udpConn.WriteToUDP tests --- + +func TestUDPConn_WriteToUDP_Supported(t *testing.T) { + pc := &fakeWriteUDPPacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + addr := &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234} + n, err := uc.WriteToUDP([]byte("hello"), addr) + assert.NoError(t, err) + assert.Equal(t, 5, n) +} + +func TestUDPConn_WriteToUDP_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + addr := &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234} + _, err := uc.WriteToUDP([]byte("hello"), addr) + assert.Equal(t, errUnsupport, err) +} + +// --- udpConn.WriteMsgUDP tests --- + +func TestUDPConn_WriteMsgUDP_Supported(t *testing.T) { + pc := &fakeWriteUDPPacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + addr := &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234} + n, oobn, err := uc.WriteMsgUDP([]byte("hello"), nil, addr) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.Equal(t, 0, oobn) +} + +func TestUDPConn_WriteMsgUDP_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + addr := &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234} + _, _, err := uc.WriteMsgUDP([]byte("hello"), nil, addr) + assert.Equal(t, errUnsupport, err) +} + +// --- udpConn.SyscallConn tests --- + +func TestUDPConn_SyscallConn_Supported(t *testing.T) { + pc := &fakeSyscallPacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + rc, err := uc.SyscallConn() + assert.NoError(t, err) + assert.NotNil(t, rc) +} + +func TestUDPConn_SyscallConn_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + _, err := uc.SyscallConn() + assert.Equal(t, errUnsupport, err) +} + +// --- udpConn.SetDSCP tests --- + +func TestUDPConn_SetDSCP_Supported(t *testing.T) { + pc := &fakeSetDSCPPacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + err := uc.(*udpConn).SetDSCP(46) + assert.NoError(t, err) +} + +func TestUDPConn_SetDSCP_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + // SetDSCP returns nil on unsupported (best-effort) + err := uc.(*udpConn).SetDSCP(46) + assert.NoError(t, err) +} + +// --- udpConn.Context tests --- + +func TestUDPConn_Context_Supported(t *testing.T) { + testCtx := context.WithValue(context.Background(), "key", "val") + pc := &fakeContextPacketConn{fakePacketConn: fakePacketConn{}, ctx: testCtx} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + result := uc.(*udpConn).Context() + assert.Equal(t, testCtx, result) +} + +func TestUDPConn_Context_NotSupported(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + result := uc.(*udpConn).Context() + assert.Nil(t, result) +} + +// --- fake implementations for net.Conn --- + +type fakeConn struct { + raddr net.Addr + data []byte + pos int + ctx context.Context + closed bool +} + +func (c *fakeConn) Read(b []byte) (n int, err error) { + if c.pos >= len(c.data) { + return 0, io.EOF + } + n = copy(b, c.data[c.pos:]) + c.pos += n + return n, nil +} + +func (c *fakeConn) Write(b []byte) (n int, err error) { + return len(b), nil +} + +func (c *fakeConn) Close() error { + c.closed = true + return nil +} + +// Context implements xctx.Context - returns the stored context, or background if nil. +func (c *fakeConn) Context() context.Context { + if c.ctx != nil { + return c.ctx + } + return context.Background() +} + +func (c *fakeConn) LocalAddr() net.Addr { + return &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 9999} +} + +func (c *fakeConn) RemoteAddr() net.Addr { + if c.raddr == nil { + return &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 9999} + } + return c.raddr +} + +func (c *fakeConn) SetDeadline(t time.Time) error { return nil } +func (c *fakeConn) SetReadDeadline(t time.Time) error { return nil } +func (c *fakeConn) SetWriteDeadline(t time.Time) error { return nil } + +// fakeSyscallConn implements syscall.Conn +type fakeSyscallConn struct { + fakeConn +} + +func (c *fakeSyscallConn) SyscallConn() (syscall.RawConn, error) { + return &fakeRawConn{}, nil +} + +type fakeRawConn struct{} + +func (r *fakeRawConn) Control(f func(fd uintptr)) error { f(0); return nil } +func (r *fakeRawConn) Read(f func(fd uintptr) (done bool)) error { f(0); return nil } +func (r *fakeRawConn) Write(f func(fd uintptr) (done bool)) error { f(0); return nil } + +// plainNetConn is a net.Conn that does NOT implement xctx.Context. +type plainNetConn struct { + raddr net.Addr +} + +func (c *plainNetConn) Read(b []byte) (n int, err error) { return 0, nil } +func (c *plainNetConn) Write(b []byte) (n int, err error) { return len(b), nil } +func (c *plainNetConn) Close() error { return nil } +func (c *plainNetConn) LocalAddr() net.Addr { return &net.TCPAddr{} } +func (c *plainNetConn) RemoteAddr() net.Addr { return c.raddr } +func (c *plainNetConn) SetDeadline(t time.Time) error { return nil } +func (c *plainNetConn) SetReadDeadline(t time.Time) error { return nil } +func (c *plainNetConn) SetWriteDeadline(t time.Time) error { return nil } + +// fakeCloseReadConn implements xio.CloseRead +type fakeCloseReadConn struct { + fakeConn +} + +func (c *fakeCloseReadConn) CloseRead() error { return nil } + +// fakeCloseWriteConn implements xio.CloseWrite +type fakeCloseWriteConn struct { + fakeConn +} + +func (c *fakeCloseWriteConn) CloseWrite() error { return nil } + +// fakeContextConn implements xctx.Context +type fakeContextConn struct { + fakeConn + ctx context.Context +} + +func (c *fakeContextConn) Context() context.Context { return c.ctx } + +// --- fake implementations for net.PacketConn --- + +type fakePacketConn struct { + addr net.Addr + data []byte + readErr error +} + +func (pc *fakePacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + if pc.readErr != nil { + return 0, nil, pc.readErr + } + n = copy(p, pc.data) + if pc.addr == nil { + addr = &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 9999} + } else { + addr = pc.addr + } + return n, addr, nil +} + +func (pc *fakePacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + return len(p), nil +} + +func (pc *fakePacketConn) Close() error { return nil } +func (pc *fakePacketConn) LocalAddr() net.Addr { return &net.UDPAddr{} } +func (pc *fakePacketConn) SetDeadline(t time.Time) error { return nil } +func (pc *fakePacketConn) SetReadDeadline(t time.Time) error { return nil } +func (pc *fakePacketConn) SetWriteDeadline(t time.Time) error { return nil } + +// dynamicPacketConn for testing deny loops +type dynamicPacketConn struct { + addr net.Addr + data []byte + onRead func() error +} + +func (pc *dynamicPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + if err := pc.onRead(); err != nil { + return 0, nil, err + } + n = copy(p, pc.data) + return n, pc.addr, nil +} + +func (pc *dynamicPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + return len(p), nil +} + +func (pc *dynamicPacketConn) Close() error { return nil } +func (pc *dynamicPacketConn) LocalAddr() net.Addr { return &net.UDPAddr{} } +func (pc *dynamicPacketConn) SetDeadline(t time.Time) error { return nil } +func (pc *dynamicPacketConn) SetReadDeadline(t time.Time) error { return nil } +func (pc *dynamicPacketConn) SetWriteDeadline(t time.Time) error { return nil } + +// fakeReaderPacketConn implements io.Reader on top of PacketConn +type fakeReaderPacketConn struct { + data []byte + pos int +} + +func (pc *fakeReaderPacketConn) Read(b []byte) (n int, err error) { + if pc.pos >= len(pc.data) { + return 0, io.EOF + } + n = copy(b, pc.data[pc.pos:]) + pc.pos += n + return n, nil +} + +func (pc *fakeReaderPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + return 0, nil, io.EOF +} + +func (pc *fakeReaderPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + return len(p), nil +} + +func (pc *fakeReaderPacketConn) Close() error { return nil } +func (pc *fakeReaderPacketConn) LocalAddr() net.Addr { return &net.UDPAddr{} } +func (pc *fakeReaderPacketConn) SetDeadline(t time.Time) error { return nil } +func (pc *fakeReaderPacketConn) SetReadDeadline(t time.Time) error { return nil } +func (pc *fakeReaderPacketConn) SetWriteDeadline(t time.Time) error { return nil } + +// fakeWriterPacketConn implements io.Writer on top of PacketConn +type fakeWriterPacketConn struct{} + +func (pc *fakeWriterPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + return 0, nil, io.EOF +} + +func (pc *fakeWriterPacketConn) Write(b []byte) (n int, err error) { + return len(b), nil +} + +func (pc *fakeWriterPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + return len(p), nil +} + +func (pc *fakeWriterPacketConn) Close() error { return nil } +func (pc *fakeWriterPacketConn) LocalAddr() net.Addr { return &net.UDPAddr{} } +func (pc *fakeWriterPacketConn) SetDeadline(t time.Time) error { return nil } +func (pc *fakeWriterPacketConn) SetReadDeadline(t time.Time) error { return nil } +func (pc *fakeWriterPacketConn) SetWriteDeadline(t time.Time) error { return nil } + +// fakeReadUDPPacketConn implements udp.ReadUDP +type fakeReadUDPPacketConn struct { + addr *net.UDPAddr + data []byte + readErr error +} + +func (pc *fakeReadUDPPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + if pc.readErr != nil { + return 0, nil, pc.readErr + } + return copy(p, pc.data), pc.addr, nil +} + +func (pc *fakeReadUDPPacketConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) { + if pc.readErr != nil { + return 0, nil, pc.readErr + } + return copy(b, pc.data), pc.addr, nil +} + +func (pc *fakeReadUDPPacketConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) { + if pc.readErr != nil { + return 0, 0, 0, nil, pc.readErr + } + return copy(b, pc.data), 0, 0, pc.addr, nil +} + +func (pc *fakeReadUDPPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + return len(p), nil +} + +func (pc *fakeReadUDPPacketConn) Close() error { return nil } +func (pc *fakeReadUDPPacketConn) LocalAddr() net.Addr { return &net.UDPAddr{} } +func (pc *fakeReadUDPPacketConn) SetDeadline(t time.Time) error { return nil } +func (pc *fakeReadUDPPacketConn) SetReadDeadline(t time.Time) error { return nil } +func (pc *fakeReadUDPPacketConn) SetWriteDeadline(t time.Time) error { return nil } + +// dynamicReadUDPPacketConn for testing deny loops in ReadUDP methods +type dynamicReadUDPPacketConn struct { + addr *net.UDPAddr + data []byte + onRead func() error +} + +func (pc *dynamicReadUDPPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + if err := pc.onRead(); err != nil { + return 0, nil, err + } + return copy(p, pc.data), pc.addr, nil +} + +func (pc *dynamicReadUDPPacketConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) { + if err := pc.onRead(); err != nil { + return 0, nil, err + } + return copy(b, pc.data), pc.addr, nil +} + +func (pc *dynamicReadUDPPacketConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) { + if err := pc.onRead(); err != nil { + return 0, 0, 0, nil, err + } + return copy(b, pc.data), 0, 0, pc.addr, nil +} + +func (pc *dynamicReadUDPPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + return len(p), nil +} + +func (pc *dynamicReadUDPPacketConn) Close() error { return nil } +func (pc *dynamicReadUDPPacketConn) LocalAddr() net.Addr { return &net.UDPAddr{} } +func (pc *dynamicReadUDPPacketConn) SetDeadline(t time.Time) error { return nil } +func (pc *dynamicReadUDPPacketConn) SetReadDeadline(t time.Time) error { return nil } +func (pc *dynamicReadUDPPacketConn) SetWriteDeadline(t time.Time) error { return nil } + +// fakeWriteUDPPacketConn implements udp.WriteUDP +type fakeWriteUDPPacketConn struct{} + +func (pc *fakeWriteUDPPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + return 0, nil, io.EOF +} + +func (pc *fakeWriteUDPPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + return len(p), nil +} + +func (pc *fakeWriteUDPPacketConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) { + return len(b), nil +} + +func (pc *fakeWriteUDPPacketConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) { + return len(b), 0, nil +} + +func (pc *fakeWriteUDPPacketConn) Close() error { return nil } +func (pc *fakeWriteUDPPacketConn) LocalAddr() net.Addr { return &net.UDPAddr{} } +func (pc *fakeWriteUDPPacketConn) SetDeadline(t time.Time) error { return nil } +func (pc *fakeWriteUDPPacketConn) SetReadDeadline(t time.Time) error { return nil } +func (pc *fakeWriteUDPPacketConn) SetWriteDeadline(t time.Time) error { return nil } + +// fakeRemoteAddrPacketConn implements xnet.RemoteAddr +type fakeRemoteAddrPacketConn struct { + fakePacketConn + raddr net.Addr +} + +func (pc *fakeRemoteAddrPacketConn) RemoteAddr() net.Addr { return pc.raddr } + +// fakeSetBufferPacketConn implements xnet.SetBuffer +type fakeSetBufferPacketConn struct { + fakePacketConn +} + +func (pc *fakeSetBufferPacketConn) SetReadBuffer(n int) error { return nil } +func (pc *fakeSetBufferPacketConn) SetWriteBuffer(n int) error { return nil } + +// fakeSetDSCPPacketConn implements xnet.SetDSCP +type fakeSetDSCPPacketConn struct { + fakePacketConn +} + +func (pc *fakeSetDSCPPacketConn) SetDSCP(n int) error { return nil } + +// fakeSyscallPacketConn implements xnet.SyscallConn +type fakeSyscallPacketConn struct { + fakePacketConn +} + +func (pc *fakeSyscallPacketConn) SyscallConn() (syscall.RawConn, error) { + return &fakeRawConn{}, nil +} + +// fakeContextPacketConn implements xctx.Context +type fakeContextPacketConn struct { + fakePacketConn + ctx context.Context +} + +func (pc *fakeContextPacketConn) Context() context.Context { return pc.ctx } + +// --- admission helpers --- + +type alwaysAdmitAdmission struct{} + +func (alwaysAdmitAdmission) Admit(ctx context.Context, network, addr string, opts ...admission.Option) bool { + return true +} + +type alwaysDenyAdmission struct{} + +func (alwaysDenyAdmission) Admit(ctx context.Context, network, addr string, opts ...admission.Option) bool { + return false +} + +// Compile-time checks +var ( + _ net.Conn = (*fakeConn)(nil) + _ net.PacketConn = (*fakePacketConn)(nil) + _ io.Reader = (*fakeReaderPacketConn)(nil) + _ io.Writer = (*fakeWriterPacketConn)(nil) + _ udp.ReadUDP = (*fakeReadUDPPacketConn)(nil) + _ udp.WriteUDP = (*fakeWriteUDPPacketConn)(nil) +) + +// Also test that generic types compile +var ( + _ syscall.Conn = (*fakeSyscallConn)(nil) + _ xio.CloseRead = (*fakeCloseReadConn)(nil) + _ xio.CloseWrite = (*fakeCloseWriteConn)(nil) + _ xctx.Context = (*fakeContextConn)(nil) + _ xnet.SetBuffer = (*fakeSetBufferPacketConn)(nil) + _ xnet.SyscallConn = (*fakeSyscallPacketConn)(nil) + _ xnet.RemoteAddr = (*fakeRemoteAddrPacketConn)(nil) + _ xnet.SetDSCP = (*fakeSetDSCPPacketConn)(nil) + _ xctx.Context = (*fakeContextPacketConn)(nil) +) + +// Test that errors are handled correctly +func TestPacketConn_ReadFrom_AdmitWithNilAdmission(t *testing.T) { + pc := &fakePacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + } + // packetConn with nil admission should skip admission check + wrapped := &packetConn{PacketConn: pc, admission: nil} + buf := make([]byte, 10) + n, addr, err := wrapped.ReadFrom(buf) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.NotNil(t, addr) +} + +func TestUDPConn_ReadFrom_AdmitWithNilAdmission(t *testing.T) { + pc := &fakePacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + } + uc := &udpConn{PacketConn: pc, admission: nil} + buf := make([]byte, 10) + n, addr, err := uc.ReadFrom(buf) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.NotNil(t, addr) +} + +func TestUDPConn_ReadFromUDP_NilAdmission(t *testing.T) { + pc := &fakeReadUDPPacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + } + uc := &udpConn{PacketConn: pc, admission: nil} + buf := make([]byte, 10) + n, addr, err := uc.ReadFromUDP(buf) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.NotNil(t, addr) +} + +func TestUDPConn_ReadMsgUDP_NilAdmission(t *testing.T) { + pc := &fakeReadUDPPacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + } + uc := &udpConn{PacketConn: pc, admission: nil} + buf := make([]byte, 10) + n, _, _, addr, err := uc.ReadMsgUDP(buf, nil) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.NotNil(t, addr) +} + +// Test that admission check is NOT performed on WriteTo (writes are not admission-checked) +func TestUDPConn_WriteTo_NoAdmissionCheck(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysDenyAdmission{}, pc) + addr := &net.UDPAddr{IP: net.ParseIP("10.0.0.1"), Port: 1234} + n, err := uc.WriteTo([]byte("hello"), addr) + // Should succeed even with deny admission; writes are not checked + assert.NoError(t, err) + assert.Equal(t, 5, n) +} + +// Test SetDeadline forwarding +func TestUDPConn_SetDeadline(t *testing.T) { + pc := &fakePacketConn{} + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + assert.NoError(t, uc.SetDeadline(time.Time{})) +} + +// Test syscall.Conn integration via the interface +func TestServerConn_SyscallConn_Integration(t *testing.T) { + sc := &fakeSyscallConn{fakeConn: fakeConn{raddr: &net.TCPAddr{}}} + conn := WrapConn(alwaysAdmitAdmission{}, sc) + + // Try to access syscall.Conn via type assertion on the wrapper + type syscallConnGetter interface { + SyscallConn() (syscall.RawConn, error) + } + scc, ok := conn.(syscallConnGetter) + assert.True(t, ok) + rc, err := scc.SyscallConn() + assert.NoError(t, err) + assert.NotNil(t, rc) +} + +// Test Close method forwarding +func TestServerConn_Close(t *testing.T) { + c := &fakeConn{raddr: &net.TCPAddr{}} + sc := WrapConn(alwaysAdmitAdmission{}, c) + assert.NoError(t, sc.Close()) +} + +// Test LocalAddr forwarding +func TestServerConn_LocalAddr(t *testing.T) { + c := &fakeConn{raddr: &net.TCPAddr{}} + sc := WrapConn(alwaysAdmitAdmission{}, c) + addr := sc.LocalAddr() + assert.NotNil(t, addr) +} + +// Test Write forwarding +func TestServerConn_Write(t *testing.T) { + c := &fakeConn{raddr: &net.TCPAddr{}} + sc := WrapConn(alwaysAdmitAdmission{}, c) + n, err := sc.Write([]byte("hello")) + assert.NoError(t, err) + assert.Equal(t, 5, n) +} + +// Test SetDeadline forwarding for serverConn +func TestServerConn_SetDeadline(t *testing.T) { + c := &fakeConn{raddr: &net.TCPAddr{}} + sc := WrapConn(alwaysAdmitAdmission{}, c) + assert.NoError(t, sc.SetDeadline(time.Time{})) + assert.NoError(t, sc.SetReadDeadline(time.Time{})) + assert.NoError(t, sc.SetWriteDeadline(time.Time{})) +} + +// Test that deny loop ReadFromUDP returns error +func TestUDPConn_ReadFromUDP_DenyLoop(t *testing.T) { + callCount := 0 + pc := &dynamicReadUDPPacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + onRead: func() error { + callCount++ + if callCount >= 3 { + return io.ErrClosedPipe + } + return nil + }, + } + uc := WrapUDPConn(alwaysDenyAdmission{}, pc) + buf := make([]byte, 10) + _, _, err := uc.ReadFromUDP(buf) + assert.Equal(t, io.ErrClosedPipe, err) +} + +func TestUDPConn_ReadMsgUDP_DenyLoop(t *testing.T) { + callCount := 0 + pc := &dynamicReadUDPPacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + onRead: func() error { + callCount++ + if callCount >= 3 { + return io.ErrClosedPipe + } + return nil + }, + } + uc := WrapUDPConn(alwaysDenyAdmission{}, pc) + buf := make([]byte, 10) + _, _, _, _, err := uc.ReadMsgUDP(buf, nil) + assert.Equal(t, io.ErrClosedPipe, err) +} + +// Test dynamicPacketConn deny loop with ReadFrom (same as ReadFromUDP flow) +func TestUDPConn_ReadFrom_DenyLoop(t *testing.T) { + callCount := 0 + pc := &dynamicPacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + onRead: func() error { + callCount++ + if callCount >= 3 { + return io.ErrClosedPipe + } + return nil + }, + } + uc := WrapUDPConn(alwaysDenyAdmission{}, pc) + buf := make([]byte, 10) + _, _, err := uc.ReadFrom(buf) + assert.Equal(t, io.ErrClosedPipe, err) +} + +// Test ReadFromUDP returns error on first read +func TestUDPConn_ReadFromUDP_ReadError(t *testing.T) { + pc := &fakeReadUDPPacketConn{ + readErr: errors.New("custom error"), + } + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + buf := make([]byte, 10) + _, _, err := uc.ReadFromUDP(buf) + assert.EqualError(t, err, "custom error") +} + +// Test ReadMsgUDP returns error on first read +func TestUDPConn_ReadMsgUDP_ReadError(t *testing.T) { + pc := &fakeReadUDPPacketConn{ + readErr: errors.New("custom error"), + } + uc := WrapUDPConn(alwaysAdmitAdmission{}, pc) + buf := make([]byte, 10) + _, _, _, _, err := uc.ReadMsgUDP(buf, nil) + assert.EqualError(t, err, "custom error") +} + +// Test ReadFromUDP with admit but no nil admission check (since WrapUDPConn always sets admission) +func TestUDPConn_ReadFromUDP_AdmitDenyLoopCoverage(t *testing.T) { + callCount := 0 + pc := &dynamicReadUDPPacketConn{ + addr: &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + data: []byte("hello"), + onRead: func() error { + callCount++ + if callCount > 5 { + return io.ErrClosedPipe + } + return nil + }, + } + uc := &udpConn{PacketConn: pc, admission: alwaysDenyAdmission{}} + buf := make([]byte, 10) + _, _, err := uc.ReadFromUDP(buf) + assert.Equal(t, io.ErrClosedPipe, err) +} + +func TestPacketConn_Close(t *testing.T) { + pc := &fakePacketConn{} + wrapped := WrapPacketConn(alwaysAdmitAdmission{}, pc) + assert.NoError(t, wrapped.Close()) +} + +func TestPacketConn_LocalAddr(t *testing.T) { + pc := &fakePacketConn{} + wrapped := WrapPacketConn(alwaysAdmitAdmission{}, pc) + addr := wrapped.LocalAddr() + assert.NotNil(t, addr) +} + +func TestPacketConn_SetDeadline(t *testing.T) { + pc := &fakePacketConn{} + wrapped := WrapPacketConn(alwaysAdmitAdmission{}, pc) + assert.NoError(t, wrapped.SetDeadline(time.Time{})) + assert.NoError(t, wrapped.SetReadDeadline(time.Time{})) + assert.NoError(t, wrapped.SetWriteDeadline(time.Time{})) +} + +func TestPacketConn_WriteTo(t *testing.T) { + pc := &fakePacketConn{} + wrapped := WrapPacketConn(alwaysAdmitAdmission{}, pc) + addr := &net.UDPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234} + n, err := wrapped.WriteTo([]byte("hello"), addr) + assert.NoError(t, err) + assert.Equal(t, 5, n) +} diff --git a/admission/wrapper/listener_test.go b/admission/wrapper/listener_test.go new file mode 100644 index 00000000..30af8a3c --- /dev/null +++ b/admission/wrapper/listener_test.go @@ -0,0 +1,234 @@ +package wrapper + +import ( + "context" + "errors" + "net" + "testing" + + "github.com/go-gost/core/admission" + xctx "github.com/go-gost/x/ctx" + "github.com/stretchr/testify/assert" +) + +// --- WrapListener tests --- + +func TestWrapListener_NilAdmission(t *testing.T) { + ln := &fakeListener{} + result := WrapListener("myservice", nil, ln) + assert.Equal(t, ln, result) // should return the original listener +} + +func TestWrapListener_WithAdmission(t *testing.T) { + ln := &fakeListener{} + result := WrapListener("myservice", alwaysAdmitAdmission{}, ln) + assert.IsType(t, &listener{}, result) +} + +// --- listener.Accept tests --- + +func TestListener_Accept_Admit(t *testing.T) { + c := &fakeConn{raddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}} + ln := &fakeListener{conns: []net.Conn{c}} + wrapped := WrapListener("myservice", alwaysAdmitAdmission{}, ln) + + conn, err := wrapped.Accept() + assert.NoError(t, err) + assert.NotNil(t, conn) + assert.False(t, c.closed) +} + +func TestListener_Accept_Deny(t *testing.T) { + c := &fakeConn{raddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}} + denied := &fakeConn{raddr: &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 5678}} + ln := &fakeListener{conns: []net.Conn{denied, c}} + wrapped := WrapListener("myservice", &denyFirstConnAdmission{}, ln) + + // First connection is denied, second is admitted + conn, err := wrapped.Accept() + assert.NoError(t, err) + assert.Equal(t, c, conn) + assert.True(t, denied.closed) +} + +func TestListener_Accept_Error(t *testing.T) { + ln := &fakeListener{acceptErr: errors.New("accept failed")} + wrapped := WrapListener("myservice", alwaysAdmitAdmission{}, ln) + + _, err := wrapped.Accept() + assert.EqualError(t, err, "accept failed") +} + +func TestListener_Accept_WithContextSrcAddr(t *testing.T) { + srcAddr := &net.TCPAddr{IP: net.ParseIP("1.2.3.4"), Port: 1234} + c := &fakeConn{ + raddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + ctx: xctx.ContextWithSrcAddr(context.Background(), srcAddr), + } + ln := &fakeListener{conns: []net.Conn{c}} + + // Use an admission that records the address it receives + recorder := &recordingAdmission{admit: true} + wrapped := WrapListener("myservice", recorder, ln) + + conn, err := wrapped.Accept() + assert.NoError(t, err) + assert.NotNil(t, conn) + // Should have used the src addr from context (String() includes port) + assert.Equal(t, "1.2.3.4:1234", recorder.lastAddr) +} + +func TestListener_Accept_WithContextNilContext(t *testing.T) { + c := &fakeConn{ + raddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}, + ctx: nil, // returns nil context + } + ln := &fakeListener{conns: []net.Conn{c}} + + recorder := &recordingAdmission{admit: true} + wrapped := WrapListener("myservice", recorder, ln) + + conn, err := wrapped.Accept() + assert.NoError(t, err) + assert.NotNil(t, conn) + // Should fall back to RemoteAddr (String() includes port) + assert.Equal(t, "192.168.1.1:1234", recorder.lastAddr) +} + +func TestListener_Accept_NoContextInterface(t *testing.T) { + // A plain net.Conn without xctx.Context support + c := &plainNetConn{raddr: &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 9999}} + ln := &fakeListener{conns: []net.Conn{c}} + + recorder := &recordingAdmission{admit: true} + wrapped := WrapListener("myservice", recorder, ln) + + conn, err := wrapped.Accept() + assert.NoError(t, err) + assert.NotNil(t, conn) + // Should use RemoteAddr (String() includes port) + assert.Equal(t, "10.0.0.1:9999", recorder.lastAddr) +} + +func TestListener_Accept_MultipleDeniesThenAdmit(t *testing.T) { + denied1 := &fakeConn{raddr: &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 1}} + denied2 := &fakeConn{raddr: &net.TCPAddr{IP: net.ParseIP("10.0.0.2"), Port: 2}} + admitted := &fakeConn{raddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 3}} + ln := &fakeListener{conns: []net.Conn{denied1, denied2, admitted}} + + // Deny the first two, admit the third + count := 0 + conditional := &conditionalAdmission{ + fn: func(ctx context.Context, network, addr string, opts ...admission.Option) bool { + count++ + return count > 2 + }, + } + wrapped := WrapListener("myservice", conditional, ln) + + conn, err := wrapped.Accept() + assert.NoError(t, err) + assert.Equal(t, admitted, conn) + assert.True(t, denied1.closed) + assert.True(t, denied2.closed) + assert.False(t, admitted.closed) +} + +// Test that Close is forwarded to the underlying listener +func TestListener_Close(t *testing.T) { + ln := &fakeListener{} + wrapped := WrapListener("myservice", alwaysAdmitAdmission{}, ln) + err := wrapped.Close() + assert.NoError(t, err) + assert.True(t, ln.closed) +} + +// Test that Addr is forwarded to the underlying listener +func TestListener_Addr(t *testing.T) { + ln := &fakeListener{} + wrapped := WrapListener("myservice", alwaysAdmitAdmission{}, ln) + addr := wrapped.Addr() + assert.NotNil(t, addr) +} + +// --- fake types for listener tests --- + +type fakeListener struct { + conns []net.Conn + pos int + acceptErr error + closed bool +} + +func (ln *fakeListener) Accept() (net.Conn, error) { + if ln.acceptErr != nil { + return nil, ln.acceptErr + } + if ln.pos >= len(ln.conns) { + return nil, errors.New("no more connections") + } + c := ln.conns[ln.pos] + ln.pos++ + return c, nil +} + +func (ln *fakeListener) Close() error { + ln.closed = true + return nil +} + +func (ln *fakeListener) Addr() net.Addr { + return &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 8080} +} + +// --- admission helpers for listener tests --- + +type recordingAdmission struct { + admit bool + lastAddr string +} + +func (r *recordingAdmission) Admit(ctx context.Context, network, addr string, opts ...admission.Option) bool { + r.lastAddr = addr + return r.admit +} + +type denyFirstConnAdmission struct { + count int +} + +func (d *denyFirstConnAdmission) Admit(ctx context.Context, network, addr string, opts ...admission.Option) bool { + d.count++ + return d.count > 1 // deny first, admit rest +} + +type conditionalAdmission struct { + fn func(ctx context.Context, network, addr string, opts ...admission.Option) bool +} + +func (c *conditionalAdmission) Admit(ctx context.Context, network, addr string, opts ...admission.Option) bool { + return c.fn(ctx, network, addr, opts...) +} + +// Compile-time checks +var ( + _ net.Listener = (*fakeListener)(nil) + _ net.Conn = (*fakeConn)(nil) + _ xctx.Context = (*fakeConn)(nil) + _ net.Conn = (*plainNetConn)(nil) +) + +// Test that plainNetConn does NOT implement xctx.Context +func TestPlainNetConn_NoContext(t *testing.T) { + c := &plainNetConn{raddr: &net.TCPAddr{}} + _, ok := any(c).(xctx.Context) + assert.False(t, ok) +} + +// Test that fakeConn DOES implement xctx.Context +func TestFakeConn_HasContext(t *testing.T) { + c := &fakeConn{raddr: &net.TCPAddr{}} + _, ok := any(c).(xctx.Context) + assert.True(t, ok) +} +