Crussell/backend/handlers/payments/charge_helpers.go

package payments

import (
	"bytes"
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"

	"crussell/db"
	"crussell/internal/square"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// resolveChargeSource resolves the Square payment source for a card charge,
// shared by CreateBookingPayment, CreateTipPayment, BuyGiftCard, and the
// CreateTerminalPayment saved-card branch (which passes the booking user's id
// and a nil new-card token).
//
// New-card path (cnon: nonce): the nonce is used DIRECTLY for one-off charges
// (no card-on-file is created — the old tokenize-then-charge flow left orphan
// cards at Square). When saveCard is true the user's Square customer is
// provisioned FIRST and the card is tokenized against it (a ccof: source MUST
// carry its customer — R6), then saved via SaveCardForUser.
//
// Saved-card path (ccof:): a saved-card row predating P14 has an empty
// square_customer_id; the user's Square customer is lazily provisioned and
// persisted on the row BEFORE charging (a ccof: source can never be charged
// without a CustomerID).
//
// On any error the helper writes the HTTP response and returns ok=false — the
// caller must return immediately.
func resolveChargeSource(ctx context.Context, w http.ResponseWriter, svc *PaymentService, userID string, newCardToken, cardID *string, saveCard bool, notFoundMsg string) (sourceID string, savedCardID *string, squareCustomerID string, ok bool) {
	if newCardToken != nil && *newCardToken != "" {
		if saveCard {
			sqCustomerID, custErr := svc.EnsureSquareCustomer(ctx, userID)
			if custErr != nil {
				log.Printf("Failed to provision Square customer for user %s: %v", userID, custErr)
				http.Error(w, "Failed to process card", http.StatusInternalServerError)
				return "", nil, "", false
			}
			cardOnFile, err := SquareClient.CreateCardOnFile(ctx, userID, *newCardToken, sqCustomerID)
			if err != nil {
				log.Printf("Failed to create card on file: %v", err)
				http.Error(w, "Failed to process card", http.StatusInternalServerError)
				return "", nil, "", false
			}
			sourceID = cardOnFile.CardID
			squareCustomerID = sqCustomerID
			// CreateCardOnFile runs before the charge. If the subsequent payment
			// fails, this card-on-file is intentionally NOT deleted: the pending
			// record's retry re-creates it via the deterministic sha256
			// idempotency key (the SAVE path), and Square returns the same card —
			// deleting it would break that retry.
			savedRowID, saveErr := svc.SaveCardForUser(ctx, userID, sqCustomerID, cardOnFile.CardID, cardOnFile.Brand, cardOnFile.Last4, cardOnFile.ExpMonth, cardOnFile.ExpYear, cardOnFile.Fingerprint)
			if saveErr != nil {
				log.Printf("Failed to save card: %v", saveErr)
				// The Square card was JUST created by this call (CreateCardOnFile
				// above) but the local DB save failed, so the card-on-file is
				// orphaned at Square — no user_saved_cards row references it, yet
				// it is a live, chargeable card. Best-effort cleanup: disable it
				// so it cannot be charged without a DB row. This is deliberately
				// NOT the payment-failure path below — that path intentionally
				// keeps the card so the pending record's retry re-creates it via
				// the deterministic sha256 idempotency key. Here the save never
				// landed, so there is no retry to preserve. A cleanup failure must
				// never fail the charge: log the redacted card id so the orphan is
				// auditable for manual cleanup.
				if delErr := SquareClient.DeleteCardOnFile(ctx, cardOnFile.CardID); delErr != nil {
					log.Printf("WARN: created Square card %s not disabled after local save failed — orphan card-on-file requires manual cleanup: %v", square.TokenPrefix(cardOnFile.CardID), delErr)
				}
			} else {
				savedCardID = &savedRowID
			}
		} else {
			// One-off new-card charge: use the cnon: nonce DIRECTLY as the
			// source. No card-on-file is created (nothing to orphan, no
			// customer needed).
			sourceID = *newCardToken
		}
		if savedCardID == nil && saveCard {
			// sourceID is cardOnFile.CardID on this branch (the ccof card).
			log.Printf("Card was not saved despite save_card=true for user %s (card %s)", userID, square.TokenPrefix(sourceID))
		}
		return sourceID, savedCardID, squareCustomerID, true
	}
	if cardID != nil {
		card, err := svc.GetCardByID(ctx, *cardID, userID)
		if err != nil {
			if errors.Is(err, pgx.ErrNoRows) {
				http.Error(w, notFoundMsg, http.StatusNotFound)
				return "", nil, "", false
			}
			log.Printf("Failed to get card: %v", err)
			http.Error(w, "internal server error", http.StatusInternalServerError)
			return "", nil, "", false
		}
		if card.SquareCustomerID == "" {
			if userID == "" {
				// Defensive parity with the original saved-card block: a card
				// with no bookable owner cannot be provisioned. Unreachable in
				// practice — GetCardByID above filters on user_id and would
				// have 404'd for an empty owner.
				http.Error(w, "Saved card has no owner and cannot be charged", http.StatusBadRequest)
				return "", nil, "", false
			}
			provisioned, provErr := svc.EnsureSquareCustomerForSavedCard(ctx, *cardID, userID)
			if provErr != nil {
				log.Printf("Failed to provision Square customer for saved card %s (user %s): %v", *cardID, userID, provErr)
				http.Error(w, "Failed to process card", http.StatusInternalServerError)
				return "", nil, "", false
			}
			card.SquareCustomerID = provisioned
		}
		return card.SquareCardID, cardID, card.SquareCustomerID, true
	}
	// Neither a new-card token nor a saved card — validation upstream
	// (ValidateCardInfo) guarantees one of them is present.
	return "", nil, "", false
}

// acquireBookingPaymentLock acquires a pinned pool connection and a bounded
// try-lock (R6) on lockKey, serializing payment attempts per booking (the core
// defence against the two-tab double-payment race). A blocking pg_advisory_lock
// would hold the pinned pool connection for the full Square round-trip of
// whichever request holds the lock; the bounded try-lock loop gives up after
// ~3s and surfaces a 409 instead of exhausting the pool. On any failure the
// helper writes the HTTP response and returns ok=false — the caller must
// return. On success the caller MUST defer releaseBookingPaymentLock(pinConn,
// lockKey): the lock and connection stay held for the whole handler so
// pg_advisory_unlock runs on the SAME session that acquired the lock.
func acquireBookingPaymentLock(ctx context.Context, w http.ResponseWriter, lockKey, conflictMsg string) (*pgxpool.Conn, bool) {
	pinConn, err := db.Conn.Acquire(ctx)
	if err != nil {
		log.Printf("Failed to acquire connection for payment lock (%s): %v", lockKey, err)
		http.Error(w, "internal server error", http.StatusInternalServerError)
		return nil, false
	}
	lockOK, err := acquireAdvisoryLock(ctx, pinConn, lockKey)
	if err != nil {
		pinConn.Release()
		log.Printf("Failed to acquire payment serialization lock %s: %v", lockKey, err)
		http.Error(w, "internal server error", http.StatusInternalServerError)
		return nil, false
	}
	if !lockOK {
		pinConn.Release()
		log.Printf("Payment serialization lock %s not acquired within bound — a payment is already in progress", lockKey)
		http.Error(w, conflictMsg, http.StatusConflict)
		return nil, false
	}
	return pinConn, true
}

// releaseBookingPaymentLock releases the advisory lock acquired by
// acquireBookingPaymentLock and returns the pinned connection to the pool.
// Both run on the same session that holds the lock.
func releaseBookingPaymentLock(pinConn *pgxpool.Conn, lockKey string) {
	if _, err := pinConn.Exec(context.Background(), `
		SELECT pg_advisory_unlock(hashtext($1))
	`, lockKey); err != nil {
		log.Printf("Failed to release payment serialization lock %s: %v", lockKey, err)
	}
	pinConn.Release()
}

// writeChargeSnapshot stores the verbatim request JSON so the sweep can replay
// the charge with an IDENTICAL body under the same key (M1). The write is
// immutability-guarded: it records the FIRST attempt's body and stays immutable
// so a nonce-changing retry can never redirect the sweep's replay away from the
// original charge. Reuse paths that legitimately refresh the snapshot (a new
// source on a pending-reuse retry) do so via the Go-reencrypt refresh
// (refreshTillSnapshotSource / the gift-card reuse refresh), not by overwriting
// the guard. table is 'payments' or 'till_sales'; label names the flow for log
// messages (e.g. "payment", "tip payment", "gift-card payment", "till sale").
// Best-effort: failures are logged and the row stays snapshot-less — the sweep
// overrides the replay source from the live square_source_id column.
func writeChargeSnapshot(ctx context.Context, q db.Querier, table, rowID string, body any, label string) {
	snap, mErr := json.Marshal(body)
	if mErr != nil {
		log.Printf("Failed to marshal square_request_snapshot for %s %s: %v", label, rowID, mErr)
		return
	}
	stored, eErr := encryptSnapshot(snap)
	if eErr != nil {
		log.Printf("Failed to encrypt square_request_snapshot for %s %s: %v", label, rowID, eErr)
		return
	}
	if _, sErr := q.Exec(ctx, `UPDATE `+table+` SET square_request_snapshot = $1 WHERE id = $2 AND (square_request_snapshot IS NULL OR square_request_snapshot = '')`, string(stored), rowID); sErr != nil {
		log.Printf("Failed to store square_request_snapshot for %s %s: %v", label, rowID, sErr)
	}
}

// writeChargeSnapshotUnconditional stores the verbatim request JSON on a
// payments/till_sales row WITHOUT the first-attempt immutability guard that
// writeChargeSnapshot applies. It is used ONLY by the gift-card (BuyGiftCard)
// and till-sale (CreateTillSale) charge sites — the two sites that were
// UNCONDITIONAL before the writeChargeSnapshot consolidation (Loop A, finding 1
// regression check 4a). Their pending-reuse branches deliberately refresh
// square_request_snapshot in the SAME transaction as the square_source_id
// refresh (the Go-reencrypt refresh in giftcards.go / refreshTillSnapshotSource
// in till.go, B6), and this post-commit write stores the FRESH full body for
// the CURRENT attempt. The guard would wrongly skip this write on the reuse
// path when the in-transaction refresh failed best-effort — the row would then
// keep the stale first-attempt body while the live square_source_id column
// already points at the new source (the pre-consolidation code explicitly
// warned against "fixing" these sites into the guarded form). The booking/tip/
// terminal saved-card paths keep the guarded writeChargeSnapshot: their reuse
// paths do NOT refresh the snapshot, so the guard correctly records the
// immutable first-attempt body.
func writeChargeSnapshotUnconditional(ctx context.Context, q db.Querier, table, rowID string, body any, label string) {
	snap, mErr := json.Marshal(body)
	if mErr != nil {
		log.Printf("Failed to marshal square_request_snapshot for %s %s: %v", label, rowID, mErr)
		return
	}
	stored, eErr := encryptSnapshot(snap)
	if eErr != nil {
		log.Printf("Failed to encrypt square_request_snapshot for %s %s: %v", label, rowID, eErr)
		return
	}
	if _, sErr := q.Exec(ctx, `UPDATE `+table+` SET square_request_snapshot = $1 WHERE id = $2`, string(stored), rowID); sErr != nil {
		log.Printf("Failed to store square_request_snapshot for %s %s: %v", label, rowID, sErr)
	}
}

// recheckBookingPayable re-reads the booking status after a Square charge
// succeeded (R9): a concurrent cancellation/eviction can move the booking out
// of a payable state between the pre-charge status check and the charge
// completing. A charge landing on a cancelled/lapsed/no-show booking must NOT
// be recorded as a completed payment — the cancellation refund path computes
// refunds from completed payments and would silently exclude it. Returns the
// re-read status and whether a completed payment is still allowed; the caller
// owns the CRITICAL logging, the mark-failed write (whose target and
// transaction semantics differ per path), and the 409 conflict response.
func recheckBookingPayable(ctx context.Context, q db.Querier, bookingID string) (string, bool, error) {
	var status string
	// FOR UPDATE (C5): a concurrent cancellation takes the same row lock and
	// commits before this transaction commits, so the recheck cannot observe a
	// status that changes between the read and the commit. Without the lock a
	// cancellation could slip in between, leaving a completed payment on a
	// cancelled booking with no refund.
	if err := q.QueryRow(ctx, `SELECT status FROM bookings WHERE id = $1 FOR UPDATE`, bookingID).Scan(&status); err != nil {
		return "", false, err
	}
	return status, bookingStatusAllowsCompletedPayment(status), nil
}

// postChargeRecheck re-reads the booking status after a Square charge
// succeeded and, when the booking is no longer payable, marks the payment row
// failed, commits the caller's transaction, writes the 409 conflict response
// and returns false — the caller must abort. The recheck and the failed mark
// run in ONE transaction so the FOR UPDATE row lock taken inside
// recheckBookingPayable persists to commit (C5). Shared by the booking, tip
// and terminal saved-card paths so the R9 recheck cannot drift between them.
// On the not-payable branch the caller's deferred rollback becomes a harmless
// no-op (the commit already closed the transaction). chargeNoun labels the
// CRITICAL log (e.g. "payment", "tip"); conflictMsg is the 409 body.
func postChargeRecheck(ctx context.Context, w http.ResponseWriter, tx pgx.Tx, bookingID, paymentID, sqStatus, sqPayID, chargeNoun, conflictMsg string) (bool, error) {
	recheckStatus, payable, err := recheckBookingPayable(ctx, tx, bookingID)
	if err != nil {
		log.Printf("CRITICAL: Square payment %s (ID=%s) was processed but re-reading booking %s status failed: %v — manual reconciliation required", sqStatus, sqPayID, bookingID, err)
		http.Error(w, "internal server error", http.StatusInternalServerError)
		return false, err
	}
	if !payable {
		log.Printf("CRITICAL: Square payment %s (ID=%s) for booking %s was processed but booking is now %q — marking %s failed; money taken at Square MUST be refunded manually",
			sqStatus, sqPayID, bookingID, recheckStatus, chargeNoun)
		if _, upErr := tx.Exec(ctx, `UPDATE payments SET status = 'failed' WHERE id = $1`, paymentID); upErr != nil {
			log.Printf("CRITICAL: Square payment %s (ID=%s) landed on %q booking %s but marking %s failed errored: %v — manual reconciliation required",
				sqStatus, sqPayID, recheckStatus, bookingID, chargeNoun, upErr)
		}
		if cErr := tx.Commit(ctx); cErr != nil {
			log.Printf("CRITICAL: Square payment %s (ID=%s) landed on %q booking %s and committing the failed mark errored: %v — manual reconciliation required",
				sqStatus, sqPayID, recheckStatus, bookingID, cErr)
		}
		http.Error(w, conflictMsg, http.StatusConflict)
		return false, nil
	}
	return true, nil
}

// snapshotEncMarker prefixes the at-rest encrypted form of a stored
// square_request_snapshot (PII: buyer email + ccof tokens) so decryptSnapshot
// can distinguish encrypted values from plaintext (dev/mock environments and
// legacy pre-encryption rows). The marker itself is not secret.
const snapshotEncMarker = "enc:v1:"

// snapshotEncKeyWarningOnce throttles the missing-key CRITICAL log to one line
// per process: a deployment without a usable SNAPSHOT_ENC_KEY falls back to
// plaintext (money-safety first — the replayable snapshot must not be lost),
// and the single loud warning makes the misconfiguration impossible to miss.
var snapshotEncKeyWarningOnce sync.Once

// snapshotEncKey parses the AES-256-GCM key from the SNAPSHOT_ENC_KEY
// environment variable (base64-encoded 32 bytes). It is read on every call so
// tests can flip the env; the parse is cheap and encryption happens once per
// payment.
func snapshotEncKey() ([]byte, error) {
	raw := strings.TrimSpace(os.Getenv("SNAPSHOT_ENC_KEY"))
	if raw == "" {
		return nil, errors.New("SNAPSHOT_ENC_KEY is not set")
	}
	decoded, err := base64.StdEncoding.DecodeString(raw)
	if err != nil {
		return nil, fmt.Errorf("SNAPSHOT_ENC_KEY is not valid base64: %w", err)
	}
	if len(decoded) != 32 {
		return nil, fmt.Errorf("SNAPSHOT_ENC_KEY must decode to 32 bytes for AES-256, got %d", len(decoded))
	}
	return decoded, nil
}

// encryptSnapshot returns the snapshot body ready for storage. In dev/mock
// environments it returns the body unchanged (no key required, tests keep
// passing); in non-mock environments (SQUARE_ENVIRONMENT production/sandbox —
// the same gate IsExplicitDevOrMockEnv drives) it AES-256-GCM-encrypts the
// body and returns "enc:v1:" + base64(nonce || ciphertext) so the PII at rest
// (buyer email, ccof card tokens) is encrypted. The transformation is
// lossless: decryptSnapshot recovers the ORIGINAL bytes exactly, which Square's
// identical-body idempotency replay depends on. A missing/unusable key in a
// non-mock deployment falls back to plaintext with a one-time CRITICAL log —
// breaking the replayable snapshot to protect PII would strand pending rows,
// so money-safety wins over best-effort hardening.
func encryptSnapshot(body []byte) ([]byte, error) {
	if IsExplicitDevOrMockEnv() {
		return body, nil
	}
	key, err := snapshotEncKey()
	if err != nil {
		snapshotEncKeyWarningOnce.Do(func() {
			log.Printf("CRITICAL: %v — storing square_request_snapshot PLAINTEXT; set SNAPSHOT_ENC_KEY to a base64-encoded 32-byte key in non-mock deployments", err)
		})
		return body, nil
	}
	gcm, err := newSnapshotGCM(key)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("failed to read snapshot encryption nonce: %w", err)
	}
	sealed := gcm.Seal(nonce, nonce, body, nil)
	out := append([]byte(snapshotEncMarker), []byte(base64.StdEncoding.EncodeToString(sealed))...)
	return out, nil
}

// decryptSnapshot reverses encryptSnapshot for a stored square_request_snapshot.
// Marker-prefixed values are base64-decoded and AES-256-GCM-decrypted back to
// the byte-identical original request body (the sweep's by-key replay depends
// on this); values without the marker (dev/mock plaintext or legacy
// pre-encryption rows) are returned unchanged. Exporting it lets the sweep's
// stale-pending reconcile decrypt stored snapshots before replay.
func decryptSnapshot(data []byte) ([]byte, error) {
	if !bytes.HasPrefix(data, []byte(snapshotEncMarker)) {
		return data, nil
	}
	key, err := snapshotEncKey()
	if err != nil {
		return nil, fmt.Errorf("cannot decrypt stored square_request_snapshot: %w", err)
	}
	gcm, err := newSnapshotGCM(key)
	if err != nil {
		return nil, err
	}
	sealed, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(string(data), snapshotEncMarker))
	if err != nil {
		return nil, fmt.Errorf("stored square_request_snapshot is not valid base64: %w", err)
	}
	nonceSize := gcm.NonceSize()
	if len(sealed) < nonceSize {
		return nil, errors.New("stored square_request_snapshot ciphertext is too short")
	}
	nonce, ciphertext := sealed[:nonceSize], sealed[nonceSize:]
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("stored square_request_snapshot failed AES-GCM authentication: %w", err)
	}
	return plaintext, nil
}

// newSnapshotGCM builds the AES-256-GCM AEAD for the given 32-byte key.
func newSnapshotGCM(key []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("failed to init snapshot AES cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to init snapshot AES-GCM: %w", err)
	}
	return gcm, nil
}