CosmicClash/server/store/server_connection_sql.go

package store

import (
	"bytes"
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"time"

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

const ServerConnectionIdempotencyScope = "server.connection"

const ServerConnectionIdempotencyInsertSQL = `INSERT INTO idempotency_keys
    (scope, idempotency_key, payload_digest, result)
VALUES ($1, $2, $3, '{}'::jsonb) ON CONFLICT (scope, idempotency_key) DO NOTHING`

const ServerConnectionIdempotencySelectSQL = `SELECT payload_digest, result
FROM idempotency_keys WHERE scope = $1 AND idempotency_key = $2 FOR UPDATE`

const ServerConnectionLeaseSQL = `SELECT mp.connection_generation, mp.connected_at, mp.disconnected_at, assn.expires_at
FROM match_participants mp
JOIN matches m ON m.match_id = mp.match_id
JOIN allocations a ON a.allocation_id = $3 AND a.match_id = m.match_id AND a.server_id = m.server_id
JOIN assignments assn ON assn.match_id = mp.match_id AND assn.player_id = mp.player_id
    AND assn.allocation_id = a.allocation_id AND assn.server_id = m.server_id
WHERE mp.match_id = $1 AND mp.player_id = $4 AND mp.participation_active
  AND m.server_id = $2 AND m.state IN ('ASSIGNMENT_READY', 'ASSIGNED', 'CONNECTING', 'LIVE')
  AND a.state = 'ALLOCATED'
FOR UPDATE OF mp`

const ServerConnectionAdmitSQL = `UPDATE match_participants
SET connection_generation = $3, connected_at = COALESCE(connected_at, $4), disconnected_at = NULL
WHERE match_id = $1 AND player_id = $2 AND connection_generation = $5
RETURNING connection_generation`

const ServerConnectionDisconnectSQL = `UPDATE match_participants
SET disconnected_at = $4
WHERE match_id = $1 AND player_id = $2 AND connection_generation = $3
  AND connected_at IS NOT NULL AND disconnected_at IS NULL
RETURNING connection_generation`

type connectionReceipt struct {
	Generation uint64 `json:"generation"`
}

// ClaimPlayerConnection atomically acquires the next durable connection
// generation. expectedGeneration is server-owned state, never a client claim.
// A reconnect is legal only after the exact previous generation was durably
// disconnected and while its 60-second grace period remains open.
func ClaimPlayerConnection(ctx context.Context, db *sql.DB, binding domain.WorkloadBinding, playerID string, expectedGeneration uint64, idempotencyKey string, now time.Time) (uint64, error) {
	if err := validateConnectionMutation(db, binding, playerID, idempotencyKey, now); err != nil {
		return 0, err
	}
	digest := connectionDigest("connect", binding, playerID, expectedGeneration)
	var claimed uint64
	err := RunSerializable(ctx, db, DefaultSerializableAttempts, func(ctx context.Context, tx *sql.Tx) error {
		replay, generation, err := beginConnectionMutation(ctx, tx, idempotencyKey, digest)
		if err != nil {
			return err
		}
		if replay {
			current, connectedAt, disconnectedAt, _, err := lockConnectionLease(ctx, tx, binding, playerID)
			if err != nil {
				return err
			}
			if current != generation || !connectedAt.Valid || disconnectedAt.Valid {
				return domain.ErrConflict
			}
			claimed = generation
			return nil
		}
		current, connectedAt, disconnectedAt, assignmentExpiry, err := lockConnectionLease(ctx, tx, binding, playerID)
		if err != nil {
			return err
		}
		// A fresh process has no in-memory generation. It may recover only a
		// durably disconnected lease; an active row still fences it. All
		// nonzero expectations remain exact CAS operations.
		if (current != expectedGeneration && !(expectedGeneration == 0 && current > 0 && disconnectedAt.Valid)) || expectedGeneration == ^uint64(0) {
			return domain.ErrConflict
		}
		if current == 0 {
			if connectedAt.Valid || disconnectedAt.Valid || !now.Before(assignmentExpiry) {
				return domain.ErrConflict
			}
		} else if !connectedAt.Valid || !disconnectedAt.Valid || now.Before(disconnectedAt.Time) || now.Sub(disconnectedAt.Time) > domain.RankedReconnectGrace {
			return domain.ErrConflict
		}
		claimed = current + 1
		if err := tx.QueryRowContext(ctx, ServerConnectionAdmitSQL, binding.MatchID, playerID, claimed, now, current).Scan(&claimed); err != nil {
			if err == sql.ErrNoRows {
				return domain.ErrConflict
			}
			return err
		}
		return finishConnectionMutation(ctx, tx, idempotencyKey, claimed)
	})
	return claimed, err
}

// RecordPlayerDisconnected closes exactly one active generation. A delayed
// disconnect from an older peer can therefore never evict a reclaimed lease.
func RecordPlayerDisconnected(ctx context.Context, db *sql.DB, binding domain.WorkloadBinding, playerID string, generation uint64, idempotencyKey string, now time.Time) error {
	if err := validateConnectionMutation(db, binding, playerID, idempotencyKey, now); err != nil {
		return err
	}
	if generation == 0 {
		return fmt.Errorf("invalid server disconnect receipt")
	}
	digest := connectionDigest("disconnect", binding, playerID, generation)
	return RunSerializable(ctx, db, DefaultSerializableAttempts, func(ctx context.Context, tx *sql.Tx) error {
		replay, _, err := beginConnectionMutation(ctx, tx, idempotencyKey, digest)
		if err != nil || replay {
			return err
		}
		current, connectedAt, disconnectedAt, _, err := lockConnectionLease(ctx, tx, binding, playerID)
		if err != nil {
			return err
		}
		if current != generation || !connectedAt.Valid || disconnectedAt.Valid || now.Before(connectedAt.Time) {
			return domain.ErrConflict
		}
		var recorded uint64
		if err := tx.QueryRowContext(ctx, ServerConnectionDisconnectSQL, binding.MatchID, playerID, generation, now).Scan(&recorded); err != nil {
			if err == sql.ErrNoRows {
				return domain.ErrConflict
			}
			return err
		}
		return finishConnectionMutation(ctx, tx, idempotencyKey, recorded)
	})
}

func validateConnectionMutation(db *sql.DB, binding domain.WorkloadBinding, playerID, key string, now time.Time) error {
	if db == nil || binding.AllocationID == "" || binding.MatchID == "" || binding.ServerID == "" || playerID == "" || len(key) < 16 || len(key) > 128 || now.IsZero() {
		return fmt.Errorf("invalid server connection receipt")
	}
	return nil
}

func connectionDigest(operation string, binding domain.WorkloadBinding, playerID string, generation uint64) [sha256.Size]byte {
	payload := []byte(operation + "\x00" + binding.AllocationID + "\x00" + binding.MatchID + "\x00" + binding.ServerID + "\x00" + playerID + "\x00")
	encoded := make([]byte, 8)
	binary.BigEndian.PutUint64(encoded, generation)
	return sha256.Sum256(append(payload, encoded...))
}

func beginConnectionMutation(ctx context.Context, tx *sql.Tx, key string, digest [sha256.Size]byte) (bool, uint64, error) {
	inserted, err := tx.ExecContext(ctx, ServerConnectionIdempotencyInsertSQL, ServerConnectionIdempotencyScope, key, digest[:])
	if err != nil {
		return false, 0, err
	}
	changed, err := inserted.RowsAffected()
	if err != nil || changed != 0 {
		return false, 0, err
	}
	var prior, result []byte
	if err := tx.QueryRowContext(ctx, ServerConnectionIdempotencySelectSQL, ServerConnectionIdempotencyScope, key).Scan(&prior, &result); err != nil {
		return false, 0, err
	}
	if !bytes.Equal(prior, digest[:]) {
		return false, 0, domain.ErrConflict
	}
	var receipt connectionReceipt
	if err := json.Unmarshal(result, &receipt); err != nil || receipt.Generation == 0 {
		return false, 0, domain.ErrConflict
	}
	return true, receipt.Generation, nil
}

func lockConnectionLease(ctx context.Context, tx *sql.Tx, binding domain.WorkloadBinding, playerID string) (uint64, sql.NullTime, sql.NullTime, time.Time, error) {
	var generation uint64
	var connectedAt, disconnectedAt sql.NullTime
	var assignmentExpiry time.Time
	err := tx.QueryRowContext(ctx, ServerConnectionLeaseSQL, binding.MatchID, binding.ServerID, binding.AllocationID, playerID).Scan(&generation, &connectedAt, &disconnectedAt, &assignmentExpiry)
	if err == sql.ErrNoRows {
		err = domain.ErrConflict
	}
	return generation, connectedAt, disconnectedAt, assignmentExpiry, err
}

func finishConnectionMutation(ctx context.Context, tx *sql.Tx, key string, generation uint64) error {
	result, err := json.Marshal(connectionReceipt{Generation: generation})
	if err != nil {
		return err
	}
	_, err = tx.ExecContext(ctx, `UPDATE idempotency_keys SET result = $3 WHERE scope = $1 AND idempotency_key = $2`, ServerConnectionIdempotencyScope, key, result)
	return err
}