207 lines
6.4 KiB
Go
207 lines
6.4 KiB
Go
// Package store persists gateway sessions in Postgres.
|
|
package store
|
|
|
|
import (
|
|
"context"
|
|
_ "embed"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
//go:embed schema.sql
|
|
var schema string
|
|
|
|
// ErrNotFound is returned when a token does not match a live session.
|
|
var ErrNotFound = errors.New("store: session not found")
|
|
var ErrDeviceLimit = errors.New("store: device limit reached")
|
|
|
|
type Session struct {
|
|
TokenHash []byte
|
|
EmbyUserID string
|
|
EmbyToken string
|
|
Username string
|
|
ServerID string
|
|
DeviceID string
|
|
DeviceName string
|
|
LastSeenAt time.Time
|
|
}
|
|
|
|
type Store struct {
|
|
pool *pgxpool.Pool
|
|
}
|
|
|
|
func Open(ctx context.Context, databaseURL string) (*Store, error) {
|
|
pool, err := pgxpool.New(ctx, databaseURL)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("store: connect: %w", err)
|
|
}
|
|
if err := pool.Ping(ctx); err != nil {
|
|
pool.Close()
|
|
return nil, fmt.Errorf("store: ping: %w", err)
|
|
}
|
|
return &Store{pool: pool}, nil
|
|
}
|
|
|
|
func (s *Store) Close() { s.pool.Close() }
|
|
|
|
func (s *Store) Ping(ctx context.Context) error { return s.pool.Ping(ctx) }
|
|
|
|
// Migrate applies the schema. It is idempotent, so it runs on every boot.
|
|
func (s *Store) Migrate(ctx context.Context) error {
|
|
if _, err := s.pool.Exec(ctx, schema); err != nil {
|
|
return fmt.Errorf("store: migrate: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// CreateSession enforces a user's device allowance under a per-user transaction lock.
|
|
// Re-authenticating the same stable device replaces its token and never consumes a slot.
|
|
// The replaced hash is returned so its Redis entry can be invalidated immediately.
|
|
func (s *Store) CreateSession(ctx context.Context, sess Session, maxClients int) ([]byte, int, error) {
|
|
tx, err := s.pool.Begin(ctx)
|
|
if err != nil {
|
|
return nil, 0, fmt.Errorf("store: begin session: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
|
|
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, sess.EmbyUserID); err != nil {
|
|
return nil, 0, fmt.Errorf("store: lock user sessions: %w", err)
|
|
}
|
|
|
|
var previousHash []byte
|
|
err = tx.QueryRow(ctx, `
|
|
SELECT token_hash FROM sessions
|
|
WHERE emby_user_id = $1 AND device_id = $2`,
|
|
sess.EmbyUserID, sess.DeviceID).Scan(&previousHash)
|
|
if err != nil && !errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, 0, fmt.Errorf("store: find device session: %w", err)
|
|
}
|
|
existingDevice := err == nil
|
|
|
|
var activeClients int
|
|
if err := tx.QueryRow(ctx,
|
|
`SELECT count(*) FROM sessions WHERE emby_user_id = $1`,
|
|
sess.EmbyUserID).Scan(&activeClients); err != nil {
|
|
return nil, 0, fmt.Errorf("store: count user sessions: %w", err)
|
|
}
|
|
if !existingDevice && activeClients >= maxClients {
|
|
return nil, activeClients, ErrDeviceLimit
|
|
}
|
|
|
|
_, err = tx.Exec(ctx, `
|
|
INSERT INTO sessions (
|
|
token_hash, emby_user_id, emby_token, username, server_id, device_id, device_name
|
|
)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
|
ON CONFLICT (emby_user_id, device_id) DO UPDATE SET
|
|
token_hash = EXCLUDED.token_hash,
|
|
emby_token = EXCLUDED.emby_token,
|
|
username = EXCLUDED.username,
|
|
server_id = EXCLUDED.server_id,
|
|
device_name = EXCLUDED.device_name,
|
|
last_seen_at = now()`,
|
|
sess.TokenHash, sess.EmbyUserID, sess.EmbyToken, sess.Username,
|
|
sess.ServerID, sess.DeviceID, sess.DeviceName)
|
|
if err != nil {
|
|
return nil, 0, fmt.Errorf("store: create session: %w", err)
|
|
}
|
|
if !existingDevice {
|
|
activeClients++
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return nil, 0, fmt.Errorf("store: commit session: %w", err)
|
|
}
|
|
return previousHash, activeClients, nil
|
|
}
|
|
|
|
func (s *Store) SessionByTokenHash(ctx context.Context, hash []byte) (Session, error) {
|
|
var sess Session
|
|
err := s.pool.QueryRow(ctx, `
|
|
SELECT token_hash, emby_user_id, emby_token, username, server_id, device_id, device_name, last_seen_at
|
|
FROM sessions WHERE token_hash = $1`, hash).
|
|
Scan(&sess.TokenHash, &sess.EmbyUserID, &sess.EmbyToken, &sess.Username,
|
|
&sess.ServerID, &sess.DeviceID, &sess.DeviceName, &sess.LastSeenAt)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return Session{}, ErrNotFound
|
|
}
|
|
if err != nil {
|
|
return Session{}, fmt.Errorf("store: load session: %w", err)
|
|
}
|
|
return sess, nil
|
|
}
|
|
|
|
// Touch records activity. Cheap enough to call on the auth path, and it is what the
|
|
// idle-expiry sweep reads.
|
|
func (s *Store) Touch(ctx context.Context, hash []byte) error {
|
|
_, err := s.pool.Exec(ctx, `UPDATE sessions SET last_seen_at = now() WHERE token_hash = $1`, hash)
|
|
return err
|
|
}
|
|
|
|
func (s *Store) DeleteSession(ctx context.Context, hash []byte) error {
|
|
_, err := s.pool.Exec(ctx, `DELETE FROM sessions WHERE token_hash = $1`, hash)
|
|
return err
|
|
}
|
|
|
|
// DeleteIdleSessions retires tokens unused for longer than idle, returning how many went.
|
|
func (s *Store) DeleteIdleSessions(ctx context.Context, idle time.Duration) (int64, error) {
|
|
tag, err := s.pool.Exec(ctx,
|
|
`DELETE FROM sessions WHERE last_seen_at < now() - $1::interval`,
|
|
fmt.Sprintf("%d seconds", int64(idle.Seconds())))
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return tag.RowsAffected(), nil
|
|
}
|
|
|
|
// TrimSessionsToLimit brings data created under an older, more generous policy back
|
|
// within the current allowance. The most recently active devices survive.
|
|
func (s *Store) TrimSessionsToLimit(ctx context.Context, maxClients int) ([]Session, error) {
|
|
rows, err := s.pool.Query(ctx, `
|
|
WITH ranked AS (
|
|
SELECT token_hash,
|
|
row_number() OVER (
|
|
PARTITION BY emby_user_id
|
|
ORDER BY last_seen_at DESC, created_at DESC, token_hash DESC
|
|
) AS device_rank
|
|
FROM sessions
|
|
),
|
|
retired AS (
|
|
DELETE FROM sessions current
|
|
USING ranked
|
|
WHERE current.token_hash = ranked.token_hash
|
|
AND ranked.device_rank > $1
|
|
RETURNING current.token_hash, current.emby_user_id, current.emby_token,
|
|
current.username, current.server_id, current.device_id,
|
|
current.device_name, current.last_seen_at
|
|
)
|
|
SELECT token_hash, emby_user_id, emby_token, username, server_id,
|
|
device_id, device_name, last_seen_at
|
|
FROM retired`,
|
|
maxClients,
|
|
)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("store: trim sessions: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
|
|
var retired []Session
|
|
for rows.Next() {
|
|
var sess Session
|
|
if err := rows.Scan(
|
|
&sess.TokenHash, &sess.EmbyUserID, &sess.EmbyToken, &sess.Username,
|
|
&sess.ServerID, &sess.DeviceID, &sess.DeviceName, &sess.LastSeenAt,
|
|
); err != nil {
|
|
return nil, fmt.Errorf("store: scan trimmed session: %w", err)
|
|
}
|
|
retired = append(retired, sess)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, fmt.Errorf("store: trim sessions rows: %w", err)
|
|
}
|
|
return retired, nil
|
|
}
|