Skip to content

Commit 2b528e1

Browse files
sitolesitole
committed
Support traffic access token on orch proxy, integration tests
1 parent ab78f40 commit 2b528e1

File tree

4 files changed

+184
-3
lines changed

4 files changed

+184
-3
lines changed

‎packages/orchestrator/internal/proxy/proxy.go‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"go.uber.org/zap"
1212

1313
"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox"
14+
"github.com/e2b-dev/infra/packages/shared/pkg/consts"
1415
"github.com/e2b-dev/infra/packages/shared/pkg/env"
1516
"github.com/e2b-dev/infra/packages/shared/pkg/logger"
1617
reverseproxy "github.com/e2b-dev/infra/packages/shared/pkg/proxy"
@@ -23,6 +24,8 @@ const (
2324
// Also it's a good practice to set it to higher values as you progress in the stack
2425
// https://cloud.google.com/load-balancing/docs/https#timeouts_and_retries%23:~:text=The%20load%20balancer%27s%20backend%20keepalive,is%20greater%20than%20600%20seconds
2526
idleTimeout = 620 * time.Second
27+
28+
accessTokenHeader = "x-e2b-traffic-access-token"
2629
)
2730

2831
type SandboxProxy struct {
@@ -48,6 +51,22 @@ func NewSandboxProxy(meterProvider metric.MeterProvider, port uint16, sandboxes
4851
return nil, reverseproxy.NewErrSandboxNotFound(sandboxId)
4952
}
5053

54+
var accessToken *string = nil
55+
if net := sbx.Config.Network; net != nil && net.GetIngress() != nil {
56+
accessToken = net.GetIngress().TrafficAccessToken
57+
}
58+
59+
// Handle traffic access token validation.
60+
// We are skipping envd port as it has its own access validation mechanism.
61+
if accessToken != nil && int64(port) != consts.DefaultEnvdServerPort {
62+
accessTokenRaw := r.Header.Get(accessTokenHeader)
63+
if accessTokenRaw == "" {
64+
return nil, reverseproxy.ErrMissingTrafficAccessToken
65+
} else if accessTokenRaw != *accessToken {
66+
return nil, reverseproxy.ErrInvalidTrafficAccessToken
67+
}
68+
}
69+
5170
url := &url.URL{
5271
Scheme: "http",
5372
Host: fmt.Sprintf("%s:%d", sbx.Slot.HostIPString(), port),

‎packages/shared/pkg/proxy/errors.go‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,14 @@
11
package proxy
22

3-
import "errors"
4-
5-
var ErrInvalidHost = errors.New("invalid url host")
3+
import (
4+
"errors"
5+
)
6+
7+
var (
8+
ErrInvalidHost = errors.New("invalid url host")
9+
ErrInvalidTrafficAccessToken = errors.New("invalid traffic access token")
10+
ErrMissingTrafficAccessToken = errors.New("missing traffic access token")
11+
)
612

713
type InvalidSandboxPortError struct {
814
Port string

‎packages/shared/pkg/proxy/handler.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,20 @@ func handler(p *pool.ProxyPool, getDestination func(r *http.Request) (*pool.Dest
5858
return
5959
}
6060

61+
if errors.Is(err, ErrMissingTrafficAccessToken) {
62+
zap.L().Warn("traffic access token header is missing", zap.String("host", r.Host))
63+
http.Error(w, "Sandbox is secured with traffic access token. Access token header is missing", http.StatusForbidden)
64+
65+
return
66+
}
67+
68+
if errors.Is(err, ErrInvalidTrafficAccessToken) {
69+
zap.L().Warn("traffic access token is invalid", zap.String("host", r.Host))
70+
http.Error(w, "Sandbox is secured with traffic access token. Provided access token is invalid.", http.StatusForbidden)
71+
72+
return
73+
}
74+
6175
if err != nil {
6276
zap.L().Error("failed to route request", zap.Error(err), zap.String("host", r.Host))
6377
http.Error(w, fmt.Sprintf("Unexpected error when routing request: %s", err), http.StatusInternalServerError)

‎tests/integration/internal/tests/proxies/traffic_access_token_test.go‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
package api
2+
3+
import (
4+
"encoding/json"
5+
"io"
6+
"net/http"
7+
"net/url"
8+
"testing"
9+
"time"
10+
11+
"github.com/stretchr/testify/assert"
12+
"github.com/stretchr/testify/require"
13+
14+
"github.com/e2b-dev/infra/packages/shared/pkg/consts"
15+
"github.com/e2b-dev/infra/tests/integration/internal/api"
16+
"github.com/e2b-dev/infra/tests/integration/internal/setup"
17+
"github.com/e2b-dev/infra/tests/integration/internal/utils"
18+
)
19+
20+
func TestSandboxWithEnabledTrafficAccessTokenButMissingHeader(t *testing.T) {
21+
c := setup.GetAPIClient()
22+
23+
sbxNetAllowPublic := false
24+
sbxNet := &api.SandboxNetworkConfig{
25+
AllowPublicAccess: &sbxNetAllowPublic,
26+
}
27+
sbx := utils.SetupSandboxWithCleanup(t, c, utils.WithNetwork(sbxNet))
28+
require.NotNil(t, sbx.TrafficAccessToken)
29+
30+
url, err := url.Parse(setup.EnvdProxy)
31+
require.NoError(t, err)
32+
33+
client := &http.Client{
34+
Timeout: 1000 * time.Second,
35+
}
36+
37+
resp := waitForStatus(t, client, sbx, url, 8080, nil, http.StatusForbidden)
38+
require.NoError(t, err)
39+
require.NotNil(t, resp)
40+
require.Equal(t, http.StatusForbidden, resp.StatusCode)
41+
42+
body, err := io.ReadAll(resp.Body)
43+
require.NoError(t, err)
44+
require.NoError(t, resp.Body.Close())
45+
46+
assert.Equal(t, "text/plain; charset=utf-8", resp.Header.Get("Content-Type"))
47+
assert.Equal(t, "Sandbox is secured with traffic access token. Access token header is missing\n", string(body))
48+
}
49+
50+
func TestSandboxWithEnabledTrafficAccessTokenButInvalidHeader(t *testing.T) {
51+
c := setup.GetAPIClient()
52+
53+
sbxNetAllowPublic := false
54+
sbxNet := &api.SandboxNetworkConfig{
55+
AllowPublicAccess: &sbxNetAllowPublic,
56+
}
57+
sbx := utils.SetupSandboxWithCleanup(t, c, utils.WithNetwork(sbxNet))
58+
require.NotNil(t, sbx.TrafficAccessToken)
59+
url, err := url.Parse(setup.EnvdProxy)
60+
require.NoError(t, err)
61+
62+
client := &http.Client{
63+
Timeout: 1000 * time.Second,
64+
}
65+
66+
headers := &http.Header{"x-e2b-traffic-access-token": []string{"abcd"}}
67+
resp := waitForStatus(t, client, sbx, url, 8080, headers, http.StatusForbidden)
68+
require.NoError(t, err)
69+
require.NotNil(t, resp)
70+
require.Equal(t, http.StatusForbidden, resp.StatusCode)
71+
72+
body, err := io.ReadAll(resp.Body)
73+
require.NoError(t, err)
74+
require.NoError(t, resp.Body.Close())
75+
76+
assert.Equal(t, "text/plain; charset=utf-8", resp.Header.Get("Content-Type"))
77+
assert.Equal(t, "Sandbox is secured with traffic access token. Provided access token is invalid.\n", string(body))
78+
}
79+
80+
func TestSandboxWithEnabledTrafficAccessToken(t *testing.T) {
81+
c := setup.GetAPIClient()
82+
83+
sbxNetAllowPublic := false
84+
sbxNet := &api.SandboxNetworkConfig{
85+
AllowPublicAccess: &sbxNetAllowPublic,
86+
}
87+
sbx := utils.SetupSandboxWithCleanup(t, c, utils.WithNetwork(sbxNet))
88+
require.NotNil(t, sbx.TrafficAccessToken)
89+
90+
url, err := url.Parse(setup.EnvdProxy)
91+
require.NoError(t, err)
92+
93+
client := &http.Client{
94+
Timeout: 1000 * time.Second,
95+
}
96+
97+
port := 8080
98+
headers := &http.Header{"x-e2b-traffic-access-token": []string{*sbx.TrafficAccessToken}}
99+
resp := waitForStatus(t, client, sbx, url, port, headers, http.StatusBadGateway)
100+
require.NoError(t, err)
101+
require.NotNil(t, resp)
102+
require.Equal(t, http.StatusBadGateway, resp.StatusCode)
103+
104+
assert.Equal(t, "application/json; charset=utf-8", resp.Header.Get("Content-Type"))
105+
106+
// Parse error response
107+
var errorResp struct {
108+
Message string `json:"message"`
109+
SandboxID string `json:"sandboxId"`
110+
Port int `json:"port"`
111+
}
112+
err = json.NewDecoder(resp.Body).Decode(&errorResp)
113+
require.NoError(t, err)
114+
require.NoError(t, resp.Body.Close())
115+
assert.Equal(t, "The sandbox is running but port is not open", errorResp.Message)
116+
assert.Equal(t, sbx.SandboxID, errorResp.SandboxID)
117+
assert.Equal(t, port, errorResp.Port)
118+
}
119+
120+
func TestEnvdPortIsNotAffectedByTrafficAccessToken(t *testing.T) {
121+
c := setup.GetAPIClient()
122+
123+
sbxNetAllowPublic := false
124+
sbxNet := &api.SandboxNetworkConfig{
125+
AllowPublicAccess: &sbxNetAllowPublic,
126+
}
127+
sbx := utils.SetupSandboxWithCleanup(t, c, utils.WithNetwork(sbxNet))
128+
require.NotNil(t, sbx.TrafficAccessToken)
129+
130+
url, err := url.Parse(setup.EnvdProxy)
131+
require.NoError(t, err)
132+
133+
client := &http.Client{
134+
Timeout: 1000 * time.Second,
135+
}
136+
137+
resp := waitForStatus(t, client, sbx, url, int(consts.DefaultEnvdServerPort), nil, http.StatusNotFound)
138+
require.NoError(t, err)
139+
require.NotNil(t, resp)
140+
require.NoError(t, resp.Body.Close())
141+
require.Equal(t, http.StatusNotFound, resp.StatusCode)
142+
}

0 commit comments

Comments
 (0)