memby/server/internal/store/store.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
	ClientVersion  string
	ClientProtocol 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) }

// RecordSearch stores a normalized query for future per-user ranking analysis.
func (s *Store) RecordSearch(ctx context.Context, userID, query string) error {
	_, err := s.pool.Exec(ctx,
		`WITH inserted AS (
			INSERT INTO search_history (emby_user_id, query) VALUES ($1, $2)
			RETURNING id
		)
		DELETE FROM search_history
		WHERE emby_user_id = $1
		  AND occurred_at < now() - interval '30 days'`,
		userID, query)
	return err
}

// RecentSearches returns a user's distinct queries in most-recently-used order.
// Case-only duplicates collapse to the spelling used most recently.
func (s *Store) RecentSearches(
	ctx context.Context,
	userID string,
	since time.Time,
	limit int,
) ([]string, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT query
		FROM (
			SELECT DISTINCT ON (lower(query)) query, occurred_at
			FROM search_history
			WHERE emby_user_id = $1 AND occurred_at >= $2
			ORDER BY lower(query), occurred_at DESC
		) AS latest
		ORDER BY occurred_at DESC
		LIMIT $3`,
		userID, since, limit)
	if err != nil {
		return nil, fmt.Errorf("store: recent searches: %w", err)
	}
	defer rows.Close()

	queries := make([]string, 0, limit)
	for rows.Next() {
		var query string
		if err := rows.Scan(&query); err != nil {
			return nil, fmt.Errorf("store: scan recent search: %w", err)
		}
		queries = append(queries, query)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("store: read recent searches: %w", err)
	}
	return queries, nil
}

// 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,
			client_version, client_protocol
		)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
		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,
			client_version  = EXCLUDED.client_version,
			client_protocol = EXCLUDED.client_protocol,
			last_seen_at = now()`,
		sess.TokenHash, sess.EmbyUserID, sess.EmbyToken, sess.Username,
		sess.ServerID, sess.DeviceID, sess.DeviceName, sess.ClientVersion, sess.ClientProtocol)
	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, client_version, client_protocol, 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.ClientVersion,
			&sess.ClientProtocol, &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
}

// UpdateSessionClientIdentity remembers the last non-empty identity supplied by a TV.
// Headerless image requests can then still be attributed to the correct app build.
func (s *Store) UpdateSessionClientIdentity(
	ctx context.Context,
	hash []byte,
	version, protocol string,
) error {
	_, err := s.pool.Exec(ctx, `
		UPDATE sessions
		SET client_version = CASE WHEN $2 <> '' THEN $2 ELSE client_version END,
		    client_protocol = CASE WHEN $3 <> '' THEN $3 ELSE client_protocol END,
		    last_seen_at = now()
		WHERE token_hash = $1`,
		hash, version, protocol)
	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.client_version, current.client_protocol,
				current.last_seen_at
		)
		SELECT token_hash, emby_user_id, emby_token, username, server_id,
			device_id, device_name, client_version, client_protocol, 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.ClientVersion,
			&sess.ClientProtocol, &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
}