feat: sticky session routing via X-Session-ID header #1998
+364
−12
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,83 @@ | ||
| package auth | ||
|
|
||
| import ( | ||
| "sync" | ||
| "time" | ||
| ) | ||
|
|
||
| // stickyStore maintains session-to-auth bindings so that requests carrying the | ||
| // same session ID are routed to the same auth/account. Entries expire after a | ||
| // configurable TTL and are garbage-collected by Cleanup. | ||
| // | ||
| // maxEntries caps the number of stored bindings to prevent memory exhaustion | ||
| // from untrusted X-CLIProxyAPI-Session-ID headers. | ||
| type stickyStore struct { | ||
| mu sync.RWMutex | ||
| entries map[string]stickyEntry | ||
| maxEntries int | ||
| } | ||
|
|
||
| type stickyEntry struct { | ||
| authID string | ||
| expiresAt time.Time | ||
| } | ||
|
|
||
| // stickyMaxEntries is the upper bound on stored session bindings. | ||
| const stickyMaxEntries = 10_000 | ||
|
|
||
| // StickyMaxSessionIDLen limits the accepted X-CLIProxyAPI-Session-ID length. | ||
| const StickyMaxSessionIDLen = 256 | ||
|
|
||
| func newStickyStore() *stickyStore { | ||
| return &stickyStore{entries: make(map[string]stickyEntry), maxEntries: stickyMaxEntries} | ||
| } | ||
|
|
||
| // Get returns the bound auth ID for the given session, if it exists and has not | ||
| // expired. | ||
| func (s *stickyStore) Get(sessionID string) (string, bool) { | ||
| s.mu.RLock() | ||
| defer s.mu.RUnlock() | ||
| e, ok := s.entries[sessionID] | ||
| if !ok || time.Now().After(e.expiresAt) { | ||
| return "", false | ||
| } | ||
| return e.authID, true | ||
| } | ||
|
|
||
| // Set binds a session to an auth ID with the specified TTL. | ||
| // If the store is at capacity, the write is silently dropped. | ||
| func (s *stickyStore) Set(sessionID, authID string, ttl time.Duration) { | ||
| s.mu.Lock() | ||
| if _, exists := s.entries[sessionID]; !exists && len(s.entries) >= s.maxEntries { | ||
| s.mu.Unlock() | ||
| return | ||
| } | ||
| s.entries[sessionID] = stickyEntry{authID: authID, expiresAt: time.Now().Add(ttl)} | ||
| s.mu.Unlock() | ||
|
Comment on lines
+49
to
+56
Contributor
There was a problem hiding this comment.
The Furthermore, the
Author
There was a problem hiding this comment. Fixed. Added |
||
| } | ||
|
|
||
| // Delete removes the binding for the given session ID. | ||
| func (s *stickyStore) Delete(sessionID string) { | ||
| s.mu.Lock() | ||
| delete(s.entries, sessionID) | ||
| s.mu.Unlock() | ||
| } | ||
|
|
||
| // Cleanup removes all expired entries. | ||
| func (s *stickyStore) Cleanup() { | ||
| now := time.Now() | ||
| s.mu.Lock() | ||
| for k, e := range s.entries { | ||
| if now.After(e.expiresAt) { | ||
| delete(s.entries, k) | ||
| } | ||
| } | ||
| s.mu.Unlock() | ||
| } | ||
|
|
||
| // Len returns the number of entries (including possibly-expired ones). | ||
| func (s *stickyStore) Len() int { | ||
| s.mu.RLock() | ||
| defer s.mu.RUnlock() | ||
| return len(s.entries) | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The
X-Session-IDheader is extracted and used directly as a key in thestickyStorewithout any length validation or sanitization. This untrusted input contributes to the potential memory exhaustion vulnerability in thestickyStore.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed. Added 256-char length limit — headers exceeding it return empty (no binding created).