Use cache with TTL for memory attribute session store

andrewheberle · Sep 20, 2024 · eefe139 · eefe139
1 parent da25c8a
commit eefe139
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 22 deletions.
diff --git a/go.mod b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/crewjam/saml v0.4.14
 	github.com/golang-jwt/jwt/v4 v4.5.0
 	github.com/jackc/pgx/v5 v5.6.0
+	github.com/karlseguin/ccache/v3 v3.0.5
 	github.com/oklog/run v1.1.0
 	github.com/russellhaering/goxmldsig v1.4.0
 	github.com/spf13/cobra v1.8.1

diff --git a/go.sum b/go.sum
@@ -34,6 +34,8 @@ github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFr
 github.com/jonboulle/clockwork v0.2.2/go.mod h1:Pkfl5aHPm1nk2H9h0bjmnJD/BcgbGXUBGnn1kMkgxc8=
 github.com/jonboulle/clockwork v0.4.0 h1:p4Cf1aMWXnXAUh8lVfewRBx1zaTSYKrKMF2g3ST4RZ4=
 github.com/jonboulle/clockwork v0.4.0/go.mod h1:xgRqUGwRcjKCO1vbZUEtSLrqKoPSsUpK7fnezOII0kc=
+github.com/karlseguin/ccache/v3 v3.0.5 h1:hFX25+fxzNjsRlREYsoGNa2LoVEw5mPF8wkWq/UnevQ=
+github.com/karlseguin/ccache/v3 v3.0.5/go.mod h1:qxC372+Qn+IBj8Pe3KvGjHPj0sWwEF7AeZVhsNPZ6uY=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=

diff --git a/pkg/sp/memorystore.go b/pkg/sp/memorystore.go
@@ -3,57 +3,49 @@ package sp
 import (
 	"fmt"
 	"log/slog"
-	"sync"
+	"time"
 
 	"github.com/crewjam/saml/samlsp"
+	"github.com/karlseguin/ccache/v3"
 )
 
 type MemoryAttributeStore struct {
-	store map[string]samlsp.Attributes
-	mu    sync.RWMutex
+	ttl   time.Duration
+	store *ccache.Cache[samlsp.Attributes]
 }
 
-func NewMemoryAttributeStore() (*MemoryAttributeStore, error) {
+func NewMemoryAttributeStore(ttl time.Duration) (*MemoryAttributeStore, error) {
 	return &MemoryAttributeStore{
-		store: make(map[string]samlsp.Attributes),
+		store: ccache.New(ccache.Configure[samlsp.Attributes]()),
 	}, nil
 }
 
 func (s *MemoryAttributeStore) Get(id string) (samlsp.Attributes, error) {
-	s.mu.RLock()
-	defer s.mu.RUnlock()
+	if item := s.store.Get(id); item != nil {
+		slog.Debug("getting attributes from store", "id", id, "attrs", item.Value())
 
-	if attrs, found := s.store[id]; found {
-		slog.Debug("getting attributes from store", "id", id, "attrs", attrs)
-
-		return attrs, nil
+		return item.Value(), nil
 	}
 
 	return nil, fmt.Errorf("not found")
 }
 
 func (s *MemoryAttributeStore) Set(id string, attrs samlsp.Attributes) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
 	if s.store == nil {
-		s.store = make(map[string]samlsp.Attributes)
+		s.store = ccache.New(ccache.Configure[samlsp.Attributes]())
 	}
 
 	slog.Debug("setting attributes in store", "id", id, "attrs", attrs)
 
-	s.store[id] = attrs
+	s.store.Set(id, attrs, s.ttl)
 }
 
 func (s *MemoryAttributeStore) Delete(id string) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
 	if s.store == nil {
 		return
 	}
 
 	slog.Debug("deleting attributes in store", "id", id)
 
-	delete(s.store, id)
+	s.store.Delete(id)
 }
diff --git a/pkg/sp/sp.go b/pkg/sp/sp.go
@@ -65,9 +65,9 @@ func NewServiceProvider(cert, key string, root *url.URL, options ...ServiceProvi
 		return nil, fmt.Errorf("metadata was not set")
 	}
 
-	// set default store
+	// set default store with a 1-hour TTL
 	if serviceProvider.store == nil {
-		serviceProvider.store, _ = NewMemoryAttributeStore()
+		serviceProvider.store, _ = NewMemoryAttributeStore(time.Hour * 1)
 	}
 
 	// samlsp options