v0.1.0 (#17)

2026-04-16 19:24:25 +00:00
parent 172413e91e
commit 605ad666a0
239 changed files with 19966 additions and 3454 deletions
--- a/internal/envdclient/client.go
+++ b/internal/envdclient/client.go
@ -3,6 +3,7 @@ package envdclient
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 	"log/slog"
@ -268,6 +269,82 @@ func (c *Client) ReadFile(ctx context.Context, path string) ([]byte, error) {
 	return data, nil
 }

+// PrepareSnapshot calls envd's POST /snapshot/prepare endpoint, which quiesces
+// continuous goroutines (port scanner, forwarder) and forces a GC cycle before
+// Firecracker takes a VM snapshot. This ensures the Go runtime's page allocator
+// is in a consistent state when vCPUs are frozen.
+//
+// Best-effort: the caller should log a warning on error but not abort the pause.
+func (c *Client) PrepareSnapshot(ctx context.Context) error {
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.base+"/snapshot/prepare", nil)
+	if err != nil {
+		return fmt.Errorf("create request: %w", err)
+	}
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return fmt.Errorf("prepare snapshot: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusNoContent {
+		respBody, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("prepare snapshot: status %d: %s", resp.StatusCode, string(respBody))
+	}
+
+	return nil
+}
+
+// PostInit calls envd's POST /init endpoint, which triggers a re-read of
+// Firecracker MMDS metadata. This updates WRENN_SANDBOX_ID, WRENN_TEMPLATE_ID
+// env vars and the corresponding files under /run/wrenn/ inside the guest.
+// Must be called after snapshot restore so envd picks up the new sandbox's metadata.
+func (c *Client) PostInit(ctx context.Context) error {
+	return c.PostInitWithDefaults(ctx, "", nil)
+}
+
+// PostInitWithDefaults calls envd's POST /init endpoint with optional default
+// user and environment variables. These are applied to envd's defaults so all
+// subsequent process executions use them.
+func (c *Client) PostInitWithDefaults(ctx context.Context, defaultUser string, envVars map[string]string) error {
+	var body io.Reader
+	if defaultUser != "" || len(envVars) > 0 {
+		payload := make(map[string]any)
+		if defaultUser != "" {
+			payload["defaultUser"] = defaultUser
+		}
+		if len(envVars) > 0 {
+			payload["envVars"] = envVars
+		}
+		data, err := json.Marshal(payload)
+		if err != nil {
+			return fmt.Errorf("marshal init body: %w", err)
+		}
+		body = bytes.NewReader(data)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.base+"/init", body)
+	if err != nil {
+		return fmt.Errorf("create request: %w", err)
+	}
+	if body != nil {
+		req.Header.Set("Content-Type", "application/json")
+	}
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return fmt.Errorf("post init: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusNoContent {
+		respBody, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("post init: status %d: %s", resp.StatusCode, string(respBody))
+	}
+
+	return nil
+}
+
 // ListDir lists directory contents inside the sandbox.
 func (c *Client) ListDir(ctx context.Context, path string, depth uint32) (*envdpb.ListDirResponse, error) {
 	req := connect.NewRequest(&envdpb.ListDirRequest{
@ -282,3 +359,30 @@ func (c *Client) ListDir(ctx context.Context, path string, depth uint32) (*envdp

 	return resp.Msg, nil
 }
+
+// MakeDir creates a directory inside the sandbox.
+func (c *Client) MakeDir(ctx context.Context, path string) (*envdpb.MakeDirResponse, error) {
+	req := connect.NewRequest(&envdpb.MakeDirRequest{
+		Path: path,
+	})
+
+	resp, err := c.filesystem.MakeDir(ctx, req)
+	if err != nil {
+		return nil, fmt.Errorf("make dir: %w", err)
+	}
+
+	return resp.Msg, nil
+}
+
+// Remove removes a file or directory inside the sandbox.
+func (c *Client) Remove(ctx context.Context, path string) error {
+	req := connect.NewRequest(&envdpb.RemoveRequest{
+		Path: path,
+	})
+
+	if _, err := c.filesystem.Remove(ctx, req); err != nil {
+		return fmt.Errorf("remove: %w", err)
+	}
+
+	return nil
+}
--- a/internal/envdclient/health.go
+++ b/internal/envdclient/health.go
@ -2,7 +2,9 @@ package envdclient

 import (
 	"context"
+	"encoding/json"
 	"fmt"
+	"io"
 	"log/slog"
 	"net/http"
 	"time"
@ -31,6 +33,38 @@ func (c *Client) WaitUntilReady(ctx context.Context) error {
 	}
 }

+// FetchVersion queries envd's health endpoint and returns the reported version.
+func (c *Client) FetchVersion(ctx context.Context) (string, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.healthURL, nil)
+	if err != nil {
+		return "", fmt.Errorf("build health request: %w", err)
+	}
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return "", fmt.Errorf("fetch envd version: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusNoContent {
+		return "", fmt.Errorf("health check returned %d", resp.StatusCode)
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil || len(body) == 0 {
+		return "", nil // envd may not support version reporting yet
+	}
+
+	var data struct {
+		Version string `json:"version"`
+	}
+	if err := json.Unmarshal(body, &data); err != nil {
+		return "", nil // non-JSON response, old envd
+	}
+
+	return data.Version, nil
+}
+
 // healthCheck sends a single GET /health request to envd.
 func (c *Client) healthCheck(ctx context.Context) error {
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.healthURL, nil)
--- a/internal/envdclient/process.go
+++ b/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
+}
--- a/internal/envdclient/pty.go
+++ b/internal/envdclient/pty.go
@ -0,0 +1,220 @@
+package envdclient
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+
+	"connectrpc.com/connect"
+
+	envdpb "git.omukk.dev/wrenn/wrenn/proto/envd/gen"
+)
+
+// PtyEvent represents a single event from a PTY output stream.
+type PtyEvent struct {
+	Type     string // "started", "output", "end"
+	PID      uint32
+	Data     []byte
+	ExitCode int32
+	Error    string
+}
+
+// PtyStart starts a new PTY process in the guest and returns a channel of events.
+// The tag is the stable identifier used to reconnect via PtyConnect.
+// The channel is closed when the process ends or ctx is cancelled.
+// NOTE: The user parameter from PtyAttachRequest is not yet supported by envd's
+// ProcessConfig proto. When envd adds user support, thread it through here.
+func (c *Client) PtyStart(ctx context.Context, tag, cmd string, args []string, cols, rows uint32, envs map[string]string, cwd string) (<-chan PtyEvent, error) {
+	stdin := true
+	cfg := &envdpb.ProcessConfig{
+		Cmd:  cmd,
+		Args: args,
+		Envs: envs,
+	}
+	if cwd != "" {
+		cfg.Cwd = &cwd
+	}
+
+	req := connect.NewRequest(&envdpb.StartRequest{
+		Process: cfg,
+		Pty: &envdpb.PTY{
+			Size: &envdpb.PTY_Size{
+				Cols: cols,
+				Rows: rows,
+			},
+		},
+		Tag:   &tag,
+		Stdin: &stdin,
+	})
+
+	stream, err := c.process.Start(ctx, req)
+	if err != nil {
+		return nil, fmt.Errorf("pty start: %w", err)
+	}
+
+	return drainPtyStream(ctx, &startStream{s: stream}, true), nil
+}
+
+// PtyConnect re-attaches to an existing PTY process by tag.
+// Returns a channel of output events starting from the current point.
+func (c *Client) PtyConnect(ctx context.Context, tag string) (<-chan PtyEvent, error) {
+	req := connect.NewRequest(&envdpb.ConnectRequest{
+		Process: &envdpb.ProcessSelector{
+			Selector: &envdpb.ProcessSelector_Tag{Tag: tag},
+		},
+	})
+
+	stream, err := c.process.Connect(ctx, req)
+	if err != nil {
+		return nil, fmt.Errorf("pty connect: %w", err)
+	}
+
+	return drainPtyStream(ctx, &connectStream{s: stream}, false), nil
+}
+
+// PtySendInput sends raw bytes to the PTY process identified by tag.
+func (c *Client) PtySendInput(ctx context.Context, tag string, data []byte) error {
+	req := connect.NewRequest(&envdpb.SendInputRequest{
+		Process: &envdpb.ProcessSelector{
+			Selector: &envdpb.ProcessSelector_Tag{Tag: tag},
+		},
+		Input: &envdpb.ProcessInput{
+			Input: &envdpb.ProcessInput_Pty{Pty: data},
+		},
+	})
+
+	if _, err := c.process.SendInput(ctx, req); err != nil {
+		return fmt.Errorf("pty send input: %w", err)
+	}
+	return nil
+}
+
+// PtyResize updates the terminal dimensions for the PTY process identified by tag.
+func (c *Client) PtyResize(ctx context.Context, tag string, cols, rows uint32) error {
+	req := connect.NewRequest(&envdpb.UpdateRequest{
+		Process: &envdpb.ProcessSelector{
+			Selector: &envdpb.ProcessSelector_Tag{Tag: tag},
+		},
+		Pty: &envdpb.PTY{
+			Size: &envdpb.PTY_Size{
+				Cols: cols,
+				Rows: rows,
+			},
+		},
+	})
+
+	if _, err := c.process.Update(ctx, req); err != nil {
+		return fmt.Errorf("pty resize: %w", err)
+	}
+	return nil
+}
+
+// PtyKill sends SIGKILL to the PTY process identified by tag.
+func (c *Client) PtyKill(ctx context.Context, tag string) error {
+	req := connect.NewRequest(&envdpb.SendSignalRequest{
+		Process: &envdpb.ProcessSelector{
+			Selector: &envdpb.ProcessSelector_Tag{Tag: tag},
+		},
+		Signal: envdpb.Signal_SIGNAL_SIGKILL,
+	})
+
+	if _, err := c.process.SendSignal(ctx, req); err != nil {
+		return fmt.Errorf("pty kill: %w", err)
+	}
+	return nil
+}
+
+// eventStream is an interface covering both StartResponse and ConnectResponse streams.
+type eventStream interface {
+	Receive() bool
+	Err() error
+	Close() error
+}
+
+type startStream struct {
+	s *connect.ServerStreamForClient[envdpb.StartResponse]
+}
+
+func (s *startStream) Receive() bool { return s.s.Receive() }
+func (s *startStream) Err() error    { return s.s.Err() }
+func (s *startStream) Close() error  { return s.s.Close() }
+func (s *startStream) Event() *envdpb.ProcessEvent {
+	return s.s.Msg().GetEvent()
+}
+
+type connectStream struct {
+	s *connect.ServerStreamForClient[envdpb.ConnectResponse]
+}
+
+func (s *connectStream) Receive() bool { return s.s.Receive() }
+func (s *connectStream) Err() error    { return s.s.Err() }
+func (s *connectStream) Close() error  { return s.s.Close() }
+func (s *connectStream) Event() *envdpb.ProcessEvent {
+	return s.s.Msg().GetEvent()
+}
+
+type eventProvider interface {
+	eventStream
+	Event() *envdpb.ProcessEvent
+}
+
+// drainPtyStream reads events from either a Start or Connect stream and maps
+// them into PtyEvent values on a channel.
+func drainPtyStream(ctx context.Context, stream eventProvider, expectStart bool) <-chan PtyEvent {
+	ch := make(chan PtyEvent, 16)
+	go func() {
+		defer close(ch)
+		defer stream.Close()
+
+		for stream.Receive() {
+			event := stream.Event()
+			if event == nil {
+				continue
+			}
+
+			var ev PtyEvent
+			switch e := event.GetEvent().(type) {
+			case *envdpb.ProcessEvent_Start:
+				if expectStart {
+					ev = PtyEvent{Type: "started", PID: e.Start.GetPid()}
+				} else {
+					continue
+				}
+
+			case *envdpb.ProcessEvent_Data:
+				switch o := e.Data.GetOutput().(type) {
+				case *envdpb.ProcessEvent_DataEvent_Pty:
+					ev = PtyEvent{Type: "output", Data: o.Pty}
+				case *envdpb.ProcessEvent_DataEvent_Stdout:
+					ev = PtyEvent{Type: "output", Data: o.Stdout}
+				case *envdpb.ProcessEvent_DataEvent_Stderr:
+					ev = PtyEvent{Type: "output", Data: o.Stderr}
+				default:
+					continue
+				}
+
+			case *envdpb.ProcessEvent_End:
+				ev = PtyEvent{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("pty stream error", "error", err)
+		}
+	}()
+
+	return ch
+}