test(admission): add unit tests achieving 100% coverage for wrapper, 97.6% for core, 96.4% for plugin

This commit is contained in:
ginuerzh
2026-05-22 19:03:47 +08:00
parent b9d74c1b8e
commit 0c37224530
5 changed files with 2474 additions and 0 deletions
+727
View File
@@ -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)
+199
View File
@@ -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)
+161
View File
@@ -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)
File diff suppressed because it is too large Load Diff
+234
View File
@@ -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)
}