124 lines
3.3 KiB
Go
124 lines
3.3 KiB
Go
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)
|
|
}
|