feat: persist and transfer script command payloads

This commit is contained in:
2026-09-06 11:11:36 +00:00
parent 26ac4c1cac
commit 3f84d3b2f1
18 changed files with 1155 additions and 31 deletions
+14 -5
View File
@@ -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)
+11 -3
View File
@@ -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()