feat: persist and transfer script command payloads

internal/server/store/command.go

@@ -1,7 +1,9 @@
 package store
 
 import (
+	"bytes"
 	"context"
+	"crypto/sha256"
 	"crypto/subtle"
 	"database/sql"
 	"errors"
@@ -9,7 +11,9 @@ import (
 	"time"
 
 	"github.com/klauspost/compress/zstd"
+	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
 	"github.com/rvbox/rvbox/internal/domain"
+	"google.golang.org/protobuf/proto"
 )
 
 var (
@@ -28,6 +32,12 @@ type QueueCommandInput struct {
 	QueueExpiryTime *time.Time
 	ImmutableSHA256 [32]byte
 	ExecutionSpec   []byte
+	// ScriptPresent distinguishes a valid zero-byte script from a command-text
+	// request with no script payload. When true, ScriptContent is checked
+	// against the descriptor embedded in ExecutionSpec and stored as a
+	// command-owned compressed payload.
+	ScriptPresent bool
+	ScriptContent []byte
 }
 
 type QueueCommandResult struct {
@@ -47,7 +57,18 @@ func (store *Store) QueueCommand(ctx context.Context, input QueueCommandInput) (
 	if err != nil {
 		return QueueCommandResult{}, err
 	}
-	charge, err := EstimateCharge(ChargeInput{EncodedBytes: uint64(len(stored)), SQLiteRows: 1, IndexEntries: 2})
+	scriptStored, scriptDescriptor, err := validateAndCompressScript(input)
+	if err != nil {
+		return QueueCommandResult{}, err
+	}
+	rows, indexes := uint64(1), uint64(2)
+	encodedBytes := uint64(len(stored))
+	if input.ScriptPresent {
+		rows++
+		indexes++
+		encodedBytes += uint64(len(scriptStored))
+	}
+	charge, err := EstimateCharge(ChargeInput{EncodedBytes: encodedBytes, SQLiteRows: rows, IndexEntries: indexes})
 	if err != nil {
 		return QueueCommandResult{}, err
 	}
@@ -94,7 +115,7 @@ FROM clients JOIN storage_counters ON storage_counters.singleton = 1 WHERE clien
 	decision, err := CheckReservation(store.quotaLimits, ReservationState{
 		ClientTotalCharged: clientCharged, ServerTotalCharged: serverCharged,
 		CloseoutRemaining: store.quotaLimits.CloseoutReserveBytes, FilesystemFreeBytes: freeBytes,
-	}, ReservationRequest{ChargedBytes: charge, PhysicalBytes: uint64(len(stored))})
+	}, ReservationRequest{ChargedBytes: charge, PhysicalBytes: encodedBytes})
 	if err != nil {
 		return QueueCommandResult{}, err
 	}
@@ -109,10 +130,17 @@ closeout_remaining_bytes, immutable_request_sha256, execution_spec,
 execution_spec_raw_bytes, execution_spec_stored_bytes, execution_spec_compression
 ) VALUES (?, ?, ?, ?, ?, 1, 1, ?, ?, ?, ?, ?, ?, ?, 2)`,
 		input.IssueUUID[:], input.ClientID, input.IssueTime.UTC().UnixNano(), input.ReceiptTime.UTC().UnixNano(), expiry,
-		len(stored), charge, decision.CloseoutRemaining, input.ImmutableSHA256[:], stored, len(input.ExecutionSpec), len(stored))
+		encodedBytes, charge, decision.CloseoutRemaining, input.ImmutableSHA256[:], stored, len(input.ExecutionSpec), len(stored))
 	if err != nil {
 		return QueueCommandResult{}, err
 	}
+	if input.ScriptPresent {
+		if _, err := tx.ExecContext(ctx, `INSERT INTO command_payloads (
+issue_uuid, kind, raw_bytes, stored_bytes, compression, sha256, inline_data, segment_path
+) VALUES (?, 'script', ?, ?, 2, ?, ?, NULL)`, input.IssueUUID[:], scriptDescriptor.GetSizeBytes(), len(scriptStored), scriptDescriptor.GetSha256(), scriptStored); err != nil {
+			return QueueCommandResult{}, err
+		}
+	}
 	if _, err := tx.ExecContext(ctx, `UPDATE clients SET charged_bytes = ? WHERE client_id = ?`, decision.ClientTotalCharged, input.ClientID); err != nil {
 		return QueueCommandResult{}, err
 	}
@@ -129,9 +157,37 @@ func validCommandInput(input QueueCommandInput) bool {
 	if isZeroUUID([16]byte(input.IssueUUID)) || input.ClientID == "" || len(input.ClientID) > 128 || input.IssueTime.IsZero() || input.ReceiptTime.IsZero() || len(input.ExecutionSpec) == 0 || allZero(input.ImmutableSHA256[:]) {
 		return false
 	}
+	if !input.ScriptPresent && len(input.ScriptContent) != 0 {
+		return false
+	}
 	return input.QueueExpiryTime == nil || input.QueueExpiryTime.After(input.IssueTime)
 }
 
+const maxStoredScriptBytes = 10 << 20
+
+func validateAndCompressScript(input QueueCommandInput) ([]byte, *rvboxv1.ScriptDescriptor, error) {
+	if !input.ScriptPresent {
+		return nil, nil, nil
+	}
+	var spec rvboxv1.ExecutionSpec
+	if err := proto.Unmarshal(input.ExecutionSpec, &spec); err != nil {
+		return nil, nil, fmt.Errorf("decode execution spec for script payload: %w", err)
+	}
+	descriptor := spec.GetScript()
+	if descriptor == nil || descriptor.GetSizeBytes() > maxStoredScriptBytes || len(descriptor.GetSha256()) != sha256.Size || uint64(len(input.ScriptContent)) != descriptor.GetSizeBytes() {
+		return nil, nil, errors.New("script payload does not match execution descriptor")
+	}
+	digest := sha256.Sum256(input.ScriptContent)
+	if !bytes.Equal(digest[:], descriptor.GetSha256()) {
+		return nil, nil, errors.New("script payload digest does not match execution descriptor")
+	}
+	stored, err := compressCommandSpec(input.ScriptContent)
+	if err != nil {
+		return nil, nil, err
+	}
+	return stored, descriptor, nil
+}
+
 func compressCommandSpec(spec []byte) ([]byte, error) {
 	encoder, err := zstd.NewWriter(nil, zstd.WithEncoderConcurrency(1))
 	if err != nil {

internal/server/store/command_test.go

@@ -5,12 +5,15 @@ import (
 	"context"
 	"crypto/sha256"
 	"errors"
+	"fmt"
 	"path/filepath"
 	"testing"
 	"time"
 
 	"github.com/klauspost/compress/zstd"
+	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
 	"github.com/rvbox/rvbox/internal/domain"
+	"google.golang.org/protobuf/proto"
 )
 
 func TestQueueCommandDurableIdempotency_HP_DISPATCH_01(t *testing.T) {
@@ -123,6 +126,45 @@ func TestClaimDispatchExpiresAndFencesRequeue_HP_DISPATCH_02(t *testing.T) {
 	}
 }
 
+func TestQueueAndClaimScriptPayloadIsDurable_HP_SCRIPT_02(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	opened, err := Open(ctx, Options{DataDir: filepath.Join(t.TempDir(), "state"), BusyTimeout: time.Second})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = opened.Close() })
+	if _, err := opened.RegisterClientSession(ctx, ClientRegistration{ClientID: "win-script", Platform: 2, Architecture: "amd64", DaemonVersion: "test", DaemonCWD: `C:\`, SupportedShells: []byte{1}, ClientInstanceID: [16]byte{11}, SessionID: [16]byte{12}, ConnectedAt: time.Now().UTC()}); err != nil {
+		t.Fatal(err)
+	}
+	body := []byte("Write-Output 'script payload'\r\n")
+	digest := sha256.Sum256(body)
+	spec, err := proto.Marshal(&rvboxv1.ExecutionSpec{ShellType: rvboxv1.ShellType_SHELL_POWERSHELL, Source: &rvboxv1.ExecutionSpec_Script{Script: &rvboxv1.ScriptDescriptor{Filename: "script.ps1", SizeBytes: uint64(len(body)), Sha256: digest[:]}}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	issue := fixedStoreIssue(0xb1)
+	now := time.Now().UTC()
+	if _, err := opened.QueueCommand(ctx, QueueCommandInput{IssueUUID: issue, ClientID: "win-script", IssueTime: now, ReceiptTime: now, ImmutableSHA256: sha256.Sum256([]byte("script-request")), ExecutionSpec: spec, ScriptPresent: true, ScriptContent: body}); err != nil {
+		t.Fatal(err)
+	}
+	var payload []byte
+	if err := opened.DB().QueryRowContext(ctx, `SELECT inline_data FROM command_payloads WHERE issue_uuid = ? AND kind = 'script'`, issue[:]).Scan(&payload); err != nil {
+		t.Fatal(err)
+	}
+	if bytes.Equal(payload, body) {
+		t.Fatal("script payload was stored uncompressed")
+	}
+	candidate, err := opened.ClaimNextDispatch(ctx, "win-script", 7, now.Add(time.Second))
+	if err != nil || candidate == nil || !bytes.Equal(candidate.ScriptContent, body) {
+		t.Fatalf("script dispatch candidate = %#v, %v", candidate, err)
+	}
+	pending, err := opened.PendingScriptDispatches(ctx, "win-script")
+	if err != nil || len(pending) != 1 || pending[0].IssueUUID != issue || !bytes.Equal(pending[0].Body, body) || pending[0].Digest != digest {
+		t.Fatalf("reconstructed script transfer = %#v, %v", pending, err)
+	}
+}
+
 func TestRecordCommandAcceptanceFencesGeneration_HP_DISPATCH_05(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
@@ -156,3 +198,11 @@ func TestRecordCommandAcceptanceFencesGeneration_HP_DISPATCH_05(t *testing.T) {
 }
 
 func timePtr(value time.Time) *time.Time { return &value }
+
+func fixedStoreIssue(last byte) domain.UUID {
+	value, err := domain.ParseUUIDv7(fmt.Sprintf("019c46f1-1d02-7000-8000-0000000000%02x", last))
+	if err != nil {
+		panic(err)
+	}
+	return value
+}

internal/server/store/dispatch.go

@@ -3,6 +3,7 @@ package store
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
 	"database/sql"
 	"errors"
 	"fmt"
@@ -10,7 +11,9 @@ import (
 	"time"
 
 	"github.com/klauspost/compress/zstd"
+	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
 	"github.com/rvbox/rvbox/internal/domain"
+	"google.golang.org/protobuf/proto"
 )
 
 const maxStoredExecutionSpecBytes = 1 << 20
@@ -25,6 +28,19 @@ type DispatchCandidate struct {
 	QueueExpiryTime *time.Time
 	ImmutableSHA256 [32]byte
 	ExecutionSpec   []byte
+	ScriptPresent   bool
+	ScriptContent   []byte
+}
+
+// PendingScript is the server-owned immutable script body for a command that
+// has already crossed the queued boundary. It is used to reconstruct the
+// session-local transfer window after a daemon or server reconnect; the
+// command row remains the source of truth for whether the work is still
+// non-terminal.
+type PendingScript struct {
+	IssueUUID domain.UUID
+	Body      []byte
+	Digest    [sha256.Size]byte
 }
 
 // ClaimNextDispatch expires old queued rows and atomically assigns the oldest
@@ -79,6 +95,10 @@ FROM commands WHERE client_id = ? AND lifecycle = 1 ORDER BY issue_time, issue_u
 	if err != nil {
 		return nil, err
 	}
+	script, scriptPresent, err := loadDispatchScript(ctx, tx, issue, spec)
+	if err != nil {
+		return nil, err
+	}
 	result, err := tx.ExecContext(ctx, `UPDATE commands SET lifecycle = 2, target_session_generation = ? WHERE issue_uuid = ? AND client_id = ? AND lifecycle = 1`, generation, issue[:], clientID)
 	if err != nil {
 		return nil, err
@@ -93,7 +113,7 @@ FROM commands WHERE client_id = ? AND lifecycle = 1 ORDER BY issue_time, issue_u
 	if err := tx.Commit(); err != nil {
 		return nil, err
 	}
-	candidate := &DispatchCandidate{IssueUUID: issue, Revision: revision, IssueTime: time.Unix(0, issueTime).UTC(), ImmutableSHA256: immutable, ExecutionSpec: spec}
+	candidate := &DispatchCandidate{IssueUUID: issue, Revision: revision, IssueTime: time.Unix(0, issueTime).UTC(), ImmutableSHA256: immutable, ExecutionSpec: spec, ScriptPresent: scriptPresent, ScriptContent: script}
 	if expiry.Valid {
 		value := time.Unix(0, expiry.Int64).UTC()
 		candidate.QueueExpiryTime = &value
@@ -101,6 +121,116 @@ FROM commands WHERE client_id = ? AND lifecycle = 1 ORDER BY issue_time, issue_u
 	return candidate, nil
 }
 
+func (store *Store) PendingScriptDispatches(ctx context.Context, clientID string) ([]PendingScript, 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 issue_uuid, execution_spec, execution_spec_raw_bytes
+FROM commands WHERE client_id = ? AND lifecycle BETWEEN 2 AND 4 ORDER BY issue_time, issue_uuid`, clientID)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	type pendingRow struct {
+		issue       domain.UUID
+		storedSpec  []byte
+		rawSpecSize uint64
+	}
+	var candidates []pendingRow
+	for rows.Next() {
+		var encodedIssue, stored []byte
+		var rawBytes uint64
+		if err := rows.Scan(&encodedIssue, &stored, &rawBytes); err != nil {
+			return nil, err
+		}
+		if len(encodedIssue) != 16 {
+			return nil, ErrInvalidSegmentRecord
+		}
+		var issue domain.UUID
+		copy(issue[:], encodedIssue)
+		candidates = append(candidates, pendingRow{issue: issue, storedSpec: bytes.Clone(stored), rawSpecSize: rawBytes})
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	var pending []PendingScript
+	for _, candidate := range candidates {
+		spec, err := decompressCommandSpec(candidate.storedSpec, candidate.rawSpecSize)
+		if err != nil {
+			return nil, err
+		}
+		body, present, err := loadDispatchScript(ctx, database, candidate.issue, spec)
+		if err != nil {
+			return nil, err
+		}
+		if !present {
+			continue
+		}
+		pending = append(pending, PendingScript{IssueUUID: candidate.issue, Body: body, Digest: sha256.Sum256(body)})
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return pending, nil
+}
+
+type queryRower interface {
+	QueryRowContext(context.Context, string, ...any) *sql.Row
+}
+
+func loadDispatchScript(ctx context.Context, tx queryRower, issue domain.UUID, encodedSpec []byte) ([]byte, bool, error) {
+	var spec rvboxv1.ExecutionSpec
+	if err := proto.Unmarshal(encodedSpec, &spec); err != nil {
+		// Existing store callers may use opaque test bytes. They cannot carry a
+		// script descriptor, so there is no payload to load.
+		return nil, false, nil
+	}
+	descriptor := spec.GetScript()
+	if descriptor == nil {
+		return nil, false, nil
+	}
+	var stored []byte
+	var rawBytes, storedBytes uint64
+	var compression uint32
+	var digest []byte
+	if err := tx.QueryRowContext(ctx, `SELECT inline_data, raw_bytes, stored_bytes, compression, sha256
+FROM command_payloads WHERE issue_uuid = ? AND kind = 'script'`, issue[:]).Scan(&stored, &rawBytes, &storedBytes, &compression, &digest); err != nil {
+		return nil, false, fmt.Errorf("load persisted script payload: %w", err)
+	}
+	if compression != 2 || storedBytes != uint64(len(stored)) || rawBytes != descriptor.GetSizeBytes() || len(digest) != 32 || !bytes.Equal(digest, descriptor.GetSha256()) {
+		return nil, false, ErrInvalidSegmentRecord
+	}
+	body, err := decompressStoredPayload(stored, rawBytes, maxStoredScriptBytes)
+	if err != nil {
+		return nil, false, err
+	}
+	computed := sha256.Sum256(body)
+	if !bytes.Equal(computed[:], descriptor.GetSha256()) {
+		return nil, false, ErrInvalidSegmentRecord
+	}
+	return body, true, nil
+}
+
+func decompressStoredPayload(stored []byte, rawBytes, maximum uint64) ([]byte, error) {
+	if rawBytes > maximum || rawBytes > uint64(math.MaxInt) {
+		return nil, ErrInvalidSegmentRecord
+	}
+	decoder, err := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxMemory(maximum+1))
+	if err != nil {
+		return nil, ErrInvalidSegmentRecord
+	}
+	defer decoder.Close()
+	decoded, err := decoder.DecodeAll(stored, nil)
+	if err != nil || uint64(len(decoded)) != rawBytes || uint64(len(decoded)) > maximum {
+		return nil, ErrInvalidSegmentRecord
+	}
+	return bytes.Clone(decoded), nil
+}
+
 // RequeueDispatch reverses only a known pre-acceptance write failure. A later
 // session generation, acceptance, or reconciliation decision cannot be rolled
 // back by a stale writer.

internal/server/store/payload_recovery.go

@@ -0,0 +1,114 @@
+package store
+
+import (
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"database/sql"
+
+	"github.com/klauspost/compress/zstd"
+	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
+	"google.golang.org/protobuf/proto"
+)
+
+// checkCommandPayloads verifies the inline compressed payload table and its
+// relationship to script descriptors. A valid SQLite page is not enough: a
+// short/changed zstd frame must make the owning scope dirty before dispatch.
+func checkCommandPayloads(ctx context.Context, database *sql.DB) error {
+	rows, err := database.QueryContext(ctx, `SELECT issue_uuid, kind, raw_bytes, stored_bytes, compression, sha256, inline_data, segment_path FROM command_payloads`)
+	if err != nil {
+		return err
+	}
+	defer rows.Close()
+	for rows.Next() {
+		var issue, digest, stored []byte
+		var kind string
+		var rawBytes, storedBytes uint64
+		var compression uint32
+		var segmentPath sql.NullString
+		if err := rows.Scan(&issue, &kind, &rawBytes, &storedBytes, &compression, &digest, &stored, &segmentPath); err != nil {
+			return err
+		}
+		if len(issue) != 16 || kind != "script" || compression != 2 || rawBytes > maxStoredScriptBytes || storedBytes != uint64(len(stored)) || len(digest) != sha256.Size || segmentPath.Valid || len(stored) == 0 {
+			return ErrCommittedRangeMissing
+		}
+		decoded, err := decodeCommandPayload(stored, rawBytes, maxStoredScriptBytes)
+		if err != nil {
+			return err
+		}
+		computed := sha256.Sum256(decoded)
+		if !bytes.Equal(computed[:], digest) {
+			return ErrCommittedRangeMissing
+		}
+	}
+	if err := rows.Err(); err != nil {
+		return err
+	}
+
+	commands, err := database.QueryContext(ctx, `SELECT issue_uuid, execution_spec FROM commands`)
+	if err != nil {
+		return err
+	}
+	defer commands.Close()
+	for commands.Next() {
+		var issue, encoded []byte
+		if err := commands.Scan(&issue, &encoded); err != nil {
+			return err
+		}
+		if len(issue) != 16 {
+			return ErrCommittedRangeMissing
+		}
+		var spec rvboxv1.ExecutionSpec
+		if err := proto.Unmarshal(encoded, &spec); err != nil {
+			// Some low-level store tests intentionally use opaque bytes. Such
+			// records cannot claim a script source, and dispatch validation will
+			// reject them before any user code runs.
+			continue
+		}
+		descriptor := spec.GetScript()
+		var payloadCount uint64
+		if err := database.QueryRowContext(ctx, `SELECT count(*) FROM command_payloads WHERE issue_uuid = ? AND kind = 'script'`, issue).Scan(&payloadCount); err != nil {
+			return err
+		}
+		if descriptor == nil {
+			if payloadCount != 0 {
+				return ErrCommittedRangeMissing
+			}
+			continue
+		}
+		if payloadCount != 1 || descriptor.GetSizeBytes() > maxStoredScriptBytes || len(descriptor.GetSha256()) != sha256.Size {
+			return ErrCommittedRangeMissing
+		}
+		var rawBytes, storedBytes uint64
+		var compression uint32
+		var digest, stored []byte
+		if err := database.QueryRowContext(ctx, `SELECT raw_bytes, stored_bytes, compression, sha256, inline_data FROM command_payloads WHERE issue_uuid = ? AND kind = 'script'`, issue).Scan(&rawBytes, &storedBytes, &compression, &digest, &stored); err != nil {
+			return err
+		}
+		if rawBytes != descriptor.GetSizeBytes() || compression != 2 || len(digest) != sha256.Size || !bytes.Equal(digest, descriptor.GetSha256()) {
+			return ErrCommittedRangeMissing
+		}
+		decoded, err := decodeCommandPayload(stored, rawBytes, maxStoredScriptBytes)
+		computed := sha256.Sum256(decoded)
+		if err != nil || !bytes.Equal(computed[:], descriptor.GetSha256()) {
+			return ErrCommittedRangeMissing
+		}
+	}
+	return commands.Err()
+}
+
+func decodeCommandPayload(stored []byte, expected, maximum uint64) ([]byte, error) {
+	if expected > maximum {
+		return nil, ErrCommittedRangeMissing
+	}
+	decoder, err := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxMemory(maximum+1))
+	if err != nil {
+		return nil, ErrCommittedRangeMissing
+	}
+	defer decoder.Close()
+	decoded, err := decoder.DecodeAll(stored, nil)
+	if err != nil || uint64(len(decoded)) != expected || uint64(len(decoded)) > maximum {
+		return nil, ErrCommittedRangeMissing
+	}
+	return decoded, nil
+}

internal/server/store/recovery.go

@@ -122,6 +122,9 @@ FROM output_segments ORDER BY issue_uuid, ordinal`)
 			return report, err
 		}
 	}
+	if err := checkCommandPayloads(ctx, database); err != nil {
+		return report, err
+	}
 	if err := checkQuotaCounters(ctx, database); err != nil {
 		return report, err
 	}