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")

type Session struct {
	TokenHash  []byte
	EmbyUserID string
	EmbyToken  string
	Username   string
	ServerID   string
	DeviceID   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
}

func (s *Store) CreateSession(ctx context.Context, sess Session) error {
	_, err := s.pool.Exec(ctx, `
		INSERT INTO sessions (token_hash, emby_user_id, emby_token, username, server_id, device_id)
		VALUES ($1, $2, $3, $4, $5, $6)
		ON CONFLICT (token_hash) DO UPDATE SET
			emby_token   = EXCLUDED.emby_token,
			username     = EXCLUDED.username,
			server_id    = EXCLUDED.server_id,
			device_id    = EXCLUDED.device_id,
			last_seen_at = now()`,
		sess.TokenHash, sess.EmbyUserID, sess.EmbyToken, sess.Username, sess.ServerID, sess.DeviceID)
	if err != nil {
		return fmt.Errorf("store: create session: %w", err)
	}
	return 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, last_seen_at
		FROM sessions WHERE token_hash = $1`, hash).
		Scan(&sess.TokenHash, &sess.EmbyUserID, &sess.EmbyToken, &sess.Username,
			&sess.ServerID, &sess.DeviceID, &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
}