diff --git a/modules/pier/session.go b/modules/pier/session.go index a1b8aca7..e9a6f5a6 100644 --- a/modules/pier/session.go +++ b/modules/pier/session.go @@ -338,10 +338,24 @@ func (sr *sessionRouter) persist(ctx context.Context, ls *localSession) { if err != nil { return } - if err := sr.p.store.Put(ctx, ls.id, Record{Binding: ls.binding, State: raw, Created: ls.created}, sr.p.cfg.IdleTTL); err != nil { + rec := Record{Binding: ls.binding, State: raw, Created: ls.created} + stored := true + // The record is only rewritten, never created here: a session that ended + // while this request ran (DELETE, revocation, deletion, expiry) stays + // ended instead of being written back by its last request. + if u, ok := sr.p.store.(updater); ok { + stored, err = u.Update(ctx, ls.id, rec, sr.p.cfg.IdleTTL) + } else { + err = sr.p.store.Put(ctx, ls.id, rec, sr.p.cfg.IdleTTL) + } + if err != nil { logger().Warn("pier: session state not stored", "session", sessionPrefix(ls.id), "error", err) return } + if !stored { + sr.closeLocal(ls.id) + return + } if complete { ls.persisted.Store(true) } diff --git a/modules/pier/store.go b/modules/pier/store.go index 1a166389..136d4e3a 100644 --- a/modules/pier/store.go +++ b/modules/pier/store.go @@ -44,3 +44,11 @@ type Store interface { Publish(ctx context.Context, ev Event) error Subscribe(ctx context.Context, fn func(Event)) error } + +// updater is implemented by the stores pier ships: Update rewrites an +// existing record and never creates one (Redis SET XX), answering false when +// the record is gone, so a request that finishes after its session ended +// (DELETE, revocation, deletion, expiry) cannot write the record back. +type updater interface { + Update(ctx context.Context, id string, rec Record, ttl time.Duration) (bool, error) +} diff --git a/modules/pier/store_memory.go b/modules/pier/store_memory.go index 39425e34..fcf0ae6a 100644 --- a/modules/pier/store_memory.go +++ b/modules/pier/store_memory.go @@ -52,6 +52,19 @@ func (m *memoryStore) Get(_ context.Context, id string) (Record, bool, error) { return r.rec, true, nil } +// Update rewrites an existing record; a missing or expired one stays gone. +func (m *memoryStore) Update(_ context.Context, id string, rec Record, ttl time.Duration) (bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + r, ok := m.records[id] + if !ok || !r.expires.After(m.now()) { + delete(m.records, id) + return false, nil + } + m.records[id] = memoryRecord{rec: rec, expires: m.now().Add(ttl)} + return true, nil +} + func (m *memoryStore) Touch(_ context.Context, id string, ttl time.Duration) error { m.mu.Lock() defer m.mu.Unlock() diff --git a/modules/pier/store_redis.go b/modules/pier/store_redis.go index bc03150f..eec5a1b1 100644 --- a/modules/pier/store_redis.go +++ b/modules/pier/store_redis.go @@ -51,6 +51,27 @@ func (s *redisStore) Put(ctx context.Context, id string, rec Record, ttl time.Du return err } +// Update rewrites an existing session record (SET XX) and refreshes the +// binding set; a record that is gone stays gone. +func (s *redisStore) Update(ctx context.Context, id string, rec Record, ttl time.Duration) (bool, error) { + raw, err := json.Marshal(rec) + if err != nil { + return false, err + } + err = s.client.SetArgs(ctx, s.sessionKey(id), raw, redis.SetArgs{Mode: "XX", TTL: ttl}).Err() + if errors.Is(err, redis.Nil) { + return false, nil + } + if err != nil { + return false, err + } + pipe := s.client.TxPipeline() + pipe.SAdd(ctx, s.bindingKey(rec.Binding), id) + pipe.PExpire(ctx, s.bindingKey(rec.Binding), ttl) + _, err = pipe.Exec(ctx) + return true, err +} + func (s *redisStore) Get(ctx context.Context, id string) (Record, bool, error) { raw, err := s.client.Get(ctx, s.sessionKey(id)).Bytes() if errors.Is(err, redis.Nil) { @@ -198,6 +219,18 @@ func (l *lazyStore) Get(ctx context.Context, id string) (Record, bool, error) { return s.Get(ctx, id) } +// Update forwards to the opened store's Update (Put for a store without one). +func (l *lazyStore) Update(ctx context.Context, id string, rec Record, ttl time.Duration) (bool, error) { + s, err := l.get(ctx) + if err != nil { + return false, err + } + if u, ok := s.(updater); ok { + return u.Update(ctx, id, rec, ttl) + } + return true, s.Put(ctx, id, rec, ttl) +} + func (l *lazyStore) Touch(ctx context.Context, id string, ttl time.Duration) error { s, err := l.get(ctx) if err != nil {