codeexecutor/e2b: add native envd process client

codeexecutor/e2b/internal/envdprocess/client_test.go

@@ -0,0 +1,227 @@
+//
+// Tencent is pleased to support the open source community by making trpc-agent-go available.
+//
+// Copyright (C) 2025 Tencent.  All rights reserved.
+//
+// trpc-agent-go is licensed under the Apache License Version 2.0.
+//
+
+package envdprocess
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewClientRejectsInvalidBaseURL(t *testing.T) {
+	_, err := NewClient("not-a-url", nil, nil)
+	require.Error(t, err)
+}
+
+func TestNewClientUsesDefaultHTTPClient(t *testing.T) {
+	client, err := NewClient("https://envd.example", nil, nil)
+	require.NoError(t, err)
+	assert.NotNil(t, client.processClient)
+}
+
+func TestNewClientRejectsInsecureRemoteBaseURL(t *testing.T) {
+	_, err := NewClient("http://envd.example", nil, nil)
+	require.ErrorContains(t, err, "remote base URL must use HTTPS")
+}
+
+func TestNewClientRejectsBaseURLCredentials(t *testing.T) {
+	_, err := NewClient("https://user:secret@envd.example", nil, nil)
+	require.ErrorContains(t, err, "must not contain user credentials")
+}
+
+func TestNewClientAllowsCredentiallessLoopbackHTTP(t *testing.T) {
+	for _, baseURL := range []string{
+		"http://localhost:49983",
+		"http://127.0.0.1:49983",
+		"http://[::1]:49983",
+	} {
+		t.Run(baseURL, func(t *testing.T) {
+			client, err := NewClient(baseURL, nil, nil)
+			require.NoError(t, err)
+			assert.False(t, client.credentialsAllowed)
+		})
+	}
+}
+
+func TestNewClientRejectsHeadersOverLoopbackHTTP(t *testing.T) {
+	_, err := NewClient(
+		"http://127.0.0.1:49983",
+		nil,
+		http.Header{"X-Access-Token": {"secret"}},
+	)
+	require.ErrorContains(t, err, "configured headers require HTTPS")
+}
+
+func TestStartRejectsProcessUserOverLoopbackHTTP(t *testing.T) {
+	client, err := NewClient("http://127.0.0.1:49983", nil, nil)
+	require.NoError(t, err)
+
+	_, err = client.Start(context.Background(), Request{
+		Cmd:  "true",
+		User: "sandbox-user",
+	})
+	require.ErrorContains(t, err, "process user requires HTTPS")
+}
+
+func TestOriginBoundHTTPClientPreservesSuppliedConfiguration(t *testing.T) {
+	httpsOrigin := &http.Request{URL: mustParseURL(
+		t, "https://envd.example/process.Process/Start",
+	)}
+	callerErr := errors.New("caller redirect policy")
+	policyCalls := 0
+	suppliedClient := &http.Client{
+		Transport: http.DefaultTransport,
+		Timeout:   time.Second,
+		CheckRedirect: func(
+			*http.Request,
+			[]*http.Request,
+		) error {
+			policyCalls++
+			return callerErr
+		},
+	}
+	boundClient := newOriginBoundHTTPClient(
+		suppliedClient,
+		mustParseURL(t, "https://envd.example"),
+	)
+	require.NotSame(t, suppliedClient, boundClient)
+	assert.Equal(t, suppliedClient.Timeout, boundClient.Timeout)
+	assert.Same(t, http.DefaultTransport, suppliedClient.Transport)
+	boundTransport, ok := boundClient.Transport.(*originBoundRoundTripper)
+	require.True(t, ok)
+	assert.Same(t, http.DefaultTransport, boundTransport.base)
+
+	err := boundClient.CheckRedirect(&http.Request{URL: mustParseURL(
+		t, "https://envd.example/process.Process/Start/",
+	)}, []*http.Request{httpsOrigin})
+	require.ErrorIs(t, err, callerErr)
+	assert.Equal(t, 1, policyCalls)
+}
+
+func TestOriginBoundHTTPClientAllowsSameOriginRedirect(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(
+		w http.ResponseWriter,
+		r *http.Request,
+	) {
+		if r.URL.Path == "/start" {
+			http.Redirect(w, r, "/finish", http.StatusTemporaryRedirect)
+			return
+		}
+		assert.Equal(t, "/finish", r.URL.Path)
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	t.Cleanup(server.Close)
+
+	client := newOriginBoundHTTPClient(
+		server.Client(), mustParseURL(t, server.URL),
+	)
+	resp, err := client.Get(server.URL + "/start")
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, resp.Body.Close()) })
+	assert.Equal(t, http.StatusNoContent, resp.StatusCode)
+}
+
+func TestNewClientRejectsCrossOriginRedirect(t *testing.T) {
+	var targetRequests atomic.Int32
+	target := httptest.NewTLSServer(http.HandlerFunc(func(
+		w http.ResponseWriter,
+		_ *http.Request,
+	) {
+		targetRequests.Add(1)
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	t.Cleanup(target.Close)
+
+	origin := httptest.NewTLSServer(http.HandlerFunc(func(
+		w http.ResponseWriter,
+		r *http.Request,
+	) {
+		http.Redirect(w, r, target.URL, http.StatusTemporaryRedirect)
+	}))
+	t.Cleanup(origin.Close)
+
+	client, err := NewClient(
+		origin.URL,
+		origin.Client(),
+		http.Header{"X-Access-Token": {"secret"}},
+	)
+	require.NoError(t, err)
+
+	_, err = client.List(context.Background())
+	require.ErrorContains(t, err, "outside configured origin")
+	assert.Zero(t, targetRequests.Load())
+}
+
+func TestOriginBoundRoundTripperRejectsSchemeAndPortChanges(t *testing.T) {
+	transportCalls := 0
+	transport := roundTripperFunc(func(
+		*http.Request,
+	) (*http.Response, error) {
+		transportCalls++
+		return &http.Response{StatusCode: http.StatusNoContent}, nil
+	})
+	boundTransport := &originBoundRoundTripper{
+		base: transport,
+		origin: originFromURL(mustParseURL(
+			t, "https://envd.example",
+		)),
+	}
+
+	for _, tt := range []struct {
+		name      string
+		targetURL string
+	}{
+		{
+			name:      "HTTPSDowngrade",
+			targetURL: "http://envd.example:443/process.Process/Start",
+		},
+		{
+			name:      "DifferentPort",
+			targetURL: "https://envd.example:49983/process.Process/Start",
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			_, err := boundTransport.RoundTrip(&http.Request{
+				URL: mustParseURL(t, tt.targetURL),
+			})
+			require.ErrorContains(t, err, "outside configured origin")
+		})
+	}
+	assert.Zero(t, transportCalls)
+}
+
+func TestOriginBoundHTTPClientRetainsDefaultRedirectPolicy(t *testing.T) {
+	client := newOriginBoundHTTPClient(
+		&http.Client{}, mustParseURL(t, "https://envd.example"),
+	)
+	assert.Nil(t, client.CheckRedirect)
+}
+
+type roundTripperFunc func(*http.Request) (*http.Response, error)
+
+func (f roundTripperFunc) RoundTrip(
+	req *http.Request,
+) (*http.Response, error) {
+	return f(req)
+}
+
+func mustParseURL(t *testing.T, rawURL string) *url.URL {
+	t.Helper()
+	u, err := url.Parse(rawURL)
+	require.NoError(t, err)
+	return u
+}

codeexecutor/e2b/internal/envdprocess/operations.go

@@ -0,0 +1,30 @@
+//
+// Tencent is pleased to support the open source community by making trpc-agent-go available.
+//
+// Copyright (C) 2025 Tencent.  All rights reserved.
+//
+// trpc-agent-go is licensed under the Apache License Version 2.0.
+//
+
+package envdprocess
+
+// ProcessInfo describes a process reported by envd.
+type ProcessInfo struct {
+	PID  uint32
+	Tag  string
+	Cmd  string
+	Args []string
+	Envs map[string]string
+	Cwd  string
+}
+
+func cloneStrings(input map[string]string) map[string]string {
+	if input == nil {
+		return nil
+	}
+	output := make(map[string]string, len(input))
+	for key, value := range input {
+		output[key] = value
+	}
+	return output
+}

codeexecutor/e2b/internal/envdprocess/operations_test.go

@@ -0,0 +1,129 @@
+//
+// Tencent is pleased to support the open source community by making trpc-agent-go available.
+//
+// Copyright (C) 2025 Tencent.  All rights reserved.
+//
+// trpc-agent-go is licensed under the Apache License Version 2.0.
+//
+
+package envdprocess
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"testing"
+
+	"connectrpc.com/connect"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	processrpc "trpc.group/trpc-go/trpc-agent-go/codeexecutor/e2b/internal/envdprocess/spec"
+)
+
+func TestClientListAndKill(t *testing.T) {
+	handler := &testProcessHandler{}
+	handler.list = func(
+		_ context.Context,
+		req *connect.Request[processrpc.ListRequest],
+	) (*connect.Response[processrpc.ListResponse], error) {
+		assert.Equal(t, "token", req.Header().Get("X-Access-Token"))
+		cwd := "/tmp/list"
+		tag := "listed"
+		return connect.NewResponse(&processrpc.ListResponse{
+			Processes: []*processrpc.ProcessInfo{{
+				Pid: 203,
+				Tag: &tag,
+				Config: &processrpc.ProcessConfig{
+					Cmd:  "/bin/sh",
+					Args: []string{"-c", "sleep 60"},
+					Envs: map[string]string{"MODE": "test"},
+					Cwd:  &cwd,
+				},
+			}},
+		}), nil
+	}
+	handler.sendSignal = func(
+		_ context.Context,
+		req *connect.Request[processrpc.SendSignalRequest],
+	) (*connect.Response[processrpc.SendSignalResponse], error) {
+		assert.Equal(t, uint32(203), req.Msg.Process.GetPid())
+		return connect.NewResponse(&processrpc.SendSignalResponse{}), nil
+	}
+
+	client := newTestClient(t, handler, http.Header{
+		"X-Access-Token": []string{"token"},
+	})
+	processes, err := client.List(context.Background())
+	require.NoError(t, err)
+	require.Len(t, processes, 1)
+	assert.Equal(t, ProcessInfo{
+		PID:  203,
+		Tag:  "listed",
+		Cmd:  "/bin/sh",
+		Args: []string{"-c", "sleep 60"},
+		Envs: map[string]string{"MODE": "test"},
+		Cwd:  "/tmp/list",
+	}, processes[0])
+
+	killed, err := client.Kill(context.Background(), 203)
+	require.NoError(t, err)
+	assert.True(t, killed)
+}
+
+func TestClientKillMissingProcessReturnsFalse(t *testing.T) {
+	handler := &testProcessHandler{}
+	handler.sendSignal = func(
+		context.Context,
+		*connect.Request[processrpc.SendSignalRequest],
+	) (*connect.Response[processrpc.SendSignalResponse], error) {
+		return nil, connect.NewError(
+			connect.CodeNotFound, errors.New("process is not running"),
+		)
+	}
+	client := newTestClient(t, handler, nil)
+
+	killed, err := client.Kill(context.Background(), 404)
+	require.NoError(t, err)
+	assert.False(t, killed)
+}
+
+func TestClientOperationsRejectInvalidInput(t *testing.T) {
+	client := newTestClient(t, &testProcessHandler{}, nil)
+
+	_, err := client.Connect(nil, 1)
+	require.ErrorContains(t, err, "nil context")
+
+	_, err = client.Kill(context.Background(), 0)
+	require.ErrorContains(t, err, "pid is zero")
+
+	err = client.SendInput(context.Background(), 0, []byte("input"))
+	require.ErrorContains(t, err, "pid is zero")
+
+	err = client.CloseStdin(context.Background(), 0)
+	require.ErrorContains(t, err, "pid is zero")
+
+	var uninitialized *Client
+	_, err = uninitialized.Kill(context.Background(), 1)
+	require.ErrorContains(t, err, "client is not initialized")
+}
+
+func TestClientCloseStdinRejectsUnsupportedEnvd(t *testing.T) {
+	handler := &testProcessHandler{}
+	handler.closeStdin = func(
+		context.Context,
+		*connect.Request[processrpc.CloseStdinRequest],
+	) (*connect.Response[processrpc.CloseStdinResponse], error) {
+		t.Fatal("CloseStdin RPC must not be sent to unsupported envd")
+		return nil, nil
+	}
+	client := newTestClient(
+		t,
+		handler,
+		nil,
+		WithEnvdVersion("0.2.10"),
+	)
+
+	err := client.CloseStdin(context.Background(), 1)
+	require.ErrorContains(t, err, "close stdin requires envd >= 0.5.2")
+}