244 lines
9.8 KiB
Go
244 lines
9.8 KiB
Go
package store
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/sha256"
|
|
"database/sql"
|
|
"encoding/binary"
|
|
"errors"
|
|
"math"
|
|
"time"
|
|
|
|
"github.com/klauspost/compress/zstd"
|
|
"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
|
|
}
|
|
|
|
// StdinIntent is a server-owned input record ready for delivery to a live
|
|
// client session. It is returned in command/write order and contains a copy of
|
|
// the decompressed payload; callers must not mutate storage-owned buffers.
|
|
type StdinIntent struct {
|
|
IssueUUID domain.UUID
|
|
WriteSeq uint64
|
|
Data []byte
|
|
AppendNewline bool
|
|
Close bool
|
|
}
|
|
|
|
// PendingStdin returns unacknowledged input intents for one client. Delivery is
|
|
// intentionally tracked by the session dispatcher, not by this durable query:
|
|
// a lost connection simply causes the next session to replay the same write
|
|
// sequence, which the client acknowledges idempotently.
|
|
func (store *Store) PendingStdin(ctx context.Context, clientID string) ([]StdinIntent, error) {
|
|
if clientID == "" {
|
|
return nil, errors.New("client ID is required")
|
|
}
|
|
database, err := store.openDatabase()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
rows, err := database.QueryContext(ctx, `SELECT w.issue_uuid, w.write_seq, w.payload, w.raw_bytes, w.stored_bytes, w.compression, w.sha256, w.append_newline, w.close_intent
|
|
FROM stdin_writes w JOIN commands c ON c.issue_uuid = w.issue_uuid
|
|
WHERE c.client_id = ? AND c.lifecycle BETWEEN 1 AND 4 AND w.acknowledged = 0
|
|
ORDER BY c.issue_time, c.issue_uuid, w.write_seq`, clientID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var result []StdinIntent
|
|
for rows.Next() {
|
|
var encodedIssue, stored, digest []byte
|
|
var writeSeq, rawBytes, storedBytes uint64
|
|
var compression uint32
|
|
var appendNewline, closeIntent bool
|
|
if err := rows.Scan(&encodedIssue, &writeSeq, &stored, &rawBytes, &storedBytes, &compression, &digest, &appendNewline, &closeIntent); err != nil {
|
|
return nil, err
|
|
}
|
|
if len(encodedIssue) != len(domain.UUID{}) || writeSeq == 0 || compression != 2 || storedBytes != uint64(len(stored)) || len(digest) != sha256.Size || rawBytes > maxStoredExecutionSpecBytes {
|
|
return nil, ErrInvalidSegmentRecord
|
|
}
|
|
var data []byte
|
|
data, err = decompressStdinPayload(stored, rawBytes)
|
|
if err != nil || (!closeIntent && rawBytes == 0) || (closeIntent && rawBytes != 0) {
|
|
return nil, ErrInvalidSegmentRecord
|
|
}
|
|
computed := sha256.Sum256(data)
|
|
if string(computed[:]) != string(digest) {
|
|
return nil, ErrInvalidSegmentRecord
|
|
}
|
|
var issue domain.UUID
|
|
copy(issue[:], encodedIssue)
|
|
result = append(result, StdinIntent{IssueUUID: issue, WriteSeq: writeSeq, Data: data, AppendNewline: appendNewline, Close: closeIntent})
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func decompressStdinPayload(stored []byte, expected uint64) ([]byte, error) {
|
|
if expected > maxStoredExecutionSpecBytes {
|
|
return nil, ErrInvalidSegmentRecord
|
|
}
|
|
decoder, err := zstd.NewReader(bytes.NewReader(stored), zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxMemory(maxStoredExecutionSpecBytes+1))
|
|
if err != nil {
|
|
return nil, ErrInvalidSegmentRecord
|
|
}
|
|
defer decoder.Close()
|
|
data, err := decoder.DecodeAll(stored, nil)
|
|
if err != nil || uint64(len(data)) != expected {
|
|
return nil, ErrInvalidSegmentRecord
|
|
}
|
|
return data, nil
|
|
}
|
|
|
|
// 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
|
|
}
|