youfak/proxy-pool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add gateway sticky sessions

Browse Source
This commit is contained in:
youfak 2026-08-02 16:47:51 +08:00
parent 5389d6f886
commit 6194e3e673
17 changed files with 900 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
README.md
View File

@ -77,6 +77,10 @@ Proxy Pool 用 Controller 协调这些变化，并让 Gateway 数据面只消费
Snapshot 时 Ready，凭据材料只保留在当前节点内存 View。Controller 内成功提交的
Upstream 启停、Routing 切换和配置发布会向本进程全部在线 Worker 快照流广播刷新；
定时刷新仍作为跨进程收敛与失效保护。
- **Gateway 粘性会话**：可为已认证请求启用 X-Proxy-Session 等配置 Header；
Gateway 以 Client、Routing 和会话值派生本地有界绑定，成功建连后固定到同一
Proxy，并在代理失效、快照移除或转发失败时自动重绑。原始会话值不会写入日志、
指标、PostgreSQL 或 Redis，也不会向目标站点转发。
- **安全边界**：Gateway、Distribution 与 Admin 使用各自的认证语义，并支持
CIDR、可信代理、严格请求解析和敏感信息最小化；Admin 的读写权限与
Distribution 提取权限可按命中凭据分别收敛。Distribution 凭据还可限制单次

34
docs/configuration/reference.md
View File

@ -204,6 +204,35 @@ gateway:
Gateway、Distribution、Admin 与 Provider API 是独立认证边界。改变其中一套
不得连带改变其他入口。
### 3.5 Gateway 粘性会话
gateway.stickySession 让调用方使用一个会话头尽量固定出口 Proxy。绑定键是认证
Client、已匹配 Routing 和会话标识的组合，因此不同 Client 或不同 Routing 使用
相同会话值也不会共享出口。该功能要求 Gateway 认证开启。
~~~yaml
gateway:
auth:
mode: bearer
tokenFile: /run/secrets/gateway-token
stickySession:
enabled: true
header: X-Proxy-Session
ttl: 20s
maxEntries: 100000
~~~
- header 必须是规范 HTTP Header 名；一个请求只能携带一个值，值由字母、数字、
点、下划线、连字符组成且最长 128 字节。
- ttl 是绑定的最长时长；实际到期时间还会被 Proxy 的 usableUntil、代理过期时间
和 allocation safety margin 截短。
- maxEntries 是当前 Gateway Worker 的硬上限。缓存只保存 Proxy ID 和过期时间，
不保存 Proxy 地址、凭据或原始会话标识。
- 会话头仅用于本地选路，Gateway 在转发 HTTP 请求前删除它。Proxy 不再符合当前
Snapshot 或一次代理尝试失败时，绑定会被清除并按正常 Routing 重选。
- 该版本不把 Redis、PostgreSQL 或 gRPC 引入 Gateway 热路径；跨 Worker 粘性需要
上游负载均衡保持 Worker 亲和，后续由控制面快照协议扩展共享恢复能力。
## 4. Worker 控制面
控制面默认关闭；默认配置中的 `8443` 端口预留不表示服务已监听。控制面 Session、
@ -302,6 +331,11 @@ gateway:
retry:
maxAttempts: 2
retryMethods: [GET, HEAD]
stickySession:
enabled: true
header: X-Proxy-Session
ttl: 20s
maxEntries: 100000
destinationPolicy:
denyPrivateNetworks: true
denyLoopback: true

12
docs/design/product-design.md
View File

@ -42,10 +42,14 @@ Gateway Worker 只读取本地不可变快照并维护本地容量计数。供
1. 接入层完成认证、来源识别、限流和目标地址检查。
2. Routing 按配置顺序首条命中。
3. Dispatcher 从本地快照筛选 Upstream、协议、标签、TTL 和健康条件。
4. 原子预留 Proxy 容量，建立到上游代理的连接。
5. 建连成功后转为 Active，传输结束后释放；失败则取消预留。
6. GET/HEAD 仅在响应提交前按策略重试；CONNECT 建立后不重放。
3. 若启用并携带粘性会话头，先在本 Worker 的有界缓存中按认证 Client、Routing
和会话标识查找仍符合当前快照的 Proxy。
4. Dispatcher 从本地快照筛选 Upstream、协议、标签、TTL 和健康条件；失效会话
自动清除并按常规策略重新选择。
5. 原子预留 Proxy 容量，建立到上游代理的连接。
6. 建连成功后转为 Active，并在需要时写入不超过 Proxy 可用期的会话绑定；
传输结束后释放，失败则取消预留并清除该绑定。
7. GET/HEAD 仅在响应提交前按策略重试；CONNECT 建立后不重放。
### 4.2 独占提取

9
docs/security/security-model.md
View File

@ -22,6 +22,10 @@
不把无权请求传入管理 mutation 或 Proxy 提取流程。
- Token Client ID 使用 SHA-256 的 128 位摘要前缀，不把 Token 本身写入领域、
日志或审计键。
- 启用 Gateway 粘性会话时，绑定键由认证 Client、Routing 和会话头共同派生；
原始会话标识不写入 PostgreSQL、Redis、日志或指标，也不会转发给目标站点。
绑定仅保留在当前 Worker 的有界内存中，并在代理不再符合当前 Snapshot、到达
代理可用期或转发失败时失效。
## 3. 目标地址策略
@ -55,8 +59,9 @@ CGNAT、协议转换保留段和已知云元数据端点属于硬拒绝项。私
## 6. 审计
Extraction 审计记录至少包含 requestId、Client、来源、Proxy ID、Upstream、
提取时间和到期时间。日志脱敏不影响审计关联，但审计接口自身必须受 Admin
PostgreSQL 审计只记录 Admin 管理 mutation 的 actor、来源、资源、变更结果和
时间；它不保存 Proxy ID、代理地址、凭据、逐次 Extraction 或 Gateway 请求。
Extraction 的短期幂等结果仅位于 Redis TTL 活动池；审计查询接口必须受 Admin
权限保护并具备保留期限。
## 7. 优雅停机

10
internal/config/config.go
View File

@ -65,6 +65,7 @@ type Listener struct {
Auth Auth `yaml:"auth"`
Limits Limits `yaml:"limits"`
Retry Retry `yaml:"retry"`
StickySession StickySession `yaml:"stickySession"`
DestinationPolicy DestinationPolicy `yaml:"destinationPolicy"`
}
@ -117,6 +118,15 @@ type Retry struct {
RetryMethods []string `yaml:"retryMethods"`
}
// StickySession configures per-Worker, bounded Gateway affinity. It does not
// store Proxy details or session identifiers in PostgreSQL or Redis.
type StickySession struct {
Enabled bool `yaml:"enabled"`
Header string `yaml:"header"`
TTL Duration `yaml:"ttl"`
MaxEntries int `yaml:"maxEntries"`
}
type DestinationPolicy struct {
DenyPrivateNetworks *bool `yaml:"denyPrivateNetworks"`
DenyLoopback *bool `yaml:"denyLoopback"`

44
internal/config/config_test.go
View File

@ -409,6 +409,50 @@ func TestValidateDistributionCredentialClientPolicy(t *testing.T) {
}
}
func TestValidateGatewayStickySession(t *testing.T) {
t.Parallel()
cfg := mustLoadValidConfig(t)
cfg.Gateway.Auth = Auth{Mode: "bearer", Token: "gateway-token"}
cfg.Gateway.StickySession = StickySession{
Enabled: true, Header: "X-Proxy-Session", TTL: Duration(30 * time.Second), MaxEntries: 10_000,
}
cfg.Routing[0].Purpose = "gateway"
if err := Validate(cfg); err != nil {
t.Fatalf("Validate(valid gateway stickySession) error = %v", err)
}
cfg.Gateway.StickySession.Header = ""
if err := Validate(cfg); err == nil || !strings.Contains(err.Error(), "stickySession.header") {
t.Fatalf("Validate(missing stickySession header) error = %v", err)
}
cfg.Gateway.StickySession.Header = "X Proxy Session"
if err := Validate(cfg); err == nil || !strings.Contains(err.Error(), "stickySession.header") {
t.Fatalf("Validate(invalid stickySession header) error = %v", err)
}
cfg.Gateway.StickySession.Header = "X-Proxy-Session"
cfg.Gateway.StickySession.TTL = 0
if err := Validate(cfg); err == nil || !strings.Contains(err.Error(), "stickySession.ttl") {
t.Fatalf("Validate(zero stickySession ttl) error = %v", err)
}
cfg.Gateway.StickySession.TTL = Duration(30 * time.Second)
cfg.Gateway.Auth = Auth{Mode: "none"}
if err := Validate(cfg); err == nil || !strings.Contains(err.Error(), "requires authentication") {
t.Fatalf("Validate(unauthenticated stickySession) error = %v", err)
}
cfg.Gateway.Auth = Auth{Mode: "bearer", Token: "gateway-token"}
cfg.Distribution.StickySession = StickySession{
Enabled: true, Header: "X-Proxy-Session", TTL: Duration(30 * time.Second), MaxEntries: 10,
}
if err := Validate(cfg); err == nil || !strings.Contains(err.Error(), "only supported on gateway") {
t.Fatalf("Validate(distribution stickySession) error = %v", err)
}
}
func TestValidateMetricsListener(t *testing.T) {
t.Parallel()
cfg := mustLoadValidConfig(t)

45
internal/config/validate.go
View File

@ -4,6 +4,7 @@ import (
"fmt"
"math"
"net"
"net/http"
"net/url"
"regexp"
"strconv"
@ -109,11 +110,17 @@ func Validate(cfg *Config) error {
if err := validateGatewayClientPolicies(cfg.Gateway.Auth, gatewayRoutings); err != nil {
return err
}
if err := validateGatewayStickySession(cfg.Gateway); err != nil {
return err
}
}
if cfg.Distribution.Enabled {
if err := validateDistributionClientPolicies(cfg.Distribution.Auth, cfg.Upstreams); err != nil {
return err
}
if cfg.Distribution.StickySession.Enabled {
return fmt.Errorf("validate distribution stickySession: only supported on gateway")
}
clientIdentificationMode := cfg.Distribution.ClientIdentification.Mode
if clientIdentificationMode == "" {
clientIdentificationMode = "sourceIP"
@ -144,9 +151,47 @@ func Validate(cfg *Config) error {
return err
}
}
if cfg.Admin.Enabled && cfg.Admin.StickySession.Enabled {
return fmt.Errorf("validate admin stickySession: only supported on gateway")
}
return nil
}
func validateGatewayStickySession(listener Listener) error {
item := listener.StickySession
if !item.Enabled {
return nil
}
if listener.Auth.Mode == "" || listener.Auth.Mode == "none" {
return fmt.Errorf("validate gateway stickySession: requires authentication")
}
if !validCanonicalHeaderName(item.Header) {
return fmt.Errorf("validate gateway stickySession.header: must be a canonical HTTP header name")
}
if item.TTL.Value() <= 0 {
return fmt.Errorf("validate gateway stickySession.ttl: must be greater than zero")
}
if item.MaxEntries <= 0 || item.MaxEntries > MaximumPoolSize {
return fmt.Errorf("validate gateway stickySession.maxEntries: must be in [1, %d]", MaximumPoolSize)
}
return nil
}
func validCanonicalHeaderName(value string) bool {
if value == "" || http.CanonicalHeaderKey(value) != value {
return false
}
for _, character := range []byte(value) {
if character >= 'a' && character <= 'z' || character >= 'A' && character <= 'Z' ||
character >= '0' && character <= '9' || strings.ContainsRune("!#$%&'*+-.^_|~", rune(character)) ||
character == 96 {
continue
}
return false
}
return true
}
func validateControlPlane(item ControlPlane) error {
if err := validateClientTLS("gatewayTLS", item.GatewayTLS); err != nil {
return err

170
internal/gateway/affinity/table.go Normal file
View File

@ -0,0 +1,170 @@
// Package affinity keeps bounded, per-Worker Gateway session bindings.
//
// A binding key is derived from the authenticated client, Routing and supplied
// session identifier. The raw session value is not retained in process memory.
package affinity
import (
"container/list"
"crypto/sha256"
"encoding/binary"
"errors"
"strings"
"sync"
"time"
)
const maxSessionLength = 128
var ErrInvalidKey = errors.New("invalid gateway affinity key")
// Key is an opaque hash of a client, Routing and validated session identifier.
// It is comparable so callers can retain it for a request lifecycle.
type Key struct {
digest [sha256.Size]byte
initialized bool
}
func (key Key) valid() bool {
return key.initialized
}
// NewKey derives an opaque affinity key. Session identifiers accept the
// conservative token subset commonly used by task and order identifiers.
func NewKey(clientID, routingName, sessionID string) (Key, error) {
if strings.TrimSpace(clientID) == "" || strings.TrimSpace(routingName) == "" ||
!validSessionID(sessionID) {
return Key{}, ErrInvalidKey
}
hasher := sha256.New()
for _, part := range []string{clientID, routingName, sessionID} {
var length [4]byte
binary.BigEndian.PutUint32(length[:], uint32(len(part)))
_, _ = hasher.Write(length[:])
_, _ = hasher.Write([]byte(part))
}
var digest [sha256.Size]byte
copy(digest[:], hasher.Sum(nil))
return Key{digest: digest, initialized: true}, nil
}
func validSessionID(value string) bool {
if len(value) == 0 || len(value) > maxSessionLength {
return false
}
for index := 0; index < len(value); index++ {
character := value[index]
if character >= 'a' && character <= 'z' || character >= 'A' && character <= 'Z' ||
character >= '0' && character <= '9' || character == '.' || character == '_' || character == '-' {
continue
}
return false
}
return true
}
type Options struct {
MaxEntries int
Now func() time.Time
}
type Binding struct {
ProxyID string
ExpiresAt time.Time
}
type entry struct {
key Key
binding Binding
element *list.Element
}
// Table is a fixed-capacity LRU cache. It only stores proxy IDs and expiry
// instants, never Proxy addresses, credentials or raw session identifiers.
type Table struct {
mu sync.Mutex
maxEntries int
now func() time.Time
entries map[Key]*entry
recent *list.List
}
func NewTable(options Options) (*Table, error) {
if options.MaxEntries <= 0 {
return nil, ErrInvalidKey
}
if options.Now == nil {
options.Now = time.Now
}
return &Table{
maxEntries: options.MaxEntries,
now: options.Now,
entries: make(map[Key]*entry, options.MaxEntries),
recent: list.New(),
}, nil
}
func (table *Table) Lookup(key Key) (Binding, bool) {
if table == nil || !key.valid() {
return Binding{}, false
}
now := table.now().UTC()
table.mu.Lock()
defer table.mu.Unlock()
current, found := table.entries[key]
if !found {
return Binding{}, false
}
if !current.binding.ExpiresAt.After(now) {
table.remove(current)
return Binding{}, false
}
table.recent.MoveToFront(current.element)
return current.binding, true
}
func (table *Table) Bind(key Key, proxyID string, expiresAt time.Time) {
if table == nil || !key.valid() || strings.TrimSpace(proxyID) == "" || expiresAt.IsZero() {
return
}
expiresAt = expiresAt.UTC()
now := table.now().UTC()
if !expiresAt.After(now) {
table.Delete(key)
return
}
table.mu.Lock()
defer table.mu.Unlock()
if current, found := table.entries[key]; found {
current.binding = Binding{ProxyID: proxyID, ExpiresAt: expiresAt}
table.recent.MoveToFront(current.element)
return
}
for len(table.entries) >= table.maxEntries {
oldest := table.recent.Back()
if oldest == nil {
return
}
table.remove(oldest.Value.(*entry))
}
current := &entry{key: key, binding: Binding{ProxyID: proxyID, ExpiresAt: expiresAt}}
current.element = table.recent.PushFront(current)
table.entries[key] = current
}
func (table *Table) Delete(key Key) {
if table == nil || !key.valid() {
return
}
table.mu.Lock()
defer table.mu.Unlock()
if current, found := table.entries[key]; found {
table.remove(current)
}
}
func (table *Table) remove(current *entry) {
delete(table.entries, current.key)
table.recent.Remove(current.element)
}

108
internal/gateway/affinity/table_test.go Normal file
View File

@ -0,0 +1,108 @@
package affinity
import (
"testing"
"time"
)
func TestTableKeepsBindingsIsolatedAndExpiresThem(t *testing.T) {
t.Parallel()
now := time.Date(2026, time.August, 2, 9, 0, 0, 0, time.UTC)
table, err := NewTable(Options{
MaxEntries: 2,
Now: func() time.Time { return now },
})
if err != nil {
t.Fatalf("NewTable() error = %v", err)
}
first, err := NewKey("client-a", "route-a", "session-a")
if err != nil {
t.Fatalf("NewKey(first) error = %v", err)
}
second, err := NewKey("client-b", "route-a", "session-a")
if err != nil {
t.Fatalf("NewKey(second) error = %v", err)
}
if first == second {
t.Fatal("same session under distinct clients must have different affinity keys")
}
table.Bind(first, "proxy-a", now.Add(time.Minute))
table.Bind(second, "proxy-b", now.Add(time.Minute))
if binding, ok := table.Lookup(first); !ok || binding.ProxyID != "proxy-a" {
t.Fatalf("Lookup(first) = (%+v, %v), want proxy-a", binding, ok)
}
now = now.Add(time.Minute)
if _, ok := table.Lookup(first); ok {
t.Fatal("Lookup(first) retained an expired binding")
}
}
func TestTableEvictsLeastRecentlyUsedBindingAtCapacity(t *testing.T) {
t.Parallel()
now := time.Date(2026, time.August, 2, 9, 0, 0, 0, time.UTC)
table, err := NewTable(Options{
MaxEntries: 2,
Now: func() time.Time { return now },
})
if err != nil {
t.Fatalf("NewTable() error = %v", err)
}
first := mustKey(t, "client-a", "route-a", "session-a")
second := mustKey(t, "client-a", "route-a", "session-b")
third := mustKey(t, "client-a", "route-a", "session-c")
expiresAt := now.Add(time.Minute)
table.Bind(first, "proxy-a", expiresAt)
table.Bind(second, "proxy-b", expiresAt)
if _, ok := table.Lookup(first); !ok {
t.Fatal("Lookup(first) must retain the binding")
}
table.Bind(third, "proxy-c", expiresAt)
if _, ok := table.Lookup(second); ok {
t.Fatal("Lookup(second) retained least recently used binding after capacity eviction")
}
if binding, ok := table.Lookup(first); !ok || binding.ProxyID != "proxy-a" {
t.Fatalf("Lookup(first) = (%+v, %v), want proxy-a", binding, ok)
}
if binding, ok := table.Lookup(third); !ok || binding.ProxyID != "proxy-c" {
t.Fatalf("Lookup(third) = (%+v, %v), want proxy-c", binding, ok)
}
}
func TestNewKeyRejectsInvalidInputs(t *testing.T) {
t.Parallel()
for _, test := range []struct {
name string
clientID string
routing string
sessionID string
}{
{name: "empty client", routing: "route-a", sessionID: "session-a"},
{name: "empty routing", clientID: "client-a", sessionID: "session-a"},
{name: "empty session", clientID: "client-a", routing: "route-a"},
{name: "oversized session", clientID: "client-a", routing: "route-a", sessionID: string(make([]byte, 129))},
{name: "space session", clientID: "client-a", routing: "route-a", sessionID: "session a"},
} {
t.Run(test.name, func(t *testing.T) {
t.Parallel()
if _, err := NewKey(test.clientID, test.routing, test.sessionID); err == nil {
t.Fatal("NewKey() error = nil, want invalid input")
}
})
}
}
func mustKey(t *testing.T, clientID, routing, sessionID string) Key {
t.Helper()
key, err := NewKey(clientID, routing, sessionID)
if err != nil {
t.Fatalf("NewKey() error = %v", err)
}
return key
}

47
internal/gateway/dispatch/dispatcher.go
View File

@ -12,7 +12,10 @@ import (
"proxy-pool/internal/gateway/snapshot"
)
var ErrNoCandidate = errors.New("no local proxy candidate is available")
var (
ErrNoCandidate = errors.New("no local proxy candidate is available")
ErrPreferredUnavailable = errors.New("preferred local proxy is unavailable")
)
type Request struct {
RoutingName string
@ -25,6 +28,7 @@ type Request struct {
RequiredTags map[string]string
Exclude map[string]struct{}
SafetyMargin time.Duration
PreferredProxyID string
}
type Lease struct {
@ -70,12 +74,32 @@ func (d *Dispatcher) Acquire(request Request) (*Lease, error) {
if request.Now.IsZero() {
request.Now = time.Now().UTC()
}
if request.PreferredProxyID != "" {
return d.acquirePreferred(view, request)
}
if request.Strategy.Type != "" {
return d.acquireRouted(view, request)
}
return d.acquireFromUpstreams(view, request, request.Upstreams)
}
func (d *Dispatcher) acquirePreferred(view *snapshot.View, request Request) (*Lease, error) {
entry, found := view.EntryByID(request.PreferredProxyID, d.query(request, request.Upstreams))
if !found {
return nil, ErrPreferredUnavailable
}
reservation, reserved := entry.Runtime.Reserve()
if !reserved {
return nil, ErrNoCandidate
}
return &Lease{
Proxy: entry.Proxy,
Epoch: view.Epoch,
Version: view.Version,
reserved: reservation,
}, nil
}
// AcquireWait retries local snapshot dispatch at a bounded interval until a
// capacity slot appears, the route timeout expires, or the caller cancels.
func (d *Dispatcher) AcquireWait(ctx context.Context, request Request, timeout time.Duration) (*Lease, error) {
@ -206,15 +230,7 @@ func (d *Dispatcher) selectorState(view *snapshot.View) *routingSelectorState {
}
func (d *Dispatcher) acquireFromUpstreams(view *snapshot.View, request Request, upstreams []string) (*Lease, error) {
selection := view.Select(snapshot.Query{
Now: request.Now,
Scheme: request.Scheme,
Upstreams: upstreams,
RequiredTags: request.RequiredTags,
Exclude: request.Exclude,
SafetyMargin: request.SafetyMargin,
})
selection := view.Select(d.query(request, upstreams))
if selection.Len() == 0 {
return nil, ErrNoCandidate
}
@ -239,6 +255,17 @@ func (d *Dispatcher) acquireFromUpstreams(view *snapshot.View, request Request,
return nil, ErrNoCandidate
}
func (d *Dispatcher) query(request Request, upstreams []string) snapshot.Query {
return snapshot.Query{
Now: request.Now,
Scheme: request.Scheme,
Upstreams: upstreams,
RequiredTags: request.RequiredTags,
Exclude: request.Exclude,
SafetyMargin: request.SafetyMargin,
}
}
func maxInt() int {
return int(^uint(0) >> 1)
}

54
internal/gateway/dispatch/dispatcher_test.go
View File

@ -52,6 +52,60 @@ func TestAcquireFiltersAndReservesLocalCapacity(t *testing.T) {
}
}
func TestAcquireUsesPreferredProxyWithoutScanningOtherCandidates(t *testing.T) {
t.Parallel()
now := time.Date(2026, time.August, 2, 9, 0, 0, 0, time.UTC)
expiresAt := now.Add(time.Minute)
store := snapshot.NewStore("cluster-a", "worker-a")
proxies := []proxyDomain.Proxy{
{ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresAt},
{ID: "proxy-b", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresAt},
}
envelope := snapshot.Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, Proxies: proxies}
envelope.Checksum = snapshot.Checksum(proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(): %v", err)
}
lease, err := New(store).Acquire(Request{
Now: now, Scheme: proxyDomain.SchemeHTTP, Upstreams: []string{"upstream-a"}, PreferredProxyID: "proxy-b",
})
if err != nil {
t.Fatalf("Acquire(): %v", err)
}
if lease.Proxy.ID != "proxy-b" {
t.Fatalf("selected proxy = %q, want proxy-b", lease.Proxy.ID)
}
if err := lease.Cancel(); err != nil {
t.Fatalf("Cancel(): %v", err)
}
}
func TestAcquireReportsUnavailablePreferredProxySeparatelyFromCapacity(t *testing.T) {
t.Parallel()
now := time.Date(2026, time.August, 2, 9, 0, 0, 0, time.UTC)
expiresAt := now.Add(time.Minute)
store := snapshot.NewStore("cluster-a", "worker-a")
proxies := []proxyDomain.Proxy{
{ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresAt},
{ID: "proxy-b", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-b", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresAt},
}
envelope := snapshot.Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, Proxies: proxies}
envelope.Checksum = snapshot.Checksum(proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(): %v", err)
}
_, err := New(store).Acquire(Request{
Now: now, Scheme: proxyDomain.SchemeHTTP, Upstreams: []string{"upstream-a"}, PreferredProxyID: "proxy-b",
})
if !errors.Is(err, ErrPreferredUnavailable) {
t.Fatalf("Acquire() error = %v, want ErrPreferredUnavailable", err)
}
}
func TestAcquireRejectsSnapshotAfterOverallValidityDeadline(t *testing.T) {
now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
expiresAt := now.Add(time.Hour)

12
internal/gateway/server/bootstrap.go
View File

@ -47,11 +47,19 @@ func BuildProtection(listener config.Listener) (Protection, error) {
}
func ConfigFromListener(listener config.Listener) Config {
return Config{
result := Config{
MaxAttempts: listener.Retry.MaxAttempts,
RetryMethods: append([]string(nil), listener.Retry.RetryMethods...),
MaxConcurrentRequests: listener.Limits.MaxConcurrentConnections,
}
if listener.StickySession.Enabled {
result.StickySession = StickySessionConfig{
Header: listener.StickySession.Header,
TTL: listener.StickySession.TTL.Value(),
MaxEntries: listener.StickySession.MaxEntries,
}
}
return result
}
func TargetPolicyFromListener(listener config.Listener) (*policy.TargetPolicy, error) {
@ -82,7 +90,7 @@ func buildConfiguredAuth(listener config.Listener) (Guard, error) {
}
protection, err := httpsecurity.NewFromListener(
authListener,
httpsecurity.ClientSourceIP,
httpsecurity.ClientAuthenticatedOrSourceIP,
httpsecurity.ProxySemantics,
nil,
)

8
internal/gateway/server/bootstrap_test.go
View File

@ -6,6 +6,7 @@ import (
"net/http"
"net/http/httptest"
"testing"
"time"
"proxy-pool/internal/config"
"proxy-pool/internal/gateway/policy"
@ -170,8 +171,13 @@ func TestConfigFromListenerMapsRetryAndConcurrency(t *testing.T) {
result := ConfigFromListener(config.Listener{
Limits: config.Limits{MaxConcurrentConnections: 123},
Retry: config.Retry{MaxAttempts: 2, RetryMethods: []string{"GET", "HEAD"}},
StickySession: config.StickySession{
Enabled: true, Header: "X-Proxy-Session", TTL: config.Duration(time.Minute), MaxEntries: 456,
},
})
if result.MaxConcurrentRequests != 123 || result.MaxAttempts != 2 || len(result.RetryMethods) != 2 {
if result.MaxConcurrentRequests != 123 || result.MaxAttempts != 2 || len(result.RetryMethods) != 2 ||
result.StickySession.Header != "X-Proxy-Session" || result.StickySession.TTL != time.Minute ||
result.StickySession.MaxEntries != 456 {
t.Fatalf("handler config = %+v", result)
}
}

72
internal/gateway/server/handler.go
View File

@ -29,6 +29,7 @@ type Config struct {
SafetyMargin time.Duration
CopyBufferSize int
MaxConcurrentRequests int
StickySession StickySessionConfig
}
const requestClosingBit uint64 = 1 << 63
@ -108,6 +109,7 @@ type Handler struct {
dispatcher Dispatcher
transport ProxyTransport
outcomes OutcomeRecorder
sticky *stickySession
buffers sync.Pool
inFlight chan struct{}
forceClose atomic.Bool
@ -138,6 +140,10 @@ func New(config Config, dependencies Dependencies) (*Handler, error) {
if config.CopyBufferSize <= 0 {
config.CopyBufferSize = 32 << 10
}
sticky, err := newStickySession(config.StickySession)
if err != nil {
return nil, fmt.Errorf("create gateway handler: %w", err)
}
handler := &Handler{
config: config,
guards: [3]Guard{dependencies.Auth, dependencies.Access, dependencies.Admission},
@ -146,6 +152,7 @@ func New(config Config, dependencies Dependencies) (*Handler, error) {
dispatcher: dependencies.Dispatcher,
transport: dependencies.Transport,
outcomes: dependencies.Outcomes,
sticky: sticky,
tunnels: make(map[*activeTunnel]struct{}),
shutdownDone: make(chan struct{}),
}
@ -196,7 +203,12 @@ func (handler *Handler) ServeHTTP(writer http.ResponseWriter, request *http.Requ
writeGatewayError(writer, err)
return
}
handler.connect(writer, request, target, route)
binding, err := handler.prepareStickySession(request, &route)
if err != nil {
writeGatewayError(writer, err)
return
}
handler.connect(writer, request, target, route, binding)
return
}
if request.URL == nil || !request.URL.IsAbs() {
@ -217,7 +229,12 @@ func (handler *Handler) ServeHTTP(writer http.ResponseWriter, request *http.Requ
writeGatewayError(writer, err)
return
}
handler.forwardHTTP(writer, request, target, route)
binding, err := handler.prepareStickySession(request, &route)
if err != nil {
writeGatewayError(writer, err)
return
}
handler.forwardHTTP(writer, request, target, route, binding)
}
func (handler *Handler) connect(
@ -225,6 +242,7 @@ func (handler *Handler) connect(
request *http.Request,
target policy.Authority,
route dispatch.Request,
binding *stickyBinding,
) {
attempts := handler.attemptLimit(request)
excluded := cloneSet(route.Exclude)
@ -234,7 +252,7 @@ func (handler *Handler) connect(
route.Now = time.Now().UTC()
route.Exclude = excluded
route.SafetyMargin = handler.config.SafetyMargin
lease, err := handler.acquireRoute(request.Context(), route)
lease, err := handler.acquireRouteWithStickySession(request.Context(), &route, binding)
if err != nil {
lastErr = err
if errors.Is(err, dispatch.ErrNoCandidate) && route.OnUnavailable == routing.OnUnavailableDirect {
@ -259,6 +277,8 @@ func (handler *Handler) connect(
if err != nil {
handler.recordOutcome(lease.Proxy.ID, route.RoutingName, outcomeDomain.StageProxyHandshake, false, err, started)
finishLease(lease, false)
binding.Clear()
route.PreferredProxyID = ""
excluded[lease.Proxy.ID] = struct{}{}
var responseError *transportDomain.ProxyResponseError
if errors.As(err, &responseError) {
@ -276,9 +296,12 @@ func (handler *Handler) connect(
if err := lease.Commit(); err != nil {
_ = upstream.Close()
finishLease(lease, false)
binding.Clear()
route.PreferredProxyID = ""
lastErr = err
break
}
binding.Bind(lease.Proxy, handler.config.SafetyMargin)
handler.recordOutcome(lease.Proxy.ID, route.RoutingName, outcomeDomain.StageProxyHandshake, true, nil, started)
handler.serveTunnel(writer, request, func() { finishLease(lease, true) }, upstream, lease.Proxy.ID, route.RoutingName)
return
@ -428,6 +451,7 @@ func (handler *Handler) forwardHTTP(
request *http.Request,
target policy.Authority,
route dispatch.Request,
binding *stickyBinding,
) {
attempts := handler.attemptLimit(request)
excluded := cloneSet(route.Exclude)
@ -445,7 +469,7 @@ func (handler *Handler) forwardHTTP(
route.Now = time.Now().UTC()
route.Exclude = excluded
route.SafetyMargin = handler.config.SafetyMargin
lease, err := handler.acquireRoute(request.Context(), route)
lease, err := handler.acquireRouteWithStickySession(request.Context(), &route, binding)
if err != nil {
lastErr = err
if errors.Is(err, dispatch.ErrNoCandidate) && route.OnUnavailable == routing.OnUnavailableDirect {
@ -471,6 +495,7 @@ func (handler *Handler) forwardHTTP(
return err
}
committed.Store(true)
binding.Bind(lease.Proxy, handler.config.SafetyMargin)
return nil
}
started := time.Now().UTC()
@ -478,6 +503,8 @@ func (handler *Handler) forwardHTTP(
if err != nil {
handler.recordOutcome(lease.Proxy.ID, route.RoutingName, outcomeDomain.StageDial, false, err, started)
finishLease(lease, committed.Load())
binding.Clear()
route.PreferredProxyID = ""
excluded[lease.Proxy.ID] = struct{}{}
lastErr = err
continue
@ -486,12 +513,16 @@ func (handler *Handler) forwardHTTP(
if err := commit(); err != nil {
_ = response.Body.Close()
finishLease(lease, false)
binding.Clear()
route.PreferredProxyID = ""
handler.recordOutcome(lease.Proxy.ID, route.RoutingName, outcomeDomain.StageResponseHeaders, false, err, started)
lastErr = err
break
}
}
if response.StatusCode == http.StatusProxyAuthRequired {
binding.Clear()
route.PreferredProxyID = ""
handler.recordOutcome(lease.Proxy.ID, route.RoutingName, outcomeDomain.StageResponseHeaders, false,
&transportDomain.ProxyResponseError{StatusCode: response.StatusCode}, started)
} else {
@ -572,6 +603,39 @@ func (handler *Handler) acquireRoute(ctx context.Context, route dispatch.Request
return waiter.AcquireWait(ctx, route, route.WaitTimeout)
}
func (handler *Handler) prepareStickySession(request *http.Request, route *dispatch.Request) (*stickyBinding, error) {
if handler == nil || route == nil || handler.sticky == nil {
return nil, nil
}
binding, preferredProxyID, err := handler.sticky.prepare(request, route.RoutingName)
if err != nil {
status := http.StatusBadRequest
if errors.Is(err, ErrMissingSessionClient) {
status = http.StatusForbidden
}
return nil, &HTTPError{StatusCode: status, Cause: err}
}
route.PreferredProxyID = preferredProxyID
return binding, nil
}
func (handler *Handler) acquireRouteWithStickySession(
ctx context.Context,
route *dispatch.Request,
binding *stickyBinding,
) (*dispatch.Lease, error) {
if route == nil {
return nil, dispatch.ErrNoCandidate
}
lease, err := handler.acquireRoute(ctx, *route)
if !errors.Is(err, dispatch.ErrPreferredUnavailable) || binding == nil {
return lease, err
}
binding.Clear()
route.PreferredProxyID = ""
return handler.acquireRoute(ctx, *route)
}
func (handler *Handler) directTransport() (DirectTransport, error) {
direct, ok := handler.transport.(DirectTransport)
if !ok {

137
internal/gateway/server/handler_test.go
View File

@ -111,6 +111,143 @@ func TestHandlerRejectsGatewayRoutingOutsideCredentialPolicy(t *testing.T) {
}
}
func TestHandlerPinsAuthenticatedSessionToCommittedProxyAndStripsHeader(t *testing.T) {
t.Parallel()
dispatcher, view := dispatcherWithProxies(t, "proxy-a", "proxy-b")
auth, err := httpsecurity.New(httpsecurity.Config{
Authentication: httpsecurity.Authentication{Mode: httpsecurity.ModeBearer, Token: "gateway-token"},
ClientIdentification: httpsecurity.ClientAuthenticated,
Semantics: httpsecurity.ProxySemantics,
}, nil)
if err != nil {
t.Fatalf("New HTTP protection: %v", err)
}
transport := &fakeTransport{roundTrip: func(
_ context.Context,
_ proxyDomain.Proxy,
request *http.Request,
commit ...func() error,
) (*http.Response, error) {
if request.Header.Get("X-Proxy-Session") != "" {
t.Fatal("sticky session header was forwarded upstream")
}
if err := commit[0](); err != nil {
return nil, err
}
return &http.Response{StatusCode: http.StatusNoContent, Header: make(http.Header), Body: http.NoBody}, nil
}}
handler, err := New(Config{
StickySession: StickySessionConfig{Header: "X-Proxy-Session", TTL: time.Minute, MaxEntries: 100},
}, Dependencies{
Auth: auth,
Targets: fakeTargets{evaluateURL: func(context.Context, string) (policy.Authority, error) {
return policy.Authority{Host: "example.test", Port: 80}, nil
}},
Router: RouteFunc(func(*http.Request) (dispatch.Request, error) {
return dispatch.Request{RoutingName: "catalog", Upstreams: []string{"provider-a"}}, nil
}),
Dispatcher: dispatcher,
Transport: transport,
})
if err != nil {
t.Fatalf("New gateway handler: %v", err)
}
for range 2 {
request := httptest.NewRequest(http.MethodGet, "http://example.test/catalog", nil)
request.Header.Set("Proxy-Authorization", "Bearer gateway-token")
request.Header.Set("X-Proxy-Session", "order-182736")
response := httptest.NewRecorder()
handler.ServeHTTP(response, request)
if response.Code != http.StatusNoContent {
t.Fatalf("status = %d, want 204", response.Code)
}
}
if got := strings.Join(transport.attempts(), ","); got != "proxy-a,proxy-a" {
t.Fatalf("selected proxies = %q, want proxy-a,proxy-a", got)
}
assertNoLeakedCapacity(t, view)
}
func TestHandlerRebindsStickySessionWhenSnapshotDropsBoundProxy(t *testing.T) {
t.Parallel()
store := snapshot.NewStore("cluster-a", "worker-a")
expiresAt := time.Now().Add(time.Minute).UTC()
first := []proxyDomain.Proxy{
{ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "provider-a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresAt},
{ID: "proxy-b", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "provider-a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresAt},
}
applySnapshot := func(version uint64, proxies []proxyDomain.Proxy) {
t.Helper()
envelope := snapshot.Envelope{
ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: version, Full: true, Proxies: proxies,
}
envelope.Checksum = snapshot.Checksum(proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("Apply(%d): %v", version, err)
}
}
applySnapshot(1, first)
auth, err := httpsecurity.New(httpsecurity.Config{
Authentication: httpsecurity.Authentication{Mode: httpsecurity.ModeBearer, Token: "gateway-token"},
ClientIdentification: httpsecurity.ClientAuthenticated,
Semantics: httpsecurity.ProxySemantics,
}, nil)
if err != nil {
t.Fatalf("New HTTP protection: %v", err)
}
transport := &fakeTransport{roundTrip: func(
_ context.Context,
_ proxyDomain.Proxy,
_ *http.Request,
commit ...func() error,
) (*http.Response, error) {
if err := commit[0](); err != nil {
return nil, err
}
return &http.Response{StatusCode: http.StatusNoContent, Header: make(http.Header), Body: http.NoBody}, nil
}}
handler, err := New(Config{
StickySession: StickySessionConfig{Header: "X-Proxy-Session", TTL: time.Minute, MaxEntries: 100},
}, Dependencies{
Auth: auth,
Targets: fakeTargets{evaluateURL: func(context.Context, string) (policy.Authority, error) {
return policy.Authority{Host: "example.test", Port: 80}, nil
}},
Router: RouteFunc(func(*http.Request) (dispatch.Request, error) {
return dispatch.Request{RoutingName: "catalog", Upstreams: []string{"provider-a"}}, nil
}),
Dispatcher: dispatch.New(store),
Transport: transport,
})
if err != nil {
t.Fatalf("New gateway handler: %v", err)
}
serve := func() {
t.Helper()
request := httptest.NewRequest(http.MethodGet, "http://example.test/catalog", nil)
request.Header.Set("Proxy-Authorization", "Bearer gateway-token")
request.Header.Set("X-Proxy-Session", "order-182736")
response := httptest.NewRecorder()
handler.ServeHTTP(response, request)
if response.Code != http.StatusNoContent {
t.Fatalf("status = %d, want 204", response.Code)
}
}
serve()
applySnapshot(2, first[1:])
serve()
if got := strings.Join(transport.attempts(), ","); got != "proxy-a,proxy-b" {
t.Fatalf("selected proxies = %q, want proxy-a,proxy-b", got)
}
assertNoLeakedCapacity(t, store.Current())
}
func TestHandlerRetriesGETWithAnotherProxyBeforeResponseCommit(t *testing.T) {
t.Parallel()

123
internal/gateway/server/sticky.go Normal file
View File

@ -0,0 +1,123 @@
package server
import (
"errors"
"fmt"
"net/http"
"strings"
"time"
proxyDomain "proxy-pool/internal/domain/proxy"
"proxy-pool/internal/gateway/affinity"
"proxy-pool/internal/platform/httpsecurity"
)
var (
ErrInvalidStickySession = errors.New("invalid gateway sticky session")
ErrMissingSessionClient = errors.New("gateway sticky session requires an authenticated client")
)
// StickySessionConfig configures a bounded per-Worker affinity table.
type StickySessionConfig struct {
Header string
TTL time.Duration
MaxEntries int
}
type stickySession struct {
header string
ttl time.Duration
table *affinity.Table
}
type stickyBinding struct {
session *stickySession
key affinity.Key
}
func newStickySession(config StickySessionConfig) (*stickySession, error) {
if config.Header == "" && config.TTL == 0 && config.MaxEntries == 0 {
return nil, nil
}
if !validStickyHeaderName(config.Header) ||
config.TTL <= 0 || config.MaxEntries <= 0 {
return nil, ErrInvalidStickySession
}
table, err := affinity.NewTable(affinity.Options{MaxEntries: config.MaxEntries})
if err != nil {
return nil, fmt.Errorf("create gateway sticky session table: %w", err)
}
return &stickySession{header: config.Header, ttl: config.TTL, table: table}, nil
}
func validStickyHeaderName(value string) bool {
if value == "" || http.CanonicalHeaderKey(value) != value {
return false
}
for _, character := range []byte(value) {
if character >= 'a' && character <= 'z' || character >= 'A' && character <= 'Z' ||
character >= '0' && character <= '9' || strings.ContainsRune("!#$%&'*+-.^_|~", rune(character)) ||
character == 96 {
continue
}
return false
}
return true
}
func (session *stickySession) prepare(request *http.Request, routingName string) (*stickyBinding, string, error) {
if session == nil {
return nil, "", nil
}
values := request.Header.Values(session.header)
request.Header.Del(session.header)
if len(values) == 0 {
return nil, "", nil
}
if len(values) != 1 {
return nil, "", ErrInvalidStickySession
}
identity, authenticated := httpsecurity.IdentityFromRequest(request)
if !authenticated || identity.ClientID == "" {
return nil, "", ErrMissingSessionClient
}
key, err := affinity.NewKey(identity.ClientID, routingName, values[0])
if err != nil {
return nil, "", ErrInvalidStickySession
}
binding := &stickyBinding{session: session, key: key}
existing, found := session.table.Lookup(key)
if !found {
return binding, "", nil
}
return binding, existing.ProxyID, nil
}
func (binding *stickyBinding) Bind(selected proxyDomain.Proxy, safetyMargin time.Duration) {
if binding == nil || binding.session == nil {
return
}
now := time.Now().UTC()
expiresAt := now.Add(binding.session.ttl)
if selected.UsableUntil != nil && selected.UsableUntil.Before(expiresAt) {
expiresAt = selected.UsableUntil.UTC()
}
if selected.ExpiresAt != nil {
proxyExpiry := selected.ExpiresAt.UTC().Add(-safetyMargin)
if proxyExpiry.Before(expiresAt) {
expiresAt = proxyExpiry
}
}
if !expiresAt.After(now) {
binding.Clear()
return
}
binding.session.table.Bind(binding.key, selected.ID, expiresAt)
}
func (binding *stickyBinding) Clear() {
if binding == nil || binding.session == nil {
return
}
binding.session.table.Delete(binding.key)
}

24
internal/gateway/snapshot/store.go
View File

@ -138,6 +138,7 @@ type View struct {
credentials map[string]Credential
all []int
byID map[string]int
byScheme map[proxyDomain.Scheme][]int
byUpstream map[string][]int
byTag map[string][]int
@ -493,6 +494,27 @@ func (s Selection) EntryAt(index int) (Entry, bool) {
return entry, true
}
// EntryByID resolves one exact current Snapshot entry while preserving the
// normal allocation filters. It lets affinity routing avoid scanning a large
// local pool for a previously selected Proxy.
func (v *View) EntryByID(proxyID string, query Query) (Entry, bool) {
if query.Now.IsZero() {
query.Now = time.Now().UTC()
}
if v == nil || proxyID == "" || (!v.ValidUntil.IsZero() && !v.ValidUntil.After(query.Now)) {
return Entry{}, false
}
index, found := v.byID[proxyID]
if !found || index < 0 || index >= len(v.Entries) {
return Entry{}, false
}
entry := v.Entries[index]
if !matchesQuery(entry.Proxy, query) {
return Entry{}, false
}
return entry, true
}
func Checksum(proxies []proxyDomain.Proxy) string {
return ChecksumWithRouting(proxies, nil)
}
@ -604,11 +626,13 @@ func (v *View) buildIndexes() {
return
}
v.all = make([]int, len(v.Entries))
v.byID = make(map[string]int, len(v.Entries))
v.byScheme = make(map[proxyDomain.Scheme][]int)
v.byUpstream = make(map[string][]int)
v.byTag = make(map[string][]int)
for index, entry := range v.Entries {
v.all[index] = index
v.byID[entry.Proxy.ID] = index
v.byScheme[entry.Proxy.Scheme] = append(v.byScheme[entry.Proxy.Scheme], index)
v.byUpstream[entry.Proxy.SourceUpstream] = append(v.byUpstream[entry.Proxy.SourceUpstream], index)
for key, value := range entry.Proxy.Tags {