CosmicClash/server/store/assignment_sql.go

package store

import (
	"context"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strconv"
	"time"

	"github.com/cosmic-clash/cosmic-clash/server/domain"
)

// DurableAssignment is the persistence form of a verified assignment-ready
// projection. It deliberately keeps player ownership in the primary key and
// query predicate so another participant cannot recover its join material.
type DurableAssignment struct {
	MatchID           string
	PlayerID          string
	AllocationID      string
	ServerID          string
	Slot              int
	Region            string
	ClientBuild       string
	ProtocolVersion   int
	Transport         string
	Endpoint          string
	JoinAuthorisation string
	ManifestDigest    []byte
	ExpiresAt         time.Time
	Revision          uint64
}

type PostgresRosterStore struct{ DB *sql.DB }

func (s PostgresRosterStore) PublishRoster(ctx context.Context, assignment domain.Assignment, roster []domain.SignedJoinAuthorisation, verify func([]byte, []byte) bool) error {
	return SaveVerifiedAssignmentRoster(ctx, s.DB, assignment, roster, verify)
}

const AssignmentUpsertSQL = `INSERT INTO assignments
    (match_id, player_id, allocation_id, server_id, slot, region, client_build,
     protocol_version, transport, endpoint, join_authorisation, manifest_digest,
     expires_at, revision)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
ON CONFLICT (match_id, player_id) DO UPDATE SET
    allocation_id = EXCLUDED.allocation_id, server_id = EXCLUDED.server_id,
    slot = EXCLUDED.slot, region = EXCLUDED.region, client_build = EXCLUDED.client_build,
    protocol_version = EXCLUDED.protocol_version, transport = EXCLUDED.transport,
    endpoint = EXCLUDED.endpoint, join_authorisation = EXCLUDED.join_authorisation,
    manifest_digest = EXCLUDED.manifest_digest, expires_at = EXCLUDED.expires_at,
    revision = EXCLUDED.revision
WHERE assignments.allocation_id = EXCLUDED.allocation_id
  AND assignments.server_id = EXCLUDED.server_id
  AND assignments.slot = EXCLUDED.slot
  AND assignments.region = EXCLUDED.region
  AND assignments.client_build = EXCLUDED.client_build
  AND assignments.protocol_version = EXCLUDED.protocol_version
  AND assignments.transport = EXCLUDED.transport
  AND assignments.endpoint = EXCLUDED.endpoint
  AND assignments.join_authorisation = EXCLUDED.join_authorisation
  AND assignments.manifest_digest = EXCLUDED.manifest_digest
  AND assignments.expires_at = EXCLUDED.expires_at
  AND assignments.revision = EXCLUDED.revision`

const AssignmentSelectSQL = `SELECT a.match_id, a.player_id, a.allocation_id, a.server_id,
    a.slot, a.region, a.client_build, a.protocol_version, a.transport, a.endpoint,
    a.join_authorisation, a.manifest_digest, a.expires_at, a.revision
FROM assignments a
JOIN matches m ON m.match_id = a.match_id AND m.server_id = a.server_id
WHERE a.match_id = $1 AND a.player_id = $2 AND a.expires_at > $3
  AND m.state IN ('ASSIGNMENT_READY', 'ASSIGNED', 'CONNECTING', 'LIVE')`

const AssignmentRosterSelectSQL = `SELECT allocation_id, server_id, join_authorisation
FROM assignments
WHERE match_id = $1 AND server_id = $2 AND expires_at > $3
ORDER BY slot, player_id`

const AssignmentExpectedRosterSQL = `SELECT mp.player_id, i.steam_id, mp.slot, mp.team
FROM match_participants mp
JOIN identities i ON i.player_id = mp.player_id
JOIN matches m ON m.match_id = mp.match_id
JOIN allocations a ON a.allocation_id = m.allocation_id AND a.match_id = m.match_id AND a.server_id = m.server_id
WHERE mp.match_id = $1 AND m.allocation_id = $2 AND m.server_id = $3
  AND m.region = $4 AND m.protocol_version = $5
  AND a.region = $4 AND a.build = $6 AND a.protocol_version = $5 AND a.transport = $7
  AND a.state = 'ALLOCATED' AND mp.participation_active
ORDER BY mp.player_id
FOR UPDATE OF mp`

func validateDurableAssignment(assignment DurableAssignment) error {
	if assignment.MatchID == "" || assignment.PlayerID == "" || assignment.AllocationID == "" || assignment.ServerID == "" || assignment.Slot < 0 || assignment.Slot > 5 || (assignment.Region != "EU" && assignment.Region != "NA") || assignment.ClientBuild == "" || assignment.ProtocolVersion < 1 || (assignment.Transport != "enet" && assignment.Transport != "steam_sdr") || assignment.Endpoint == "" || assignment.JoinAuthorisation == "" || len(assignment.ManifestDigest) == 0 || assignment.ExpiresAt.IsZero() || assignment.Revision == 0 {
		return fmt.Errorf("invalid durable assignment")
	}
	return nil
}

func SaveAssignment(ctx context.Context, db *sql.DB, assignment DurableAssignment) error {
	if db == nil {
		return fmt.Errorf("invalid assignment database")
	}
	if err := validateDurableAssignment(assignment); err != nil {
		return err
	}
	result, err := db.ExecContext(ctx, AssignmentUpsertSQL, assignment.MatchID, assignment.PlayerID, assignment.AllocationID, assignment.ServerID, assignment.Slot, assignment.Region, assignment.ClientBuild, assignment.ProtocolVersion, assignment.Transport, assignment.Endpoint, assignment.JoinAuthorisation, assignment.ManifestDigest, assignment.ExpiresAt, assignment.Revision)
	if err != nil {
		return err
	}
	changed, err := result.RowsAffected()
	if err != nil {
		return err
	}
	if changed != 1 {
		return fmt.Errorf("assignment persistence conflict")
	}
	return nil
}

// SaveAssignments publishes a complete signed roster atomically. Assignment
// readiness is a match boundary: exposing only some players would let the
// control plane tell different participants incompatible stories after a
// transient database failure.
func SaveAssignments(ctx context.Context, db *sql.DB, assignments []DurableAssignment) error {
	if db == nil || len(assignments) == 0 {
		return fmt.Errorf("invalid assignment batch")
	}
	if err := validateAssignmentBatch(assignments); err != nil {
		return err
	}
	return RunSerializable(ctx, db, DefaultSerializableAttempts, func(ctx context.Context, tx *sql.Tx) error {
		return saveAssignmentsTx(ctx, tx, assignments)
	})
}

func validateAssignmentBatch(assignments []DurableAssignment) error {
	if len(assignments) == 0 {
		return fmt.Errorf("invalid assignment batch")
	}
	first := assignments[0]
	seen := make(map[string]struct{}, len(assignments))
	seenSlots := make(map[int]struct{}, len(assignments))
	for _, assignment := range assignments {
		if err := validateDurableAssignment(assignment); err != nil {
			return err
		}
		if assignment.MatchID != first.MatchID || assignment.AllocationID != first.AllocationID || assignment.ServerID != first.ServerID || assignment.Region != first.Region || assignment.ClientBuild != first.ClientBuild || assignment.ProtocolVersion != first.ProtocolVersion || assignment.Transport != first.Transport || assignment.Endpoint != first.Endpoint || string(assignment.ManifestDigest) != string(first.ManifestDigest) || assignment.Revision != first.Revision {
			return fmt.Errorf("mixed assignment batch")
		}
		if _, ok := seen[assignment.PlayerID]; ok {
			return fmt.Errorf("duplicate assignment in batch")
		}
		if _, ok := seenSlots[assignment.Slot]; ok {
			return fmt.Errorf("duplicate assignment slot in batch")
		}
		seen[assignment.PlayerID] = struct{}{}
		seenSlots[assignment.Slot] = struct{}{}
	}
	return nil
}

func saveAssignmentsTx(ctx context.Context, tx *sql.Tx, assignments []DurableAssignment) error {
	for _, assignment := range assignments {
		result, err := tx.ExecContext(ctx, AssignmentUpsertSQL, assignment.MatchID, assignment.PlayerID, assignment.AllocationID, assignment.ServerID, assignment.Slot, assignment.Region, assignment.ClientBuild, assignment.ProtocolVersion, assignment.Transport, assignment.Endpoint, assignment.JoinAuthorisation, assignment.ManifestDigest, assignment.ExpiresAt, assignment.Revision)
		if err != nil {
			return err
		}
		changed, err := result.RowsAffected()
		if err != nil {
			return err
		}
		if changed != 1 {
			return fmt.Errorf("assignment persistence conflict")
		}
	}
	return nil
}

// SaveVerifiedAssignmentRoster converts the backend-verified signed roster to
// player-scoped rows. It rechecks the claims at this persistence boundary so a
// caller cannot accidentally publish a token for another match or slot.
func SaveVerifiedAssignmentRoster(ctx context.Context, db *sql.DB, assignment domain.Assignment, roster []domain.SignedJoinAuthorisation, verify func([]byte, []byte) bool) error {
	if assignment.Allocation.State != domain.ServerAllocated || len(roster) == 0 || verify == nil {
		return fmt.Errorf("invalid verified assignment roster")
	}
	digest := domain.ManifestDigest(assignment.Manifest)
	rows := make([]DurableAssignment, 0, len(roster))
	seenPlayers := make(map[string]struct{}, len(roster))
	seenSlots := make(map[int]struct{}, len(roster))
	for _, signed := range roster {
		auth := signed.Authorisation
		if err := validateSignedRosterEntry(assignment, signed, verify); err != nil {
			return err
		}
		if _, exists := seenPlayers[auth.PlayerID]; exists {
			return fmt.Errorf("invalid signed assignment roster: duplicate player")
		}
		if _, exists := seenSlots[auth.Slot]; exists {
			return fmt.Errorf("invalid signed assignment roster: duplicate slot")
		}
		seenPlayers[auth.PlayerID] = struct{}{}
		seenSlots[auth.Slot] = struct{}{}
		envelope, err := json.Marshal(signed)
		if err != nil {
			return fmt.Errorf("encode signed assignment roster: %w", err)
		}
		rows = append(rows, DurableAssignment{
			MatchID: assignment.Allocation.MatchID, PlayerID: auth.PlayerID,
			AllocationID: assignment.Allocation.AllocationID, ServerID: assignment.Allocation.ServerID,
			Slot: auth.Slot, Region: assignment.Allocation.Region, ClientBuild: assignment.Allocation.Build,
			ProtocolVersion: assignment.Allocation.Protocol, Transport: assignment.Allocation.Transport,
			Endpoint: assignment.Endpoint, JoinAuthorisation: base64.RawURLEncoding.EncodeToString(envelope),
			ManifestDigest: digest[:], ExpiresAt: auth.ExpiresAt, Revision: 1,
		})
	}
	if db == nil {
		return fmt.Errorf("invalid assignment database")
	}
	if err := validateAssignmentBatch(rows); err != nil {
		return err
	}
	return RunSerializable(ctx, db, DefaultSerializableAttempts, func(ctx context.Context, tx *sql.Tx) error {
		if err := validateExpectedAssignmentRoster(ctx, tx, assignment, roster); err != nil {
			return err
		}
		return saveAssignmentsTx(ctx, tx, rows)
	})
}

func validateExpectedAssignmentRoster(ctx context.Context, tx *sql.Tx, assignment domain.Assignment, roster []domain.SignedJoinAuthorisation) error {
	rows, err := tx.QueryContext(ctx, AssignmentExpectedRosterSQL,
		assignment.Allocation.MatchID, assignment.Allocation.AllocationID, assignment.Allocation.ServerID,
		assignment.Allocation.Region, assignment.Allocation.Protocol, assignment.Allocation.Build, assignment.Allocation.Transport)
	if err != nil {
		return err
	}
	defer rows.Close()
	type expectedPlayer struct {
		steamID string
		slot    int
		team    int
	}
	expected := make(map[string]expectedPlayer, len(roster))
	for rows.Next() {
		var playerID string
		var player expectedPlayer
		if err := rows.Scan(&playerID, &player.steamID, &player.slot, &player.team); err != nil {
			return err
		}
		expected[playerID] = player
	}
	if err := rows.Err(); err != nil {
		return err
	}
	if err := rows.Close(); err != nil {
		return err
	}
	if len(expected) == 0 || len(expected) != len(roster) {
		return fmt.Errorf("signed assignment roster is incomplete")
	}
	for _, signed := range roster {
		auth := signed.Authorisation
		player, ok := expected[auth.PlayerID]
		if !ok || player.steamID != auth.SteamID || player.slot != auth.Slot || player.team != auth.Team {
			return fmt.Errorf("signed assignment roster does not match durable participants")
		}
	}
	return nil
}

func validateSignedRosterEntry(assignment domain.Assignment, signed domain.SignedJoinAuthorisation, verify func([]byte, []byte) bool) error {
	auth := signed.Authorisation
	if len(signed.Signature) == 0 || verify == nil || !verify(domain.JoinAuthorisationBytes(auth), signed.Signature) || auth.MatchID != assignment.Allocation.MatchID || auth.ServerID != assignment.Allocation.ServerID || auth.Protocol != strconv.Itoa(assignment.Allocation.Protocol) || auth.PlayerID == "" || auth.Slot < 0 || auth.Slot > 5 || auth.Team < 0 || auth.Team > 1 || auth.Slot/3 != auth.Team || auth.ExpiresAt.IsZero() {
		return fmt.Errorf("invalid signed assignment roster")
	}
	return nil
}

func GetAssignment(ctx context.Context, db *sql.DB, playerID, matchID string, now time.Time) (DurableAssignment, error) {
	if db == nil || playerID == "" || matchID == "" || now.IsZero() {
		return DurableAssignment{}, fmt.Errorf("invalid assignment recovery arguments")
	}
	var assignment DurableAssignment
	err := db.QueryRowContext(ctx, AssignmentSelectSQL, matchID, playerID, now).Scan(&assignment.MatchID, &assignment.PlayerID, &assignment.AllocationID, &assignment.ServerID, &assignment.Slot, &assignment.Region, &assignment.ClientBuild, &assignment.ProtocolVersion, &assignment.Transport, &assignment.Endpoint, &assignment.JoinAuthorisation, &assignment.ManifestDigest, &assignment.ExpiresAt, &assignment.Revision)
	if err != nil {
		return DurableAssignment{}, err
	}
	if err := validateDurableAssignment(assignment); err != nil {
		return DurableAssignment{}, err
	}
	return assignment, nil
}

// GetAssignmentRoster returns the complete signed roster for an allocated
// server. It is intentionally server-scoped rather than player-scoped and is
// called only after workload authentication at the API boundary. All rows
// must belong to one allocation; a partial or mixed allocation is unsafe to
// hand to the game process.
func GetAssignmentRoster(ctx context.Context, db *sql.DB, matchID, serverID string, now time.Time) ([][]byte, error) {
	if db == nil || matchID == "" || serverID == "" || now.IsZero() {
		return nil, fmt.Errorf("invalid assignment roster arguments")
	}
	rows, err := db.QueryContext(ctx, AssignmentRosterSelectSQL, matchID, serverID, now)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var allocationID string
	var roster [][]byte
	for rows.Next() {
		var rowAllocation, rowServer, encoded string
		if err := rows.Scan(&rowAllocation, &rowServer, &encoded); err != nil {
			return nil, err
		}
		if rowServer != serverID || rowAllocation == "" || (allocationID != "" && allocationID != rowAllocation) {
			return nil, fmt.Errorf("assignment roster contains mixed allocation")
		}
		decoded, err := base64.RawURLEncoding.DecodeString(encoded)
		if err != nil || len(decoded) == 0 {
			return nil, fmt.Errorf("assignment roster contains invalid envelope")
		}
		allocationID = rowAllocation
		roster = append(roster, decoded)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	if len(roster) == 0 {
		return nil, sql.ErrNoRows
	}
	return roster, nil
}

// LoadAssignmentParticipants reads the authoritative roster for an allocated
// match. It reuses AssignmentExpectedRosterSQL -- the same query
// SaveVerifiedAssignmentRoster re-validates against -- so the allocator cannot
// build a roster the persistence boundary would then reject for disagreeing
// with the durable participants.
func LoadAssignmentParticipants(ctx context.Context, db *sql.DB, allocation domain.Allocation) ([]domain.AssignmentParticipant, error) {
	if db == nil || allocation.MatchID == "" || allocation.AllocationID == "" || allocation.ServerID == "" {
		return nil, fmt.Errorf("invalid assignment participant arguments")
	}
	rows, err := db.QueryContext(ctx, AssignmentExpectedRosterSQL,
		allocation.MatchID, allocation.AllocationID, allocation.ServerID,
		allocation.Region, allocation.Protocol, allocation.Build, allocation.Transport)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var participants []domain.AssignmentParticipant
	for rows.Next() {
		var participant domain.AssignmentParticipant
		if err := rows.Scan(&participant.PlayerID, &participant.SteamID, &participant.Slot, &participant.Team); err != nil {
			return nil, err
		}
		if participant.PlayerID == "" || participant.SteamID == "" || participant.Slot < 0 || participant.Slot > 5 || participant.Team < 0 || participant.Team > 1 || participant.Slot/3 != participant.Team {
			return nil, fmt.Errorf("assignment participant %q has an invalid slot/team", participant.PlayerID)
		}
		participants = append(participants, participant)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	if len(participants) == 0 {
		return nil, fmt.Errorf("allocated match %s has no durable participants", allocation.MatchID)
	}
	return participants, nil
}

// AssignmentRosters adapts the participant loader to the allocator's
// AssignmentRosterSource interface.
type AssignmentRosters struct{ DB *sql.DB }

func (a AssignmentRosters) LoadAssignmentParticipants(ctx context.Context, allocation domain.Allocation) ([]domain.AssignmentParticipant, error) {
	return LoadAssignmentParticipants(ctx, a.DB, allocation)
}