feat: add Unix control service and rvc queue path
This commit is contained in:
@@ -0,0 +1,181 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"net"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
|
||||
"github.com/rvbox/rvbox/internal/agentproto"
|
||||
"github.com/rvbox/rvbox/internal/domain"
|
||||
"github.com/rvbox/rvbox/internal/server/store"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/credentials/insecure"
|
||||
"google.golang.org/grpc/status"
|
||||
"google.golang.org/grpc/test/bufconn"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/types/known/durationpb"
|
||||
)
|
||||
|
||||
func TestControlListAndGetViews_HP_CONTROL_01(t *testing.T) {
|
||||
service, persistence := newTestService(t)
|
||||
defer persistence.Close()
|
||||
ctx := context.Background()
|
||||
registerControlClient(t, persistence, "win-a", rvboxv1.Platform_PLATFORM_WINDOWS, rvboxv1.ShellType_SHELL_POWERSHELL, 1)
|
||||
registerControlClient(t, persistence, "win-b", rvboxv1.Platform_PLATFORM_WINDOWS, rvboxv1.ShellType_SHELL_CMD, 2)
|
||||
|
||||
first, err := service.ListClients(ctx, &rvboxv1.ListClientsRequest{PageSize: 1})
|
||||
if err != nil || len(first.GetClients()) != 1 || first.GetNextPageToken() == "" {
|
||||
t.Fatalf("first client page = %#v, %v", first, err)
|
||||
}
|
||||
second, err := service.ListClients(ctx, &rvboxv1.ListClientsRequest{PageSize: 1, PageToken: first.GetNextPageToken()})
|
||||
if err != nil || len(second.GetClients()) != 1 || second.GetClients()[0].GetClientId() != "win-b" {
|
||||
t.Fatalf("second client page = %#v, %v", second, err)
|
||||
}
|
||||
if _, err := service.ListClients(ctx, &rvboxv1.ListClientsRequest{PageSize: 1, PageToken: "bad-token"}); status.Code(err) != codes.InvalidArgument {
|
||||
t.Fatalf("bad client cursor code = %v", status.Code(err))
|
||||
}
|
||||
client, err := service.GetClient(ctx, &rvboxv1.GetClientRequest{ClientId: "win-a"})
|
||||
if err != nil || client.GetClient().GetSupportedShells()[0] != rvboxv1.ShellType_SHELL_POWERSHELL {
|
||||
t.Fatalf("get client = %#v, %v", client, err)
|
||||
}
|
||||
|
||||
issue := fixedIssue(0xa1)
|
||||
queued, err := service.RunCommand(ctx, &rvboxv1.RunCommandRequest{
|
||||
TargetClientId: "win-a", RequestId: issue.String(),
|
||||
Spec: &rvboxv1.ExecutionSpec{Source: &rvboxv1.ExecutionSpec_CommandText{CommandText: "echo hi"}},
|
||||
})
|
||||
if err != nil || queued.GetIssueUuid() != issue.String() || queued.GetLifecycle() != rvboxv1.CommandLifecycle_COMMAND_QUEUED {
|
||||
t.Fatalf("run response = %#v, %v", queued, err)
|
||||
}
|
||||
commands, err := service.ListCommands(ctx, &rvboxv1.ListCommandsRequest{ClientId: "win-a", IncludeTerminal: true})
|
||||
if err != nil || len(commands.GetCommands()) != 1 || commands.GetCommands()[0].GetSpec().GetShellType() != rvboxv1.ShellType_SHELL_POWERSHELL {
|
||||
t.Fatalf("command list = %#v, %v", commands, err)
|
||||
}
|
||||
got, err := service.GetCommand(ctx, &rvboxv1.GetCommandRequest{ClientId: "win-a", IssueUuid: issue.String()})
|
||||
if err != nil || got.GetCommand().GetIssueUuid() != issue.String() {
|
||||
t.Fatalf("command get = %#v, %v", got, err)
|
||||
}
|
||||
if _, err := service.GetCommand(ctx, &rvboxv1.GetCommandRequest{IssueUuid: "019c46f1-1d02-6000-8000-0000000000a1"}); status.Code(err) != codes.InvalidArgument {
|
||||
t.Fatalf("non-v7 command ID code = %v", status.Code(err))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunCommandRequestIDIdempotencyAndValidation_BH_CONTROL_02(t *testing.T) {
|
||||
service, persistence := newTestService(t)
|
||||
defer persistence.Close()
|
||||
ctx := context.Background()
|
||||
registerControlClient(t, persistence, "win-a", rvboxv1.Platform_PLATFORM_WINDOWS, rvboxv1.ShellType_SHELL_POWERSHELL, 3)
|
||||
issue := fixedIssue(0xa2)
|
||||
request := &rvboxv1.RunCommandRequest{TargetClientId: "win-a", RequestId: issue.String(), Spec: &rvboxv1.ExecutionSpec{Source: &rvboxv1.ExecutionSpec_CommandText{CommandText: "echo one"}}}
|
||||
if _, err := service.RunCommand(ctx, request); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if replay, err := service.RunCommand(ctx, proto.Clone(request).(*rvboxv1.RunCommandRequest)); err != nil || replay.GetIssueUuid() != issue.String() {
|
||||
t.Fatalf("exact replay = %#v, %v", replay, err)
|
||||
}
|
||||
conflict := proto.Clone(request).(*rvboxv1.RunCommandRequest)
|
||||
conflict.Spec = &rvboxv1.ExecutionSpec{Source: &rvboxv1.ExecutionSpec_CommandText{CommandText: "echo two"}}
|
||||
if _, err := service.RunCommand(ctx, conflict); status.Code(err) != codes.AlreadyExists {
|
||||
t.Fatalf("conflicting replay code = %v", status.Code(err))
|
||||
}
|
||||
badID := proto.Clone(request).(*rvboxv1.RunCommandRequest)
|
||||
badID.RequestId = "not-a-uuid"
|
||||
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))
|
||||
}
|
||||
badTTL := proto.Clone(request).(*rvboxv1.RunCommandRequest)
|
||||
badTTL.RequestId = fixedIssue(0xa4).String()
|
||||
badTTL.QueueTtl = durationpb.New(-time.Second)
|
||||
if _, err := service.RunCommand(ctx, badTTL); status.Code(err) != codes.InvalidArgument {
|
||||
t.Fatalf("negative TTL code = %v", status.Code(err))
|
||||
}
|
||||
}
|
||||
|
||||
func TestCommandPaginationBindsFilters_BH_CONTROL_03(t *testing.T) {
|
||||
service, persistence := newTestService(t)
|
||||
defer persistence.Close()
|
||||
ctx := context.Background()
|
||||
registerControlClient(t, persistence, "win-a", rvboxv1.Platform_PLATFORM_WINDOWS, rvboxv1.ShellType_SHELL_CMD, 4)
|
||||
for index := byte(0xb0); index < 0xb3; index++ {
|
||||
if _, err := service.RunCommand(ctx, &rvboxv1.RunCommandRequest{TargetClientId: "win-a", RequestId: fixedIssue(index).String(), Spec: &rvboxv1.ExecutionSpec{ShellType: rvboxv1.ShellType_SHELL_CMD, Source: &rvboxv1.ExecutionSpec_CommandText{CommandText: "echo"}}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
first, err := service.ListCommands(ctx, &rvboxv1.ListCommandsRequest{ClientId: "win-a", PageSize: 1})
|
||||
if err != nil || len(first.GetCommands()) != 1 || first.GetNextPageToken() == "" {
|
||||
t.Fatalf("first command page = %#v, %v", first, err)
|
||||
}
|
||||
changedFilter := &rvboxv1.ListCommandsRequest{ClientId: "win-a", IncludeTerminal: true, PageSize: 1, PageToken: first.GetNextPageToken()}
|
||||
if _, err := service.ListCommands(ctx, changedFilter); status.Code(err) != codes.InvalidArgument {
|
||||
t.Fatalf("changed-filter cursor code = %v", status.Code(err))
|
||||
}
|
||||
second, err := service.ListCommands(ctx, &rvboxv1.ListCommandsRequest{ClientId: "win-a", PageSize: 1, PageToken: first.GetNextPageToken()})
|
||||
if err != nil || len(second.GetCommands()) != 1 || second.GetCommands()[0].GetIssueUuid() == first.GetCommands()[0].GetIssueUuid() {
|
||||
t.Fatalf("second command page = %#v, %v", second, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestControlGRPCRoundTrip_HP_CONTROL_06(t *testing.T) {
|
||||
service, persistence := newTestService(t)
|
||||
defer persistence.Close()
|
||||
registerControlClient(t, persistence, "win-a", rvboxv1.Platform_PLATFORM_WINDOWS, rvboxv1.ShellType_SHELL_CMD, 5)
|
||||
grpcServer := grpc.NewServer()
|
||||
rvboxv1.RegisterControlServer(grpcServer, service)
|
||||
listener := bufconn.Listen(1 << 20)
|
||||
go func() { _ = grpcServer.Serve(listener) }()
|
||||
defer grpcServer.Stop()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
connection, err := grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() }), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithBlock())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer connection.Close()
|
||||
client := rvboxv1.NewControlClient(connection)
|
||||
response, err := client.RunCommand(ctx, &rvboxv1.RunCommandRequest{TargetClientId: "win-a", RequestId: fixedIssue(0xa5).String(), Spec: &rvboxv1.ExecutionSpec{ShellType: rvboxv1.ShellType_SHELL_CMD, Source: &rvboxv1.ExecutionSpec_CommandText{CommandText: "echo wire"}}})
|
||||
if err != nil || response.GetLifecycle() != rvboxv1.CommandLifecycle_COMMAND_QUEUED {
|
||||
t.Fatalf("gRPC run response = %#v, %v", response, err)
|
||||
}
|
||||
}
|
||||
|
||||
func newTestService(t *testing.T) (*Service, *store.Store) {
|
||||
t.Helper()
|
||||
persistence, err := store.Open(context.Background(), store.Options{DataDir: filepath.Join(t.TempDir(), "state"), BusyTimeout: time.Second})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
service, err := NewService(Options{Store: persistence, CursorKey: bytesKey(), Now: func() time.Time { return time.Date(2026, time.September, 6, 12, 0, 0, 0, time.UTC) }, Limits: agentproto.DefaultLimits()})
|
||||
if err != nil {
|
||||
persistence.Close()
|
||||
t.Fatal(err)
|
||||
}
|
||||
return service, persistence
|
||||
}
|
||||
|
||||
func bytesKey() []byte { return []byte("0123456789abcdef0123456789abcdef") }
|
||||
|
||||
func registerControlClient(t *testing.T, persistence *store.Store, clientID string, platform rvboxv1.Platform, shell rvboxv1.ShellType, seed byte) {
|
||||
t.Helper()
|
||||
shells, err := proto.Marshal(&rvboxv1.ClientHello{SupportedShells: []rvboxv1.ShellType{shell}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := persistence.RegisterClientSession(context.Background(), store.ClientRegistration{ClientID: clientID, Platform: uint32(platform), Architecture: "amd64", DaemonVersion: "test", DaemonCWD: `C:\\ProgramData\\RVBox`, SupportedShells: shells, ClientInstanceID: [16]byte{seed}, SessionID: [16]byte{seed + 10}, ConnectedAt: time.Date(2026, time.September, 6, 11, 0, 0, 0, time.UTC)}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func fixedIssue(last byte) domain.UUID {
|
||||
value, _ := domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000001")
|
||||
value[15] = last
|
||||
return value
|
||||
}
|
||||
Reference in New Issue
Block a user