rvbox/internal/server/store/stdin.go

package store

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/binary"
	"errors"
	"math"
	"time"

	"github.com/rvbox/rvbox/internal/domain"
)

var (
	ErrCommandTerminal = errors.New("command is already terminal")
	ErrStdinConflict   = errors.New("stdin request conflicts with an earlier request")
)

type StdinWriteInput struct {
	IssueUUID     domain.UUID
	ClientID      string
	RequestUUID   domain.UUID
	Data          []byte
	AppendNewline bool
	Close         bool
	ImmutableHash [32]byte
	OccurredAt    time.Time
}

type StdinWriteResult struct {
	WriteSeq  uint64
	Duplicate bool
}

// AppendStdin durably records one ordered stdin intent. It does not claim
// delivery to a process; an agent acknowledgement is a later protocol event.
// Reusing RequestUUID with the same immutable hash returns the original write
// sequence without adding a second row.
func (store *Store) AppendStdin(ctx context.Context, input StdinWriteInput) (StdinWriteResult, error) {
	return store.appendStdin(ctx, input)
}

func (store *Store) CloseStdin(ctx context.Context, input StdinWriteInput) (StdinWriteResult, error) {
	input.Close = true
	return store.appendStdin(ctx, input)
}

func (store *Store) appendStdin(ctx context.Context, input StdinWriteInput) (StdinWriteResult, error) {
	if input.IssueUUID == (domain.UUID{}) || input.ClientID == "" || input.RequestUUID == (domain.UUID{}) || input.OccurredAt.IsZero() || input.ImmutableHash == [32]byte{} || (!input.Close && len(input.Data) == 0) || (input.Close && len(input.Data) != 0) || len(input.Data) > 64<<10 {
		return StdinWriteResult{}, errors.New("invalid stdin write")
	}
	stored, err := compressCommandSpec(input.Data)
	if err != nil {
		return StdinWriteResult{}, err
	}
	charge, err := EstimateCharge(ChargeInput{EncodedBytes: uint64(len(stored)), SQLiteRows: 1, IndexEntries: 1})
	if err != nil {
		return StdinWriteResult{}, err
	}
	method := "append_stdin"
	if input.Close {
		method = "close_stdin"
	}
	target := input.IssueUUID.String()
	store.writeMu.Lock()
	defer store.writeMu.Unlock()
	database, err := store.openDatabase()
	if err != nil {
		return StdinWriteResult{}, err
	}
	if existing, found, lookupErr := lookupControlMutation(ctx, database, input.RequestUUID, method, target, input.ImmutableHash); lookupErr != nil {
		return StdinWriteResult{}, lookupErr
	} else if found {
		if len(existing) != 8 {
			return StdinWriteResult{}, ErrInvalidSegmentRecord
		}
		return StdinWriteResult{WriteSeq: binary.BigEndian.Uint64(existing), Duplicate: true}, nil
	}
	var lifecycle uint32
	var commandCharged, closeout, clientCharged, serverCharged uint64
	err = database.QueryRowContext(ctx, `SELECT commands.lifecycle, commands.charged_bytes, commands.closeout_remaining_bytes,
clients.charged_bytes, storage_counters.command_charged_bytes
FROM commands JOIN clients ON clients.client_id = commands.client_id
JOIN storage_counters ON storage_counters.singleton = 1
WHERE commands.issue_uuid = ? AND commands.client_id = ?`, input.IssueUUID[:], input.ClientID).Scan(&lifecycle, &commandCharged, &closeout, &clientCharged, &serverCharged)
	if errors.Is(err, sql.ErrNoRows) {
		return StdinWriteResult{}, ErrCommandNotFound
	}
	if err != nil {
		return StdinWriteResult{}, err
	}
	if lifecycle < 1 || lifecycle > 4 {
		return StdinWriteResult{}, ErrCommandTerminal
	}
	freeBytes, err := store.freeSpaceProbe.AvailableBytes(store.dataDir)
	if err != nil {
		return StdinWriteResult{}, err
	}
	reservation, err := CheckReservation(store.quotaLimits, ReservationState{CommandTotalCharged: commandCharged, CloseoutRemaining: closeout, ClientTotalCharged: clientCharged, ServerTotalCharged: serverCharged, FilesystemFreeBytes: freeBytes}, ReservationRequest{ChargedBytes: charge, PhysicalBytes: uint64(len(stored))})
	if err != nil {
		return StdinWriteResult{}, err
	}
	var writeSeq uint64
	if err := database.QueryRowContext(ctx, `SELECT COALESCE(MAX(write_seq), 0) + 1 FROM stdin_writes WHERE issue_uuid = ?`, input.IssueUUID[:]).Scan(&writeSeq); err != nil {
		return StdinWriteResult{}, err
	}
	if writeSeq == 0 || writeSeq > math.MaxInt64 {
		return StdinWriteResult{}, errors.New("stdin write sequence exhausted")
	}
	digest := sha256.Sum256(input.Data)
	appendNewline := 0
	if input.AppendNewline {
		appendNewline = 1
	}
	closeIntent := 0
	if input.Close {
		closeIntent = 1
	}
	tx, err := database.BeginTx(ctx, nil)
	if err != nil {
		return StdinWriteResult{}, err
	}
	defer tx.Rollback()
	_, err = tx.ExecContext(ctx, `INSERT INTO stdin_writes (issue_uuid, write_seq, payload, raw_bytes, stored_bytes, compression, sha256, append_newline, close_intent, acknowledged) VALUES (?, ?, ?, ?, ?, 2, ?, ?, ?, 0)`, input.IssueUUID[:], writeSeq, stored, len(input.Data), len(stored), digest[:], appendNewline, closeIntent)
	if err == nil {
		_, err = tx.ExecContext(ctx, `UPDATE commands SET charged_bytes = ?, closeout_remaining_bytes = ? WHERE issue_uuid = ? AND charged_bytes = ?`, reservation.CommandTotalCharged, reservation.CloseoutRemaining, input.IssueUUID[:], commandCharged)
	}
	if err == nil {
		_, err = tx.ExecContext(ctx, `UPDATE clients SET charged_bytes = ? WHERE client_id = ? AND charged_bytes = ?`, reservation.ClientTotalCharged, input.ClientID, clientCharged)
	}
	if err == nil {
		_, err = tx.ExecContext(ctx, `UPDATE storage_counters SET command_charged_bytes = ? WHERE singleton = 1 AND command_charged_bytes = ?`, reservation.ServerTotalCharged, serverCharged)
	}
	result := make([]byte, 8)
	if err == nil {
		binary.BigEndian.PutUint64(result, writeSeq)
		_, err = tx.ExecContext(ctx, `INSERT INTO control_mutations (request_uuid, method, owner_kind, owner_id, target, immutable_sha256, assigned_write_seq, result, created_at) VALUES (?, ?, 'command', ?, ?, ?, ?, ?, ?)`, input.RequestUUID[:], method, input.IssueUUID.String(), target, input.ImmutableHash[:], writeSeq, result, input.OccurredAt.UTC().UnixNano())
	}
	if err != nil {
		return StdinWriteResult{}, err
	}
	if err := tx.Commit(); err != nil {
		return StdinWriteResult{}, err
	}
	return StdinWriteResult{WriteSeq: writeSeq}, nil
}

func lookupControlMutation(ctx context.Context, database *sql.DB, requestUUID domain.UUID, method, target string, hash [32]byte) ([]byte, bool, error) {
	var storedMethod, storedTarget string
	var storedHash, result []byte
	err := database.QueryRowContext(ctx, `SELECT method, target, immutable_sha256, result FROM control_mutations WHERE request_uuid = ?`, requestUUID[:]).Scan(&storedMethod, &storedTarget, &storedHash, &result)
	if errors.Is(err, sql.ErrNoRows) {
		return nil, false, nil
	}
	if err != nil {
		return nil, false, err
	}
	if storedMethod != method || storedTarget != target || len(storedHash) != 32 || string(storedHash) != string(hash[:]) {
		return nil, false, ErrMutationConflict
	}
	return append([]byte(nil), result...), true, nil
}