[Feature] support custom e2b

pkg/proxy/ext_proc.go

@@ -91,9 +91,14 @@ func (s *Server) handleRequestHeaders(requestHeaders *extProcPb.ProcessingReques
 	sandboxID, sandboxPort, extraHeaders, err := s.adapter.Map(scheme, authority, path, port, headers)
 	if err != nil {
 		// Return error response instead of gRPC error
+		log.Error(err, "failed to map request to sandbox")
 		errorMsg := fmt.Sprintf("failed to map request to sandbox, URL=%s://%s%s", scheme, authority, path)
 		return s.logAndCreateErrorResponse(http.StatusInternalServerError, errorMsg, log)
 	}
+	if sandboxPort < 0 || sandboxPort > 65535 {
+		errorMsg := fmt.Sprintf("invalid sandbox port: %d", sandboxPort)
+		return s.logAndCreateErrorResponse(http.StatusBadRequest, errorMsg, log)
+	}
 	log.Info("request mapped", "sandboxID", sandboxID, "sandboxPort", sandboxPort, "extraHeaders", extraHeaders)
 
 	route, ok := s.LoadRoute(sandboxID)
@@ -111,7 +116,6 @@ func (s *Server) handleRequestHeaders(requestHeaders *extProcPb.ProcessingReques
 		extraHeaders[k] = v
 	}
 	extraHeaders[OrigDstHeader] = fmt.Sprintf("%s:%d", route.IP, sandboxPort)
-
 	return s.logAndCreateDstResponse(requestHeaders.RequestHeaders, extraHeaders, log)
 }
 

pkg/proxy/server.go

@@ -96,8 +96,8 @@ func (s *Server) Run() error {
 
 	// HTTP
 	mux := http.NewServeMux()
-	web.RegisterRoute(mux, fmt.Sprintf("%s %s", http.MethodPost, RefreshAPI), s.handleRefresh)
-	web.RegisterRoute(mux, fmt.Sprintf("%s %s", http.MethodGet, HelloAPI), s.handleHello)
+	web.RegisterRoute(mux, http.MethodPost, RefreshAPI, s.handleRefresh)
+	web.RegisterRoute(mux, http.MethodGet, HelloAPI, s.handleHello)
 	s.httpSrv = &http.Server{
 		Addr:              fmt.Sprintf(":%d", SystemPort),
 		Handler:           mux,

pkg/servers/e2b/adapters/customized_e2b.go

@@ -0,0 +1,39 @@
+package adapters
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+type CustomizedE2BAdapter struct{}
+
+const CustomPrefix = "/kruise"
+
+// Map maps paths like /kruise/sandbox1234/3000/xxx to sandboxID=sandbox1234 and port=3000
+func (a *CustomizedE2BAdapter) Map(_, _, path string, _ int, _ map[string]string) (
+	sandboxID string, sandboxPort int, extraHeaders map[string]string, err error) {
+	if len(path) <= len(CustomPrefix)+1 {
+		err = fmt.Errorf("invalid path: %s", path)
+		return
+	}
+	path = path[len(CustomPrefix)+1:] // remove prefix "/kruise/"
+	split := strings.SplitN(path, "/", 3)
+	if len(split) < 3 {
+		err = fmt.Errorf("invalid path: %s", path)
+		return
+	}
+	sandboxID = split[0]
+	sandboxPort, err = strconv.Atoi(split[1])
+	if err != nil {
+		return
+	}
+	extraHeaders = map[string]string{
+		":path": "/" + split[2],
+	}
+	return
+}
+
+func (a *CustomizedE2BAdapter) IsSandboxRequest(_, path string, _ int) bool {
+	return !strings.HasPrefix(path, CustomPrefix+"/api")
+}

pkg/servers/e2b/adapters/common.go → pkg/servers/e2b/adapters/native_e2b.go

@@ -7,13 +7,12 @@ import (
 	"strings"
 )
 
-type CommonAdapter struct {
-	Port int
-}
+type NativeE2BAdapter struct{}
 
 var hostRegex = regexp.MustCompile(`^(\d+)-([a-zA-Z0-9\-]+)\.`)
 
-func (a *CommonAdapter) Map(_, authority, _ string, _ int, _ map[string]string) (
+// Map maps authorities like 3000-sandbox1234.example.com to sandboxID=sandbox1234 and port=3000
+func (a *NativeE2BAdapter) Map(_, authority, _ string, _ int, _ map[string]string) (
 	sandboxID string, sandboxPort int, extraHeaders map[string]string, err error) {
 	matches := hostRegex.FindStringSubmatch(authority)
 	if len(matches) != 3 {
@@ -24,17 +23,12 @@ func (a *CommonAdapter) Map(_, authority, _ string, _ int, _ map[string]string)
 	// Extract port number and sandboxID
 	sandboxPort, err = strconv.Atoi(matches[1])
 	if err != nil {
-		return
+		return // impossible
 	}
 	sandboxID = matches[2]
-
-	return sandboxID, sandboxPort, extraHeaders, err
+	return
 }
 
-func (a *CommonAdapter) IsSandboxRequest(authority, _ string, _ int) bool {
+func (a *NativeE2BAdapter) IsSandboxRequest(authority, _ string, _ int) bool {
 	return !strings.HasPrefix(authority, "api.")
 }
-
-func (a *CommonAdapter) Entry() string {
-	return fmt.Sprintf("127.0.0.1:%d", a.Port)
-}

pkg/servers/e2b/adapters/unified_e2b.go

@@ -0,0 +1,47 @@
+package adapters
+
+import (
+	"fmt"
+	"strings"
+)
+
+// E2BMapper is part of proxy.RequestAdapter
+type E2BMapper interface {
+	Map(scheme, authority, path string, port int, headers map[string]string) (
+		sandboxID string, sandboxPort int, extraHeaders map[string]string, err error)
+	IsSandboxRequest(authority, path string, port int) bool
+}
+
+type E2BAdapter struct {
+	Port       int
+	native     *NativeE2BAdapter
+	customized *CustomizedE2BAdapter
+}
+
+func NewE2BAdapter(port int) *E2BAdapter {
+	return &E2BAdapter{
+		Port:       port,
+		native:     &NativeE2BAdapter{},
+		customized: &CustomizedE2BAdapter{},
+	}
+}
+
+func (a *E2BAdapter) Map(scheme, authority, path string, port int, headers map[string]string) (
+	sandboxID string, sandboxPort int, extraHeaders map[string]string, err error) {
+	return a.ChooseAdapter(path).Map(scheme, authority, path, port, headers)
+}
+
+func (a *E2BAdapter) IsSandboxRequest(authority, path string, port int) bool {
+	return a.ChooseAdapter(path).IsSandboxRequest(authority, path, port)
+}
+
+func (a *E2BAdapter) Entry() string {
+	return fmt.Sprintf("127.0.0.1:%d", a.Port)
+}
+
+func (a *E2BAdapter) ChooseAdapter(path string) E2BMapper {
+	if strings.HasPrefix(path, CustomPrefix) {
+		return a.customized
+	}
+	return a.native
+}

pkg/servers/e2b/adapters/unified_e2b_test.go

@@ -0,0 +1,222 @@
+package adapters
+
+import (
+	"context"
+	"testing"
+
+	"github.com/openkruise/agents/pkg/proxy"
+	"github.com/openkruise/agents/pkg/sandbox-manager/clients"
+	"github.com/openkruise/agents/pkg/servers/e2b/keys"
+	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+var adminKey = "admin-key"
+
+func SetUpE2BAdapter(t *testing.T) proxy.RequestAdapter {
+	client := clients.NewFakeClientSet()
+	_, err := client.K8sClient.CoreV1().Secrets("default").Create(context.Background(), &corev1.Secret{
+		ObjectMeta: metav1.ObjectMeta{
+			Name: keys.KeySecretName,
+		},
+		Data: map[string][]byte{},
+	}, metav1.CreateOptions{})
+	assert.NoError(t, err)
+	keyStore := &keys.SecretKeyStorage{
+		Namespace: "default",
+		AdminKey:  adminKey,
+		Client:    client.K8sClient,
+		Stop:      make(chan struct{}),
+	}
+	assert.NoError(t, keyStore.Init(context.Background()))
+	return NewE2BAdapter(8080)
+}
+
+func TestMap(t *testing.T) {
+	tests := []struct {
+		name      string
+		authority string
+		path      string
+		headers   map[string]string
+
+		expectErr         bool
+		expectSandboxID   string
+		expectSandboxPort int
+		expectHeaders     map[string]string
+	}{
+		{
+			name:              "native e2b adapter - valid authority with CDP port",
+			authority:         "9222-sandbox1234.example.com",
+			path:              "/",
+			headers:           map[string]string{},
+			expectSandboxID:   "sandbox1234",
+			expectSandboxPort: 9222,
+			expectErr:         false,
+		},
+		{
+			name:              "native e2b adapter - valid authority with regular port and valid token",
+			authority:         "3000-sandbox5678.example.com",
+			path:              "/",
+			headers:           map[string]string{"x-access-token": adminKey},
+			expectSandboxID:   "sandbox5678",
+			expectSandboxPort: 3000,
+			expectErr:         false,
+		},
+		{
+			name:              "native e2b adapter - valid authority with regular port and invalid token",
+			authority:         "3000-sandbox5678.example.com",
+			path:              "/",
+			headers:           map[string]string{"x-access-token": "invalid-token"},
+			expectSandboxID:   "sandbox5678",
+			expectSandboxPort: 3000,
+			expectErr:         false,
+		},
+		{
+			name:      "native e2b adapter - invalid authority format",
+			authority: "invalid-authority",
+			path:      "/",
+			headers:   map[string]string{},
+			expectErr: true,
+		},
+		{
+			name:              "customized e2b adapter - valid path with CDP port",
+			authority:         "",
+			path:              "/kruise/sandbox1234/9222/some/path",
+			headers:           map[string]string{},
+			expectSandboxID:   "sandbox1234",
+			expectSandboxPort: 9222,
+			expectHeaders:     map[string]string{":path": "/some/path"},
+			expectErr:         false,
+		},
+		{
+			name:              "customized e2b adapter - valid path with regular port and valid token",
+			authority:         "",
+			path:              "/kruise/sandbox1234/3000/some/path",
+			headers:           map[string]string{"x-access-token": adminKey},
+			expectSandboxID:   "sandbox1234",
+			expectSandboxPort: 3000,
+			expectHeaders:     map[string]string{":path": "/some/path"},
+			expectErr:         false,
+		},
+		{
+			name:              "customized e2b adapter - valid path with regular port and invalid token",
+			authority:         "",
+			path:              "/kruise/sandbox1234/3000/some/path",
+			headers:           map[string]string{"x-access-token": "invalid-token"},
+			expectSandboxID:   "sandbox1234",
+			expectSandboxPort: 3000,
+			expectHeaders:     map[string]string{":path": "/some/path"},
+			expectErr:         false,
+		},
+		{
+			name:      "customized e2b adapter - invalid path (too short)",
+			authority: "",
+			path:      "/kruise/",
+			headers:   map[string]string{},
+			expectErr: true,
+		},
+		{
+			name:      "customized e2b adapter - invalid path (missing components)",
+			authority: "",
+			path:      "/kruise/sandbox1234",
+			headers:   map[string]string{},
+			expectErr: true,
+		},
+		{
+			name:      "customized e2b adapter - invalid port number",
+			authority: "",
+			path:      "/kruise/sandbox1234/invalid-port/some/path",
+			headers:   map[string]string{},
+			expectErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			adapter := SetUpE2BAdapter(t)
+			sandboxID, sandboxPort, headers, err := adapter.Map("http", tt.authority, tt.path, 8080, tt.headers)
+			if tt.expectErr {
+				assert.Error(t, err)
+				return
+			}
+			assert.NoError(t, err)
+			if err == nil {
+				assert.Equal(t, tt.expectSandboxID, sandboxID)
+				assert.Equal(t, tt.expectSandboxPort, sandboxPort)
+				if tt.expectHeaders == nil {
+					assert.Nil(t, headers)
+				} else {
+					assert.NotNil(t, headers)
+					assert.Equal(t, len(tt.expectHeaders), len(headers))
+					for k, v := range tt.expectHeaders {
+						assert.Equal(t, v, headers[k])
+					}
+				}
+			}
+		})
+	}
+}
+
+func TestIsSandboxRequest(t *testing.T) {
+	tests := []struct {
+		name      string
+		authority string
+		path      string
+
+		expectIsSandboxRequest bool
+	}{
+		// Native E2B Adapter Tests
+		{
+			name:                   "native e2b adapter - regular sandbox request",
+			authority:              "3000-sandbox1234.example.com",
+			path:                   "/",
+			expectIsSandboxRequest: true,
+		},
+		{
+			name:                   "native e2b adapter - api request",
+			authority:              "api.example.com",
+			path:                   "/",
+			expectIsSandboxRequest: false,
+		},
+		{
+			name:                   "native e2b adapter - api subdomain request",
+			authority:              "api.something.example.com",
+			path:                   "/",
+			expectIsSandboxRequest: false,
+		},
+
+		// Customized E2B Adapter Tests
+		{
+			name:                   "customized e2b adapter - regular sandbox request",
+			authority:              "",
+			path:                   "/kruise/sandbox1234/3000/some/path",
+			expectIsSandboxRequest: true,
+		},
+		{
+			name:                   "customized e2b adapter - api request",
+			authority:              "",
+			path:                   "/kruise/api/something",
+			expectIsSandboxRequest: false,
+		},
+		{
+			name:                   "customized e2b adapter - api-like path but not under /kruise/api",
+			authority:              "",
+			path:                   "/kruise/apiserver/something",
+			expectIsSandboxRequest: false,
+		},
+		{
+			name:                   "customized e2b adapter - non-api request under kruise prefix",
+			authority:              "",
+			path:                   "/kruise/some/other/path",
+			expectIsSandboxRequest: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			adapter := SetUpE2BAdapter(t)
+			isSandbox := adapter.IsSandboxRequest(tt.authority, tt.path, 8080)
+			assert.Equal(t, tt.expectIsSandboxRequest, isSandbox)
+		})
+	}
+}