diff --git a/ctx/value.go b/ctx/value.go index 4774bd7d..e6be320d 100644 --- a/ctx/value.go +++ b/ctx/value.go @@ -1,3 +1,5 @@ +// Package ctx provides typed context keys and helpers for passing per-request +// data (session ID, addresses, client ID, hash source) through the handler chain. package ctx import ( @@ -5,20 +7,19 @@ import ( "net" ) +// Context is implemented by connections that carry a [context.Context]. type Context interface { Context() context.Context } -type WithContext interface { - WithContext(ctx context.Context) -} - type srcAddrKey struct{} +// ContextWithSrcAddr returns a copy of ctx that carries the source address addr. func ContextWithSrcAddr(ctx context.Context, addr net.Addr) context.Context { return context.WithValue(ctx, srcAddrKey{}, addr) } +// SrcAddrFromContext returns the source address stored in ctx, or nil if none is set. func SrcAddrFromContext(ctx context.Context) net.Addr { v, _ := ctx.Value(srcAddrKey{}).(net.Addr) return v @@ -26,66 +27,75 @@ func SrcAddrFromContext(ctx context.Context) net.Addr { type dstAddrKey struct{} +// ContextWithDstAddr returns a copy of ctx that carries the destination address addr. func ContextWithDstAddr(ctx context.Context, addr net.Addr) context.Context { return context.WithValue(ctx, dstAddrKey{}, addr) } +// DstAddrFromContext returns the destination address stored in ctx, or nil if none is set. func DstAddrFromContext(ctx context.Context) net.Addr { v, _ := ctx.Value(dstAddrKey{}).(net.Addr) return v } type ( - Sid string - // sidKey saves the session ID. + // Sid is a session identifier carried in the context. + Sid string sidKey struct{} ) +// String implements fmt.Stringer. func (s Sid) String() string { return string(s) } +// ContextWithSid returns a copy of ctx that carries the session ID sid. func ContextWithSid(ctx context.Context, sid Sid) context.Context { return context.WithValue(ctx, sidKey{}, sid) } +// SidFromContext returns the session ID stored in ctx, or the zero value if none is set. func SidFromContext(ctx context.Context) Sid { v, _ := ctx.Value(sidKey{}).(Sid) return v } type ( - // hashKey saves the hash source for Selector. hashKey struct{} - Hash struct { + // Hash carries the hash source used by selectors for sticky load balancing. + Hash struct { Source string } ) -func ContextWithHash(ctx context.Context, hash *Hash) context.Context { - return context.WithValue(ctx, hashKey{}, hash) +// ContextWithHash returns a copy of ctx that carries the hash source h. +func ContextWithHash(ctx context.Context, h *Hash) context.Context { + return context.WithValue(ctx, hashKey{}, h) } +// HashFromContext returns the hash source stored in ctx, or nil if none is set. func HashFromContext(ctx context.Context) *Hash { - if v, _ := ctx.Value(hashKey{}).(*Hash); v != nil { - return v - } - return nil + v, _ := ctx.Value(hashKey{}).(*Hash) + return v } type ( + // ClientID is a client identifier used for hash-based load balancing. ClientID string clientIDKey struct{} ) -func (s ClientID) String() string { - return string(s) +// String implements fmt.Stringer. +func (c ClientID) String() string { + return string(c) } +// ContextWithClientID returns a copy of ctx that carries the client ID clientID. func ContextWithClientID(ctx context.Context, clientID ClientID) context.Context { return context.WithValue(ctx, clientIDKey{}, clientID) } +// ClientIDFromContext returns the client ID stored in ctx, or the zero value if none is set. func ClientIDFromContext(ctx context.Context) ClientID { v, _ := ctx.Value(clientIDKey{}).(ClientID) return v diff --git a/ctx/value_test.go b/ctx/value_test.go new file mode 100644 index 00000000..be4312c1 --- /dev/null +++ b/ctx/value_test.go @@ -0,0 +1,211 @@ +package ctx + +import ( + "context" + "net" + "testing" +) + +func TestContextWithSrcAddr(t *testing.T) { + addr := &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 8080} + ctx := ContextWithSrcAddr(context.Background(), addr) + + got := SrcAddrFromContext(ctx) + if got != addr { + t.Errorf("SrcAddrFromContext() = %v, want %v", got, addr) + } +} + +func TestSrcAddrFromContext_Empty(t *testing.T) { + if got := SrcAddrFromContext(context.Background()); got != nil { + t.Errorf("SrcAddrFromContext(empty) = %v, want nil", got) + } +} + +func TestSrcAddrFromContext_WrongType(t *testing.T) { + ctx := context.WithValue(context.Background(), srcAddrKey{}, "not-an-addr") + if got := SrcAddrFromContext(ctx); got != nil { + t.Errorf("SrcAddrFromContext(wrong type) = %v, want nil", got) + } +} + +func TestSrcAddrFromContext_NilAddr(t *testing.T) { + ctx := ContextWithSrcAddr(context.Background(), nil) + if got := SrcAddrFromContext(ctx); got != nil { + t.Errorf("SrcAddrFromContext(nil addr) = %v, want nil", got) + } +} + +func TestContextWithDstAddr(t *testing.T) { + addr := &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 443} + ctx := ContextWithDstAddr(context.Background(), addr) + + got := DstAddrFromContext(ctx) + if got != addr { + t.Errorf("DstAddrFromContext() = %v, want %v", got, addr) + } +} + +func TestDstAddrFromContext_Empty(t *testing.T) { + if got := DstAddrFromContext(context.Background()); got != nil { + t.Errorf("DstAddrFromContext(empty) = %v, want nil", got) + } +} + +func TestDstAddrFromContext_NilAddr(t *testing.T) { + ctx := ContextWithDstAddr(context.Background(), nil) + if got := DstAddrFromContext(ctx); got != nil { + t.Errorf("DstAddrFromContext(nil addr) = %v, want nil", got) + } +} + +func TestSid_String(t *testing.T) { + sid := Sid("session-123") + if s := sid.String(); s != "session-123" { + t.Errorf("Sid.String() = %q, want %q", s, "session-123") + } +} + +func TestContextWithSid(t *testing.T) { + sid := Sid("abc-123") + ctx := ContextWithSid(context.Background(), sid) + + got := SidFromContext(ctx) + if got != sid { + t.Errorf("SidFromContext() = %v, want %v", got, sid) + } +} + +func TestSidFromContext_Empty(t *testing.T) { + if got := SidFromContext(context.Background()); got != "" { + t.Errorf("SidFromContext(empty) = %q, want empty string", got) + } +} + +func TestSidFromContext_WrongType(t *testing.T) { + ctx := context.WithValue(context.Background(), sidKey{}, 42) + if got := SidFromContext(ctx); got != "" { + t.Errorf("SidFromContext(wrong type) = %q, want empty string", got) + } +} + +func TestContextWithHash(t *testing.T) { + h := &Hash{Source: "client-ip"} + ctx := ContextWithHash(context.Background(), h) + + got := HashFromContext(ctx) + if got != h { + t.Errorf("HashFromContext() = %v, want %v", got, h) + } +} + +func TestHashFromContext_Empty(t *testing.T) { + if got := HashFromContext(context.Background()); got != nil { + t.Errorf("HashFromContext(empty) = %v, want nil", got) + } +} + +func TestHashFromContext_WrongType(t *testing.T) { + ctx := context.WithValue(context.Background(), hashKey{}, "not-a-hash") + if got := HashFromContext(ctx); got != nil { + t.Errorf("HashFromContext(wrong type) = %v, want nil", got) + } +} + +func TestHashFromContext_NilHash(t *testing.T) { + ctx := ContextWithHash(context.Background(), nil) + if got := HashFromContext(ctx); got != nil { + t.Errorf("HashFromContext(nil hash) = %v, want nil", got) + } +} + +func TestClientID_String(t *testing.T) { + cid := ClientID("client-456") + if s := cid.String(); s != "client-456" { + t.Errorf("ClientID.String() = %q, want %q", s, "client-456") + } +} + +func TestContextWithClientID(t *testing.T) { + cid := ClientID("user-789") + ctx := ContextWithClientID(context.Background(), cid) + + got := ClientIDFromContext(ctx) + if got != cid { + t.Errorf("ClientIDFromContext() = %v, want %v", got, cid) + } +} + +func TestClientIDFromContext_Empty(t *testing.T) { + if got := ClientIDFromContext(context.Background()); got != "" { + t.Errorf("ClientIDFromContext(empty) = %q, want empty string", got) + } +} + +func TestClientIDFromContext_WrongType(t *testing.T) { + ctx := context.WithValue(context.Background(), clientIDKey{}, 99) + if got := ClientIDFromContext(ctx); got != "" { + t.Errorf("ClientIDFromContext(wrong type) = %q, want empty string", got) + } +} + +func TestMultipleValuesInContext(t *testing.T) { + ctx := context.Background() + ctx = ContextWithSrcAddr(ctx, &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 8080}) + ctx = ContextWithDstAddr(ctx, &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 443}) + ctx = ContextWithSid(ctx, Sid("multi-session")) + ctx = ContextWithHash(ctx, &Hash{Source: "hash-src"}) + ctx = ContextWithClientID(ctx, ClientID("multi-client")) + + if got := SrcAddrFromContext(ctx); got == nil { + t.Error("SrcAddrFromContext() should not be nil") + } + if got := DstAddrFromContext(ctx); got == nil { + t.Error("DstAddrFromContext() should not be nil") + } + if got := SidFromContext(ctx); got != "multi-session" { + t.Errorf("SidFromContext() = %q, want %q", got, "multi-session") + } + if got := HashFromContext(ctx); got == nil || got.Source != "hash-src" { + t.Errorf("HashFromContext().Source = %v, want %q", got, "hash-src") + } + if got := ClientIDFromContext(ctx); got != "multi-client" { + t.Errorf("ClientIDFromContext() = %q, want %q", got, "multi-client") + } +} + +func TestIndependentKeys(t *testing.T) { + ctx := ContextWithSrcAddr(context.Background(), &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 8080}) + + // Other keys should not be affected. + if got := SidFromContext(ctx); got != "" { + t.Errorf("SidFromContext() = %q, want empty after setting SrcAddr", got) + } + if got := HashFromContext(ctx); got != nil { + t.Errorf("HashFromContext() = %v, want nil after setting SrcAddr", got) + } +} + +type testContextConn struct { + ctx context.Context +} + +func (c *testContextConn) Context() context.Context { + return c.ctx +} + +func TestContextInterface(t *testing.T) { + ctx := context.Background() + ctx = ContextWithSid(ctx, Sid("iface-test")) + + conn := &testContextConn{ctx: ctx} + cc, ok := any(conn).(Context) + if !ok { + t.Fatal("testContextConn should implement Context interface") + } + + inner := cc.Context() + if got := SidFromContext(inner); got != "iface-test" { + t.Errorf("SidFromContext(inner) = %q, want %q", got, "iface-test") + } +}