v0.1.0 (#17)

internal/envdclient/process.go

@@ -0,0 +1,187 @@
+package envdclient
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+
+	"connectrpc.com/connect"
+
+	envdpb "git.omukk.dev/wrenn/wrenn/proto/envd/gen"
+)
+
+// ProcessInfo holds metadata about a running process inside the sandbox.
+type ProcessInfo struct {
+	PID  uint32
+	Tag  string
+	Cmd  string
+	Args []string
+}
+
+// StartBackground starts a process that runs independently of the RPC stream.
+// It opens a Start stream, reads the first StartEvent to obtain the PID,
+// then closes the stream. The process continues running inside the VM because
+// envd binds it to context.Background().
+func (c *Client) StartBackground(ctx context.Context, tag, cmd string, args []string, envs map[string]string, cwd string) (uint32, error) {
+	stdin := false
+	cfg := &envdpb.ProcessConfig{
+		Cmd:  cmd,
+		Args: args,
+		Envs: envs,
+	}
+	if cwd != "" {
+		cfg.Cwd = &cwd
+	}
+
+	req := connect.NewRequest(&envdpb.StartRequest{
+		Process: cfg,
+		Tag:     &tag,
+		Stdin:   &stdin,
+	})
+
+	stream, err := c.process.Start(ctx, req)
+	if err != nil {
+		return 0, fmt.Errorf("start background process: %w", err)
+	}
+	defer stream.Close()
+
+	// Read events until we get the StartEvent with the PID.
+	for stream.Receive() {
+		msg := stream.Msg()
+		if msg.Event == nil {
+			continue
+		}
+		if start, ok := msg.Event.GetEvent().(*envdpb.ProcessEvent_Start); ok {
+			return start.Start.GetPid(), nil
+		}
+	}
+
+	if err := stream.Err(); err != nil && err != io.EOF {
+		return 0, fmt.Errorf("start background process stream: %w", err)
+	}
+
+	return 0, fmt.Errorf("start background process: no start event received")
+}
+
+// ConnectProcess re-attaches to a running process by PID or tag and returns
+// a channel of streaming events. The channel is closed when the process ends
+// or the context is cancelled.
+func (c *Client) ConnectProcess(ctx context.Context, pid uint32, tag string) (<-chan ExecStreamEvent, error) {
+	var selector *envdpb.ProcessSelector
+	if tag != "" {
+		selector = &envdpb.ProcessSelector{
+			Selector: &envdpb.ProcessSelector_Tag{Tag: tag},
+		}
+	} else {
+		selector = &envdpb.ProcessSelector{
+			Selector: &envdpb.ProcessSelector_Pid{Pid: pid},
+		}
+	}
+
+	stream, err := c.process.Connect(ctx, connect.NewRequest(&envdpb.ConnectRequest{
+		Process: selector,
+	}))
+	if err != nil {
+		return nil, fmt.Errorf("connect process: %w", err)
+	}
+
+	ch := make(chan ExecStreamEvent, 16)
+	go func() {
+		defer close(ch)
+		defer stream.Close()
+
+		for stream.Receive() {
+			msg := stream.Msg()
+			if msg.Event == nil {
+				continue
+			}
+
+			var ev ExecStreamEvent
+			switch e := msg.Event.GetEvent().(type) {
+			case *envdpb.ProcessEvent_Start:
+				ev = ExecStreamEvent{Type: "start", PID: e.Start.GetPid()}
+
+			case *envdpb.ProcessEvent_Data:
+				switch o := e.Data.GetOutput().(type) {
+				case *envdpb.ProcessEvent_DataEvent_Stdout:
+					ev = ExecStreamEvent{Type: "stdout", Data: o.Stdout}
+				case *envdpb.ProcessEvent_DataEvent_Stderr:
+					ev = ExecStreamEvent{Type: "stderr", Data: o.Stderr}
+				default:
+					continue
+				}
+
+			case *envdpb.ProcessEvent_End:
+				ev = ExecStreamEvent{Type: "end", ExitCode: e.End.GetExitCode()}
+				if e.End.Error != nil {
+					ev.Error = e.End.GetError()
+				}
+
+			case *envdpb.ProcessEvent_Keepalive:
+				continue
+			}
+
+			select {
+			case ch <- ev:
+			case <-ctx.Done():
+				return
+			}
+		}
+
+		if err := stream.Err(); err != nil && err != io.EOF {
+			slog.Debug("connect process stream error", "error", err)
+		}
+	}()
+
+	return ch, nil
+}
+
+// ListProcesses returns all running processes inside the sandbox.
+func (c *Client) ListProcesses(ctx context.Context) ([]ProcessInfo, error) {
+	resp, err := c.process.List(ctx, connect.NewRequest(&envdpb.ListRequest{}))
+	if err != nil {
+		return nil, fmt.Errorf("list processes: %w", err)
+	}
+
+	procs := make([]ProcessInfo, 0, len(resp.Msg.Processes))
+	for _, p := range resp.Msg.Processes {
+		info := ProcessInfo{
+			PID: p.Pid,
+		}
+		if p.Tag != nil {
+			info.Tag = *p.Tag
+		}
+		if p.Config != nil {
+			info.Cmd = p.Config.Cmd
+			info.Args = p.Config.Args
+		}
+		procs = append(procs, info)
+	}
+
+	return procs, nil
+}
+
+// KillProcess sends a signal to a process identified by PID or tag.
+func (c *Client) KillProcess(ctx context.Context, pid uint32, tag string, signal envdpb.Signal) error {
+	var selector *envdpb.ProcessSelector
+	if tag != "" {
+		selector = &envdpb.ProcessSelector{
+			Selector: &envdpb.ProcessSelector_Tag{Tag: tag},
+		}
+	} else {
+		selector = &envdpb.ProcessSelector{
+			Selector: &envdpb.ProcessSelector_Pid{Pid: pid},
+		}
+	}
+
+	_, err := c.process.SendSignal(ctx, connect.NewRequest(&envdpb.SendSignalRequest{
+		Process: selector,
+		Signal:  signal,
+	}))
+	if err != nil {
+		return fmt.Errorf("kill process: %w", err)
+	}
+
+	return nil
+}