feat: persist and transfer script command payloads
This commit is contained in:
@@ -4,6 +4,7 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
@@ -241,15 +242,20 @@ func (service *Service) RunCommand(ctx context.Context, request *rvboxv1.RunComm
|
||||
return nil, controlError(codes.InvalidArgument, rvboxv1.ControlError_INVALID_ARGUMENT, "target client platform is unspecified")
|
||||
}
|
||||
}
|
||||
if _, script := spec.Source.(*rvboxv1.ExecutionSpec_Script); script {
|
||||
return nil, status.Error(codes.Unimplemented, "script command admission is not enabled until payload dispatch is implemented")
|
||||
}
|
||||
if len(request.GetScriptContent()) != 0 {
|
||||
_, scriptSource := spec.Source.(*rvboxv1.ExecutionSpec_Script)
|
||||
if !scriptSource && len(request.GetScriptContent()) != 0 {
|
||||
return nil, controlError(codes.InvalidArgument, rvboxv1.ControlError_INVALID_ARGUMENT, "script_content requires a script source")
|
||||
}
|
||||
if err := agentproto.ValidateExecutionSpec(spec, service.limits, platform); err != nil {
|
||||
return nil, controlError(codes.InvalidArgument, rvboxv1.ControlError_INVALID_ARGUMENT, err.Error())
|
||||
}
|
||||
if scriptSource {
|
||||
descriptor := spec.GetScript()
|
||||
digest := sha256.Sum256(request.GetScriptContent())
|
||||
if descriptor == nil || uint64(len(request.GetScriptContent())) != descriptor.GetSizeBytes() || !bytes.Equal(digest[:], descriptor.GetSha256()) {
|
||||
return nil, controlError(codes.InvalidArgument, rvboxv1.ControlError_INVALID_ARGUMENT, "script_content does not match the script descriptor")
|
||||
}
|
||||
}
|
||||
if !advertisedShell(client.SupportedShells, spec.GetShellType()) {
|
||||
return nil, controlError(codes.InvalidArgument, rvboxv1.ControlError_INVALID_ARGUMENT, "target client does not advertise the requested shell")
|
||||
}
|
||||
@@ -262,7 +268,9 @@ func (service *Service) RunCommand(ctx context.Context, request *rvboxv1.RunComm
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
canonical := &rvboxv1.RunCommandRequest{TargetClientId: request.GetTargetClientId(), Spec: spec, QueueTtl: request.GetQueueTtl()}
|
||||
canonical := proto.Clone(request).(*rvboxv1.RunCommandRequest)
|
||||
canonical.Spec = spec
|
||||
canonical.RequestId = ""
|
||||
encoded, err := proto.MarshalOptions{Deterministic: true}.Marshal(canonical)
|
||||
if err != nil {
|
||||
return nil, controlError(codes.Internal, rvboxv1.ControlError_INTERNAL, "canonicalize command request")
|
||||
@@ -271,6 +279,7 @@ func (service *Service) RunCommand(ctx context.Context, request *rvboxv1.RunComm
|
||||
queued, err := service.store.QueueCommand(ctx, store.QueueCommandInput{
|
||||
IssueUUID: issue, ClientID: request.GetTargetClientId(), IssueTime: now, ReceiptTime: now,
|
||||
QueueExpiryTime: expiry, ImmutableSHA256: hash, ExecutionSpec: mustMarshal(spec),
|
||||
ScriptPresent: scriptSource, ScriptContent: append([]byte(nil), request.GetScriptContent()...),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, mapStoreError(err)
|
||||
|
||||
@@ -90,9 +90,17 @@ func TestRunCommandRequestIDIdempotencyAndValidation_BH_CONTROL_02(t *testing.T)
|
||||
if _, err := service.RunCommand(ctx, badID); status.Code(err) != codes.InvalidArgument {
|
||||
t.Fatalf("bad request ID code = %v", status.Code(err))
|
||||
}
|
||||
script := &rvboxv1.RunCommandRequest{TargetClientId: "win-a", RequestId: fixedIssue(0xa3).String(), Spec: &rvboxv1.ExecutionSpec{Source: &rvboxv1.ExecutionSpec_Script{Script: &rvboxv1.ScriptDescriptor{Filename: "x.ps1", SizeBytes: 1, Sha256: sha256.New().Sum(nil)}}}, ScriptContent: []byte("x")}
|
||||
if _, err := service.RunCommand(ctx, script); status.Code(err) != codes.Unimplemented {
|
||||
t.Fatalf("script admission code = %v", status.Code(err))
|
||||
scriptBody := []byte("x")
|
||||
scriptDigest := sha256.Sum256(scriptBody)
|
||||
script := &rvboxv1.RunCommandRequest{TargetClientId: "win-a", RequestId: fixedIssue(0xa3).String(), Spec: &rvboxv1.ExecutionSpec{Source: &rvboxv1.ExecutionSpec_Script{Script: &rvboxv1.ScriptDescriptor{Filename: "x.ps1", SizeBytes: 1, Sha256: scriptDigest[:]}}}, ScriptContent: scriptBody}
|
||||
if queued, err := service.RunCommand(ctx, script); err != nil || queued.GetIssueUuid() != script.GetRequestId() {
|
||||
t.Fatalf("script admission = %#v, %v", queued, err)
|
||||
}
|
||||
badScript := proto.Clone(script).(*rvboxv1.RunCommandRequest)
|
||||
badScript.RequestId = fixedIssue(0xa7).String()
|
||||
badScript.ScriptContent = []byte("y")
|
||||
if _, err := service.RunCommand(ctx, badScript); status.Code(err) != codes.InvalidArgument {
|
||||
t.Fatalf("mismatched script code = %v", status.Code(err))
|
||||
}
|
||||
badTTL := proto.Clone(request).(*rvboxv1.RunCommandRequest)
|
||||
badTTL.RequestId = fixedIssue(0xa4).String()
|
||||
|
||||
@@ -3,10 +3,12 @@ package session
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -143,13 +145,26 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
|
||||
reservations := make(map[string]DispatchLane)
|
||||
stdinSent := make(map[string]struct{})
|
||||
signalSent := make(map[string]struct{})
|
||||
scriptTransfers := make(map[string]*scriptTransfer)
|
||||
// A dispatch may have been durably claimed immediately before a server or
|
||||
// client reconnect. Reconstruct every still-active script from the command
|
||||
// store before enabling the dispatcher so an uncertain upload is replayed
|
||||
// from its immutable body instead of being silently lost.
|
||||
pendingScripts, err := server.Store.PendingScriptDispatches(parent, hello.GetClientId())
|
||||
if err != nil {
|
||||
server.close(connection, websocket.StatusInternalError, "could not restore script transfers")
|
||||
return
|
||||
}
|
||||
for _, pending := range pendingScripts {
|
||||
scriptTransfers[pending.IssueUUID.String()] = &scriptTransfer{Body: append([]byte(nil), pending.Body...), Digest: pending.Digest}
|
||||
}
|
||||
reconciled := make(chan struct{})
|
||||
var reconcileOnce sync.Once
|
||||
if !capacity.UpdateAdvertised(0, 0, hello.GetMaxRunningCommands(), hello.GetMaxQueuedCommands()) {
|
||||
server.close(connection, websocket.StatusPolicyViolation, "invalid initial client capacity")
|
||||
return
|
||||
}
|
||||
go server.dispatchLoop(sessionContext, queue, handle.DispatchWake(), reconciled, &capacityMu, capacity, reservations, stdinSent, signalSent, hello.GetClientId(), hello.GetPlatform(), encodeSessionID(sessionID), registration.Generation, cancel)
|
||||
go server.dispatchLoop(sessionContext, queue, handle.DispatchWake(), reconciled, &capacityMu, capacity, reservations, stdinSent, signalSent, scriptTransfers, hello.GetClientId(), hello.GetPlatform(), encodeSessionID(sessionID), registration.Generation, cancel)
|
||||
encodedSessionID := encodeSessionID(sessionID)
|
||||
welcome, err := proto.Marshal(&rvboxv1.AgentEnvelope{
|
||||
SessionId: encodedSessionID, SessionGeneration: registration.Generation,
|
||||
@@ -252,6 +267,26 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
|
||||
continue
|
||||
}
|
||||
if event := envelope.GetCommandEvent(); event != nil {
|
||||
var scriptIssue domain.UUID
|
||||
if scriptStatus := event.GetScriptStatus(); scriptStatus != nil {
|
||||
var parseErr error
|
||||
scriptIssue, parseErr = domain.ParseUUIDv7(event.GetIssueUuid())
|
||||
if parseErr != nil {
|
||||
server.close(connection, websocket.StatusPolicyViolation, "invalid script status")
|
||||
return
|
||||
}
|
||||
capacityMu.Lock()
|
||||
transfer := scriptTransfers[scriptIssue.String()]
|
||||
validStatus := transfer != nil && scriptStatus.GetReceivedBytes() <= uint64(len(transfer.Body)) && scriptStatus.GetReceivedBytes() <= transfer.NextOffset
|
||||
if validStatus && scriptStatus.GetComplete() {
|
||||
validStatus = scriptStatus.GetReceivedBytes() == uint64(len(transfer.Body)) && transfer.CommitQueued
|
||||
}
|
||||
capacityMu.Unlock()
|
||||
if !validStatus {
|
||||
server.close(connection, websocket.StatusPolicyViolation, "invalid script progress")
|
||||
return
|
||||
}
|
||||
}
|
||||
appendEvent, eventErr := eventAppendFromWire(event, hello.GetClientId(), registration.Generation, server.now())
|
||||
if eventErr != nil {
|
||||
server.close(connection, websocket.StatusPolicyViolation, "invalid command event")
|
||||
@@ -281,6 +316,18 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
|
||||
capacityMu.Unlock()
|
||||
handle.SignalDispatch()
|
||||
}
|
||||
if scriptStatus := event.GetScriptStatus(); scriptStatus != nil {
|
||||
capacityMu.Lock()
|
||||
transfer := scriptTransfers[scriptIssue.String()]
|
||||
if scriptStatus.GetReceivedBytes() > transfer.AcknowledgedOffset {
|
||||
transfer.AcknowledgedOffset = scriptStatus.GetReceivedBytes()
|
||||
}
|
||||
if scriptStatus.GetComplete() {
|
||||
transfer.Completed = true
|
||||
}
|
||||
capacityMu.Unlock()
|
||||
handle.SignalDispatch()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -337,7 +384,7 @@ func eventType(event *rvboxv1.CommandEvent) uint16 {
|
||||
// enqueueNextDispatch records the queued-to-dispatched transition before
|
||||
// exposing work to the network. A full data lane is a pre-write failure, so
|
||||
// only the owning generation can put the command back into the queue.
|
||||
func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *WriterQueue, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64, beforeEnqueue func(domain.UUID)) (domain.UUID, bool, error) {
|
||||
func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *WriterQueue, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64, beforeEnqueue func(*store.DispatchCandidate)) (domain.UUID, bool, error) {
|
||||
candidate, err := server.Store.ClaimNextDispatch(ctx, clientID, generation, server.now())
|
||||
if err != nil || candidate == nil {
|
||||
return domain.UUID{}, false, err
|
||||
@@ -373,7 +420,7 @@ func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *Write
|
||||
return candidate.IssueUUID, sent, requeueErr
|
||||
}
|
||||
if beforeEnqueue != nil {
|
||||
beforeEnqueue(candidate.IssueUUID)
|
||||
beforeEnqueue(candidate)
|
||||
}
|
||||
if !queue.EnqueueData(Frame{Kind: FrameData, Payload: encoded}) {
|
||||
sent, requeueErr := requeue(ErrDispatchDataFull)
|
||||
@@ -382,6 +429,75 @@ func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *Write
|
||||
return candidate.IssueUUID, true, nil
|
||||
}
|
||||
|
||||
const scriptSendWindowBytes uint64 = 1 << 20
|
||||
|
||||
type scriptTransfer struct {
|
||||
Body []byte
|
||||
Digest [sha256.Size]byte
|
||||
NextOffset uint64
|
||||
AcknowledgedOffset uint64
|
||||
CommitQueued bool
|
||||
Completed bool
|
||||
}
|
||||
|
||||
// enqueueNextScript sends one bounded script frame. The dispatcher never
|
||||
// queues more than a 1 MiB unacknowledged window, and the data lane remains
|
||||
// bounded; a reconnect simply starts the exact payload from offset zero.
|
||||
func (server *AgentServer) enqueueNextScript(queue *WriterQueue, sessionID string, generation uint64, transfers map[string]*scriptTransfer, sentMu *sync.Mutex, chunkLimit uint64) (bool, error) {
|
||||
if chunkLimit == 0 {
|
||||
return false, errors.New("script chunk limit is zero")
|
||||
}
|
||||
sentMu.Lock()
|
||||
keys := make([]string, 0, len(transfers))
|
||||
for key := range transfers {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
for _, key := range keys {
|
||||
transfer := transfers[key]
|
||||
if transfer == nil || transfer.Completed || transfer.NextOffset-transfer.AcknowledgedOffset >= scriptSendWindowBytes {
|
||||
continue
|
||||
}
|
||||
var envelope *rvboxv1.AgentEnvelope
|
||||
if transfer.NextOffset < uint64(len(transfer.Body)) {
|
||||
end := transfer.NextOffset + chunkLimit
|
||||
if end > uint64(len(transfer.Body)) {
|
||||
end = uint64(len(transfer.Body))
|
||||
}
|
||||
chunk := transfer.Body[transfer.NextOffset:end]
|
||||
digest := sha256.Sum256(chunk)
|
||||
envelope = &rvboxv1.AgentEnvelope{SessionId: sessionID, SessionGeneration: generation, Payload: &rvboxv1.AgentEnvelope_ScriptChunk{ScriptChunk: &rvboxv1.ScriptChunk{IssueUuid: key, Offset: transfer.NextOffset, Data: append([]byte(nil), chunk...), Sha256: digest[:]}}}
|
||||
} else if !transfer.CommitQueued {
|
||||
transfer.CommitQueued = true
|
||||
envelope = &rvboxv1.AgentEnvelope{SessionId: sessionID, SessionGeneration: generation, Payload: &rvboxv1.AgentEnvelope_ScriptCommit{ScriptCommit: &rvboxv1.ScriptCommit{IssueUuid: key, SizeBytes: uint64(len(transfer.Body)), Sha256: transfer.Digest[:]}}}
|
||||
} else {
|
||||
continue
|
||||
}
|
||||
encoded, err := proto.Marshal(envelope)
|
||||
if err != nil {
|
||||
if transfer.CommitQueued && transfer.NextOffset == uint64(len(transfer.Body)) {
|
||||
transfer.CommitQueued = false
|
||||
}
|
||||
sentMu.Unlock()
|
||||
return false, err
|
||||
}
|
||||
if !queue.EnqueueData(Frame{Kind: FrameData, Payload: encoded}) {
|
||||
if transfer.CommitQueued && transfer.NextOffset == uint64(len(transfer.Body)) {
|
||||
transfer.CommitQueued = false
|
||||
}
|
||||
sentMu.Unlock()
|
||||
return false, ErrDispatchDataFull
|
||||
}
|
||||
if transfer.NextOffset < uint64(len(transfer.Body)) {
|
||||
transfer.NextOffset += uint64(len(envelope.GetScriptChunk().GetData()))
|
||||
}
|
||||
sentMu.Unlock()
|
||||
return true, nil
|
||||
}
|
||||
sentMu.Unlock()
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// enqueueNextStdin exposes one durable input intent on the essential control
|
||||
// lane. Session-local sent tracking suppresses duplicate frames while a live
|
||||
// connection remains usable; reconnecting naturally replays unacknowledged
|
||||
@@ -472,7 +588,7 @@ func (server *AgentServer) enqueueNextSignal(ctx context.Context, queue *WriterQ
|
||||
// dispatchLoop is the per-session serialized dispatcher. It waits for a
|
||||
// complete reconciliation result before consuming queued work, then coalesces
|
||||
// wakeups from local control RPCs, capacity advertisements, and acceptances.
|
||||
func (server *AgentServer) dispatchLoop(ctx context.Context, queue *WriterQueue, wake <-chan struct{}, reconciled <-chan struct{}, capacityMu *sync.Mutex, capacity *CapacityShadow, reservations map[string]DispatchLane, stdinSent, signalSent map[string]struct{}, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64, cancel context.CancelFunc) {
|
||||
func (server *AgentServer) dispatchLoop(ctx context.Context, queue *WriterQueue, wake <-chan struct{}, reconciled <-chan struct{}, capacityMu *sync.Mutex, capacity *CapacityShadow, reservations map[string]DispatchLane, stdinSent, signalSent map[string]struct{}, scriptTransfers map[string]*scriptTransfer, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64, cancel context.CancelFunc) {
|
||||
select {
|
||||
case <-reconciled:
|
||||
case <-ctx.Done():
|
||||
@@ -496,6 +612,14 @@ func (server *AgentServer) dispatchLoop(ctx context.Context, queue *WriterQueue,
|
||||
if signalQueued {
|
||||
continue
|
||||
}
|
||||
scriptQueued, scriptErr := server.enqueueNextScript(queue, sessionID, generation, scriptTransfers, capacityMu, server.limits().MaxRawChunkBytes)
|
||||
if scriptErr != nil && !errors.Is(scriptErr, ErrDispatchDataFull) {
|
||||
server.closeForDispatchFailure(cancel)
|
||||
return
|
||||
}
|
||||
if scriptQueued {
|
||||
continue
|
||||
}
|
||||
capacityMu.Lock()
|
||||
lane := capacity.Reserve()
|
||||
capacityMu.Unlock()
|
||||
@@ -503,9 +627,12 @@ func (server *AgentServer) dispatchLoop(ctx context.Context, queue *WriterQueue,
|
||||
break
|
||||
}
|
||||
inserted := false
|
||||
issue, sent, err := server.enqueueNextDispatch(ctx, queue, clientID, platform, sessionID, generation, func(issue domain.UUID) {
|
||||
issue, sent, err := server.enqueueNextDispatch(ctx, queue, clientID, platform, sessionID, generation, func(candidate *store.DispatchCandidate) {
|
||||
capacityMu.Lock()
|
||||
reservations[issue.String()] = lane
|
||||
reservations[candidate.IssueUUID.String()] = lane
|
||||
if candidate.ScriptPresent {
|
||||
scriptTransfers[candidate.IssueUUID.String()] = &scriptTransfer{Body: append([]byte(nil), candidate.ScriptContent...), Digest: sha256.Sum256(candidate.ScriptContent)}
|
||||
}
|
||||
inserted = true
|
||||
capacityMu.Unlock()
|
||||
})
|
||||
@@ -513,6 +640,7 @@ func (server *AgentServer) dispatchLoop(ctx context.Context, queue *WriterQueue,
|
||||
capacityMu.Lock()
|
||||
if inserted {
|
||||
delete(reservations, issue.String())
|
||||
delete(scriptTransfers, issue.String())
|
||||
}
|
||||
capacity.Release(lane)
|
||||
capacityMu.Unlock()
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -20,6 +23,54 @@ import (
|
||||
"google.golang.org/protobuf/types/known/timestamppb"
|
||||
)
|
||||
|
||||
func TestScriptTransferUsesChecksummedChunksAndCommit_HP_SCRIPT_03(t *testing.T) {
|
||||
t.Parallel()
|
||||
body := bytes.Repeat([]byte("abcd"), 8)
|
||||
digest := sha256.Sum256(body)
|
||||
transfer := map[string]*scriptTransfer{"019c46f1-1d02-7000-8000-0000000000b1": {Body: body, Digest: digest}}
|
||||
queue := NewWriterQueue(4, 8, 2)
|
||||
defer queue.Close()
|
||||
for offset := uint64(0); offset < uint64(len(body)); {
|
||||
queued, err := (&AgentServer{}).enqueueNextScript(queue, "session", 7, transfer, &sync.Mutex{}, 5)
|
||||
if err != nil || !queued {
|
||||
t.Fatalf("chunk at offset %d: queued=%t err=%v", offset, queued, err)
|
||||
}
|
||||
frame, err := queue.Next(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var envelope rvboxv1.AgentEnvelope
|
||||
if err := proto.Unmarshal(frame.Payload, &envelope); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
chunk := envelope.GetScriptChunk()
|
||||
if chunk == nil || chunk.GetOffset() != offset || !bytes.Equal(chunk.GetData(), body[offset:offset+uint64(len(chunk.GetData()))]) {
|
||||
t.Fatalf("chunk = %+v at offset %d", chunk, offset)
|
||||
}
|
||||
chunkDigest := sha256.Sum256(chunk.GetData())
|
||||
if !bytes.Equal(chunk.GetSha256(), chunkDigest[:]) {
|
||||
t.Fatal("chunk digest mismatch")
|
||||
}
|
||||
offset += uint64(len(chunk.GetData()))
|
||||
}
|
||||
queued, err := (&AgentServer{}).enqueueNextScript(queue, "session", 7, transfer, &sync.Mutex{}, 5)
|
||||
if err != nil || !queued {
|
||||
t.Fatalf("commit queued=%t err=%v", queued, err)
|
||||
}
|
||||
frame, err := queue.Next(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var envelope rvboxv1.AgentEnvelope
|
||||
if err := proto.Unmarshal(frame.Payload, &envelope); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
commit := envelope.GetScriptCommit()
|
||||
if commit == nil || commit.GetSizeBytes() != uint64(len(body)) || !bytes.Equal(commit.GetSha256(), digest[:]) {
|
||||
t.Fatalf("commit = %+v", commit)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentServerRegistrationAndReplacement_HP_SES_05(t *testing.T) {
|
||||
server, cleanup := newTestAgentServer(t)
|
||||
defer cleanup()
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user