proxy-pool/internal/gateway/affinity/table.go

// 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)
}