From e22f6fbeefc94b74d2dca3e778b6bb9c3594c395 Mon Sep 17 00:00:00 2001 From: Fedor Kororkov Date: Mon, 4 May 2026 15:36:20 -0400 Subject: [PATCH] Add more options to `exec` (#47) --- .cirrus.yml | 5 +-- internal/rpc/agent.pb.go | 80 ++++++++++++++++++++++++++++------------ internal/rpc/exec.go | 77 +++++++++++++++++++++++++++++++++++++- proto/rpc/agent.proto | 3 ++ 4 files changed, 137 insertions(+), 28 deletions(-) diff --git a/.cirrus.yml b/.cirrus.yml index 02629b4..b07270d 100644 --- a/.cirrus.yml +++ b/.cirrus.yml @@ -46,17 +46,16 @@ task: check_script: git diff --exit-code task: - only_if: $CIRRUS_BRANCH != '' + only_if: $CIRRUS_BRANCH != '' && $CIRRUS_PR == '' && $CIRRUS_REPO_OWNER == 'cirruslabs' name: buf push container: image: bufbuild/buf - login_script: echo $BUF_TOKEN | buf registry login --username $BUF_LOGIN --token-stdin + login_script: echo "$BUF_TOKEN" | buf registry login --token-stdin push_script: buf push --git-metadata env: - BUF_LOGIN: fkorotkov BUF_TOKEN: ENCRYPTED[!8ee7eb2504cc84b08d4a7c0dacbe103640b1feaa26d06f0df010784e872d39e65a0cdea3fc7c09b065a917a77113b96b!] task: diff --git a/internal/rpc/agent.pb.go b/internal/rpc/agent.pb.go index 56e90d1..c1b27fa 100644 --- a/internal/rpc/agent.pb.go +++ b/internal/rpc/agent.pb.go @@ -401,6 +401,9 @@ type ExecRequest_Command struct { Interactive bool `protobuf:"varint,3,opt,name=interactive,proto3" json:"interactive,omitempty"` Tty bool `protobuf:"varint,4,opt,name=tty,proto3" json:"tty,omitempty"` TerminalSize *TerminalSize `protobuf:"bytes,5,opt,name=terminal_size,json=terminalSize,proto3" json:"terminal_size,omitempty"` + Detach bool `protobuf:"varint,6,opt,name=detach,proto3" json:"detach,omitempty"` + Env map[string]string `protobuf:"bytes,7,rep,name=env,proto3" json:"env,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + Workdir string `protobuf:"bytes,8,opt,name=workdir,proto3" json:"workdir,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -470,6 +473,27 @@ func (x *ExecRequest_Command) GetTerminalSize() *TerminalSize { return nil } +func (x *ExecRequest_Command) GetDetach() bool { + if x != nil { + return x.Detach + } + return false +} + +func (x *ExecRequest_Command) GetEnv() map[string]string { + if x != nil { + return x.Env + } + return nil +} + +func (x *ExecRequest_Command) GetWorkdir() string { + if x != nil { + return x.Workdir + } + return "" +} + type ExecResponse_Exit struct { state protoimpl.MessageState `protogen:"open.v1"` Code int32 `protobuf:"varint,1,opt,name=code,proto3" json:"code,omitempty"` @@ -479,7 +503,7 @@ type ExecResponse_Exit struct { func (x *ExecResponse_Exit) Reset() { *x = ExecResponse_Exit{} - mi := &file_rpc_agent_proto_msgTypes[7] + mi := &file_rpc_agent_proto_msgTypes[8] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -491,7 +515,7 @@ func (x *ExecResponse_Exit) String() string { func (*ExecResponse_Exit) ProtoMessage() {} func (x *ExecResponse_Exit) ProtoReflect() protoreflect.Message { - mi := &file_rpc_agent_proto_msgTypes[7] + mi := &file_rpc_agent_proto_msgTypes[8] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -518,17 +542,23 @@ var File_rpc_agent_proto protoreflect.FileDescriptor const file_rpc_agent_proto_rawDesc = "" + "\n" + - "\x0frpc/agent.proto\x1a\x1bgoogle/protobuf/empty.proto\"\xd0\x02\n" + + "\x0frpc/agent.proto\x1a\x1bgoogle/protobuf/empty.proto\"\xeb\x03\n" + "\vExecRequest\x120\n" + "\acommand\x18\x01 \x01(\v2\x14.ExecRequest.CommandH\x00R\acommand\x121\n" + "\x0estandard_input\x18\x02 \x01(\v2\b.IOChunkH\x00R\rstandardInput\x128\n" + - "\x0fterminal_resize\x18\x03 \x01(\v2\r.TerminalSizeH\x00R\x0eterminalResize\x1a\x99\x01\n" + + "\x0fterminal_resize\x18\x03 \x01(\v2\r.TerminalSizeH\x00R\x0eterminalResize\x1a\xb4\x02\n" + "\aCommand\x12\x12\n" + "\x04name\x18\x01 \x01(\tR\x04name\x12\x12\n" + "\x04args\x18\x02 \x03(\tR\x04args\x12 \n" + "\vinteractive\x18\x03 \x01(\bR\vinteractive\x12\x10\n" + "\x03tty\x18\x04 \x01(\bR\x03tty\x122\n" + - "\rterminal_size\x18\x05 \x01(\v2\r.TerminalSizeR\fterminalSizeB\x06\n" + + "\rterminal_size\x18\x05 \x01(\v2\r.TerminalSizeR\fterminalSize\x12\x16\n" + + "\x06detach\x18\x06 \x01(\bR\x06detach\x12/\n" + + "\x03env\x18\a \x03(\v2\x1d.ExecRequest.Command.EnvEntryR\x03env\x12\x18\n" + + "\aworkdir\x18\b \x01(\tR\aworkdir\x1a6\n" + + "\bEnvEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01B\x06\n" + "\x04type\"\xc4\x01\n" + "\fExecResponse\x12(\n" + "\x04exit\x18\x01 \x01(\v2\x12.ExecResponse.ExitH\x00R\x04exit\x123\n" + @@ -561,7 +591,7 @@ func file_rpc_agent_proto_rawDescGZIP() []byte { return file_rpc_agent_proto_rawDescData } -var file_rpc_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 8) +var file_rpc_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 9) var file_rpc_agent_proto_goTypes = []any{ (*ExecRequest)(nil), // 0: ExecRequest (*ExecResponse)(nil), // 1: ExecResponse @@ -570,25 +600,27 @@ var file_rpc_agent_proto_goTypes = []any{ (*ResolveIPRequest)(nil), // 4: ResolveIPRequest (*ResolveIPResponse)(nil), // 5: ResolveIPResponse (*ExecRequest_Command)(nil), // 6: ExecRequest.Command - (*ExecResponse_Exit)(nil), // 7: ExecResponse.Exit + nil, // 7: ExecRequest.Command.EnvEntry + (*ExecResponse_Exit)(nil), // 8: ExecResponse.Exit } var file_rpc_agent_proto_depIdxs = []int32{ - 6, // 0: ExecRequest.command:type_name -> ExecRequest.Command - 3, // 1: ExecRequest.standard_input:type_name -> IOChunk - 2, // 2: ExecRequest.terminal_resize:type_name -> TerminalSize - 7, // 3: ExecResponse.exit:type_name -> ExecResponse.Exit - 3, // 4: ExecResponse.standard_output:type_name -> IOChunk - 3, // 5: ExecResponse.standard_error:type_name -> IOChunk - 2, // 6: ExecRequest.Command.terminal_size:type_name -> TerminalSize - 0, // 7: Agent.Exec:input_type -> ExecRequest - 4, // 8: Agent.ResolveIP:input_type -> ResolveIPRequest - 1, // 9: Agent.Exec:output_type -> ExecResponse - 5, // 10: Agent.ResolveIP:output_type -> ResolveIPResponse - 9, // [9:11] is the sub-list for method output_type - 7, // [7:9] is the sub-list for method input_type - 7, // [7:7] is the sub-list for extension type_name - 7, // [7:7] is the sub-list for extension extendee - 0, // [0:7] is the sub-list for field type_name + 6, // 0: ExecRequest.command:type_name -> ExecRequest.Command + 3, // 1: ExecRequest.standard_input:type_name -> IOChunk + 2, // 2: ExecRequest.terminal_resize:type_name -> TerminalSize + 8, // 3: ExecResponse.exit:type_name -> ExecResponse.Exit + 3, // 4: ExecResponse.standard_output:type_name -> IOChunk + 3, // 5: ExecResponse.standard_error:type_name -> IOChunk + 2, // 6: ExecRequest.Command.terminal_size:type_name -> TerminalSize + 7, // 7: ExecRequest.Command.env:type_name -> ExecRequest.Command.EnvEntry + 0, // 8: Agent.Exec:input_type -> ExecRequest + 4, // 9: Agent.ResolveIP:input_type -> ResolveIPRequest + 1, // 10: Agent.Exec:output_type -> ExecResponse + 5, // 11: Agent.ResolveIP:output_type -> ResolveIPResponse + 10, // [10:12] is the sub-list for method output_type + 8, // [8:10] is the sub-list for method input_type + 8, // [8:8] is the sub-list for extension type_name + 8, // [8:8] is the sub-list for extension extendee + 0, // [0:8] is the sub-list for field type_name } func init() { file_rpc_agent_proto_init() } @@ -612,7 +644,7 @@ func file_rpc_agent_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_rpc_agent_proto_rawDesc), len(file_rpc_agent_proto_rawDesc)), NumEnums: 0, - NumMessages: 8, + NumMessages: 9, NumExtensions: 0, NumServices: 1, }, diff --git a/internal/rpc/exec.go b/internal/rpc/exec.go index 7e501d1..0e934ab 100644 --- a/internal/rpc/exec.go +++ b/internal/rpc/exec.go @@ -14,6 +14,7 @@ import ( "os/exec" "slices" "strings" + "syscall" ) const ( @@ -36,9 +37,47 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse]) zap.S().Infof("executing %s", formatCommandAndArgs(firstExecRequestCommand.Command.Name, firstExecRequestCommand.Command.Args)) + if firstExecRequestCommand.Command.Detach && + (firstExecRequestCommand.Command.Interactive || firstExecRequestCommand.Command.Tty) { + return fmt.Errorf("detach cannot be used with interactive or tty") + } + // Execute the command - cmd := exec.CommandContext(stream.Context(), firstExecRequestCommand.Command.Name, + execCtx := stream.Context() + if firstExecRequestCommand.Command.Detach { + execCtx = context.Background() + } + + cmd := exec.CommandContext(execCtx, firstExecRequestCommand.Command.Name, firstExecRequestCommand.Command.Args...) + applyExecOverrides(cmd, firstExecRequestCommand.Command) + + if firstExecRequestCommand.Command.Detach { + cmd.Stdout = io.Discard + cmd.Stderr = io.Discard + cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} + + if err := cmd.Start(); err != nil { + return err + } + if cmd.Process != nil { + if err := cmd.Process.Release(); err != nil { + return err + } + } + + if err := stream.Send(&ExecResponse{ + Type: &ExecResponse_Exit_{ + Exit: &ExecResponse_Exit{ + Code: 0, + }, + }, + }); err != nil && !errors.Is(err, context.Canceled) { + return err + } + + return nil + } var stdin io.WriteCloser var stdout, stderr io.ReadCloser @@ -239,6 +278,42 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse]) }) } +func applyExecOverrides(cmd *exec.Cmd, command *ExecRequest_Command) { + if command.Workdir != "" { + cmd.Dir = command.Workdir + } + + if len(command.Env) > 0 { + cmd.Env = mergeEnv(command.Env) + } +} + +func mergeEnv(overrides map[string]string) []string { + if len(overrides) == 0 { + return os.Environ() + } + + envMap := make(map[string]string, len(overrides)) + for _, entry := range os.Environ() { + parts := strings.SplitN(entry, "=", 2) + if len(parts) != 2 { + continue + } + envMap[parts[0]] = parts[1] + } + + for key, value := range overrides { + envMap[key] = value + } + + merged := make([]string, 0, len(envMap)) + for key, value := range envMap { + merged = append(merged, key+"="+value) + } + + return merged +} + func formatCommandAndArgs(name string, args []string) string { var all []string diff --git a/proto/rpc/agent.proto b/proto/rpc/agent.proto index 86adf18..638fdd7 100644 --- a/proto/rpc/agent.proto +++ b/proto/rpc/agent.proto @@ -16,6 +16,9 @@ message ExecRequest { bool interactive = 3; bool tty = 4; TerminalSize terminal_size = 5; + bool detach = 6; + map env = 7; + string workdir = 8; } oneof type {