273 lines
13 KiB
Go
273 lines
13 KiB
Go
package store
|
|
|
|
import (
|
|
"bytes"
|
|
"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) {
|
|
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-client", Platform: 2, Architecture: "amd64", DaemonVersion: "test", DaemonCWD: `C:\`, SupportedShells: []byte{1},
|
|
ClientInstanceID: [16]byte{1}, SessionID: [16]byte{2}, ConnectedAt: time.Date(2026, time.September, 6, 12, 0, 0, 0, time.UTC),
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
issue, err := domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000091")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
immutable := sha256.Sum256([]byte("queue immutable request"))
|
|
now := time.Date(2026, time.September, 6, 12, 0, 1, 0, time.UTC)
|
|
input := QueueCommandInput{IssueUUID: issue, ClientID: "win-client", IssueTime: now, ReceiptTime: now, ImmutableSHA256: immutable, ExecutionSpec: []byte("serialized execution spec")}
|
|
queued, err := opened.QueueCommand(ctx, input)
|
|
if err != nil || queued.Duplicate || queued.Lifecycle != 1 || queued.Revision != 1 {
|
|
t.Fatalf("initial queue = %#v, %v", queued, err)
|
|
}
|
|
duplicate, err := opened.QueueCommand(ctx, input)
|
|
if err != nil || !duplicate.Duplicate || duplicate.Lifecycle != 1 || duplicate.Revision != 1 {
|
|
t.Fatalf("idempotent queue = %#v, %v", duplicate, err)
|
|
}
|
|
var stored []byte
|
|
var compression, rawBytes, storedBytes uint64
|
|
if err := opened.DB().QueryRow(`SELECT execution_spec, execution_spec_compression, execution_spec_raw_bytes, execution_spec_stored_bytes FROM commands WHERE issue_uuid = ?`, issue[:]).Scan(&stored, &compression, &rawBytes, &storedBytes); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
decoder, err := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
decompressed, err := decoder.DecodeAll(stored, nil)
|
|
decoder.Close()
|
|
if err != nil || !bytes.Equal(decompressed, input.ExecutionSpec) || compression != 2 || rawBytes != uint64(len(input.ExecutionSpec)) || storedBytes != uint64(len(stored)) {
|
|
t.Fatalf("stored command = raw=%q compression=%d sizes=%d/%d err=%v", decompressed, compression, rawBytes, storedBytes, err)
|
|
}
|
|
different := input
|
|
different.ImmutableSHA256 = sha256.Sum256([]byte("different request"))
|
|
if _, err := opened.QueueCommand(ctx, different); !errors.Is(err, ErrCommandConflict) {
|
|
t.Fatalf("conflicting UUID error = %v", err)
|
|
}
|
|
different = input
|
|
different.IssueUUID, err = domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000092")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
different.ClientID = "missing-client"
|
|
if _, err := opened.QueueCommand(ctx, different); !errors.Is(err, ErrClientNotFound) {
|
|
t.Fatalf("missing client error = %v", err)
|
|
}
|
|
}
|
|
|
|
func TestClaimDispatchExpiresAndFencesRequeue_HP_DISPATCH_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-client", Platform: 2, Architecture: "amd64", DaemonVersion: "test", DaemonCWD: `C:\`, SupportedShells: []byte{1},
|
|
ClientInstanceID: [16]byte{3}, SessionID: [16]byte{4}, ConnectedAt: time.Date(2026, time.September, 6, 12, 0, 0, 0, time.UTC),
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
now := time.Date(2026, time.September, 6, 12, 0, 1, 0, time.UTC)
|
|
expiredIssue, err := domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000093")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
queuedIssue, err := domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000094")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for _, input := range []QueueCommandInput{
|
|
{IssueUUID: expiredIssue, ClientID: "win-client", IssueTime: now.Add(-time.Minute), ReceiptTime: now.Add(-time.Minute), QueueExpiryTime: timePtr(now.Add(-time.Second)), ImmutableSHA256: sha256.Sum256([]byte("expired")), ExecutionSpec: []byte("expired spec")},
|
|
{IssueUUID: queuedIssue, ClientID: "win-client", IssueTime: now, ReceiptTime: now, ImmutableSHA256: sha256.Sum256([]byte("queued")), ExecutionSpec: []byte("queued spec")},
|
|
} {
|
|
if _, err := opened.QueueCommand(ctx, input); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
candidate, err := opened.ClaimNextDispatch(ctx, "win-client", 7, now)
|
|
if err != nil || candidate == nil || candidate.IssueUUID != queuedIssue || candidate.Revision != 1 || string(candidate.ExecutionSpec) != "queued spec" {
|
|
t.Fatalf("dispatch candidate = %#v, %v", candidate, err)
|
|
}
|
|
if requeued, err := opened.RequeueDispatch(ctx, queuedIssue, "win-client", 6); err != nil || requeued {
|
|
t.Fatalf("stale requeue = %t, %v", requeued, err)
|
|
}
|
|
if requeued, err := opened.RequeueDispatch(ctx, queuedIssue, "win-client", 7); err != nil || !requeued {
|
|
t.Fatalf("owned requeue = %t, %v", requeued, err)
|
|
}
|
|
candidate, err = opened.ClaimNextDispatch(ctx, "win-client", 8, now)
|
|
if err != nil || candidate == nil || candidate.IssueUUID != queuedIssue {
|
|
t.Fatalf("reclaimed candidate = %#v, %v", candidate, err)
|
|
}
|
|
var lifecycle, revision uint64
|
|
if err := opened.DB().QueryRow(`SELECT lifecycle, revision FROM commands WHERE issue_uuid = ?`, expiredIssue[:]).Scan(&lifecycle, &revision); err != nil || lifecycle != 10 || revision != 2 {
|
|
t.Fatalf("expired command lifecycle/revision = %d/%d, %v", lifecycle, revision, err)
|
|
}
|
|
}
|
|
|
|
func TestLateDispatchAcceptanceAndEventRetainExpiryContradiction_BH_DISPATCH_09(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() })
|
|
base := time.Date(2026, time.September, 6, 13, 0, 0, 0, time.UTC)
|
|
if _, err := opened.RegisterClientSession(ctx, ClientRegistration{
|
|
ClientID: "late-client", Platform: 2, Architecture: "amd64", DaemonVersion: "test", DaemonCWD: `C:\`, SupportedShells: []byte{1},
|
|
ClientInstanceID: [16]byte{21}, SessionID: [16]byte{22}, ConnectedAt: base,
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
issue := fixedStoreIssue(0xd6)
|
|
expiry := base.Add(time.Second)
|
|
if _, err := opened.QueueCommand(ctx, QueueCommandInput{IssueUUID: issue, ClientID: "late-client", IssueTime: base, ReceiptTime: base, QueueExpiryTime: &expiry, ImmutableSHA256: sha256.Sum256([]byte("late")), ExecutionSpec: []byte("spec")}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if candidate, err := opened.ClaimNextDispatch(ctx, "late-client", 7, base.Add(500*time.Millisecond)); err != nil || candidate == nil {
|
|
t.Fatalf("dispatch claim = %#v, %v", candidate, err)
|
|
}
|
|
if changed, err := opened.RecordCommandAcceptance(ctx, issue, "late-client", 7, 1, true, base.Add(2*time.Second)); err != nil || !changed {
|
|
t.Fatalf("late acceptance = %t, %v", changed, err)
|
|
}
|
|
var late int
|
|
if err := opened.DB().QueryRow(`SELECT late_after_expiry FROM commands WHERE issue_uuid = ?`, issue[:]).Scan(&late); err != nil || late != 1 {
|
|
t.Fatalf("late acceptance flag = %d, %v", late, err)
|
|
}
|
|
var audits int
|
|
if err := opened.DB().QueryRow(`SELECT COUNT(*) FROM audit_events WHERE action = 'late_after_expiry' AND issue_uuid = ?`, issue[:]).Scan(&audits); err != nil || audits != 1 {
|
|
t.Fatalf("late acceptance audit rows = %d, %v", audits, err)
|
|
}
|
|
if _, err := opened.AppendCommandEvent(ctx, EventAppend{IssueUUID: [16]byte(issue), ClientID: "late-client", SessionGeneration: 7, EventSeq: 1, ObservedUnixNano: base.Add(2 * time.Second).UnixNano(), ReceiptUnixNano: base.Add(2 * time.Second).UnixNano(), EventType: 5, Stream: 1, Compression: 1, RawLength: 1, Payload: []byte("x"), ImmutableSHA256: sha256.Sum256([]byte("event")), Output: true}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := opened.DB().QueryRow(`SELECT late_after_expiry FROM commands WHERE issue_uuid = ?`, issue[:]).Scan(&late); err != nil || late != 1 {
|
|
t.Fatalf("late event flag = %d, %v", late, err)
|
|
}
|
|
if err := opened.DB().QueryRow(`SELECT COUNT(*) FROM audit_events WHERE action = 'late_after_expiry' AND issue_uuid = ?`, issue[:]).Scan(&audits); err != nil || audits != 1 {
|
|
t.Fatalf("late event duplicated audit rows = %d, %v", audits, err)
|
|
}
|
|
}
|
|
|
|
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()
|
|
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-client", Platform: 2, Architecture: "amd64", DaemonVersion: "test", DaemonCWD: `C:\`, SupportedShells: []byte{1}, ClientInstanceID: [16]byte{5}, SessionID: [16]byte{6}, ConnectedAt: time.Now()}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
issue, err := domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000095")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := opened.QueueCommand(ctx, QueueCommandInput{IssueUUID: issue, ClientID: "win-client", IssueTime: time.Now(), ReceiptTime: time.Now(), ImmutableSHA256: sha256.Sum256([]byte("accept")), ExecutionSpec: []byte("spec")}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := opened.ClaimNextDispatch(ctx, "win-client", 3, time.Now()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := opened.RecordCommandAcceptance(ctx, issue, "win-client", 2, 1, true, time.Now()); !errors.Is(err, ErrDispatchMismatch) {
|
|
t.Fatalf("stale acceptance error = %v", err)
|
|
}
|
|
if changed, err := opened.RecordCommandAcceptance(ctx, issue, "win-client", 3, 1, true, time.Now()); err != nil || !changed {
|
|
t.Fatalf("acceptance = %t, %v", changed, err)
|
|
}
|
|
if changed, err := opened.RecordCommandAcceptance(ctx, issue, "win-client", 3, 1, true, time.Now()); err != nil || changed {
|
|
t.Fatalf("duplicate acceptance = %t, %v", changed, err)
|
|
}
|
|
rejectedIssue := fixedStoreIssue(0x96)
|
|
if _, err := opened.QueueCommand(ctx, QueueCommandInput{IssueUUID: rejectedIssue, ClientID: "win-client", IssueTime: time.Now(), ReceiptTime: time.Now(), ImmutableSHA256: sha256.Sum256([]byte("reject")), ExecutionSpec: []byte("spec")}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := opened.ClaimNextDispatch(ctx, "win-client", 3, time.Now()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
rejection := &rvboxv1.ControlError{Code: rvboxv1.ControlError_INVALID_ARGUMENT, Message: "invalid command"}
|
|
if changed, err := opened.RecordCommandAcceptanceWithRejection(ctx, rejectedIssue, "win-client", 3, 1, false, rejection, time.Now()); err != nil || !changed {
|
|
t.Fatalf("rejection acceptance = %t, %v", changed, err)
|
|
}
|
|
view, err := opened.GetCommandView(ctx, "win-client", rejectedIssue)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
stored := &rvboxv1.ControlError{}
|
|
if err := proto.Unmarshal(view.Rejection, stored); err != nil || stored.GetCode() != rejection.GetCode() || stored.GetMessage() != rejection.GetMessage() {
|
|
t.Fatalf("stored rejection = %v, %v", stored, err)
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|