From 2299f39ab960439d33955435516afd5ef34d0a0d Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Thu, 12 Mar 2026 13:54:55 +0100
Subject: [PATCH 1/9] chore(api): add auth attributes

---
 packages/api/internal/handlers/store.go       |  8 +--
 packages/auth/pkg/auth/service.go             | 10 +++-
 packages/client-proxy/internal/proxy/proxy.go | 11 +---
 packages/orchestrator/internal/proxy/proxy.go | 17 ++----
 packages/shared/pkg/logger/fields.go          | 60 +++++++++++++++++++
 packages/shared/pkg/middleware/logging.go     |  4 +-
 packages/shared/pkg/telemetry/fields.go       |  8 +++
 7 files changed, 89 insertions(+), 29 deletions(-)

diff --git a/packages/api/internal/handlers/store.go b/packages/api/internal/handlers/store.go
index 171a756ce9..a3f35d5437 100644
--- a/packages/api/internal/handlers/store.go
+++ b/packages/api/internal/handlers/store.go
@@ -272,18 +272,18 @@ func (a *APIStore) GetHealth(c *gin.Context) {
 	c.String(http.StatusServiceUnavailable, "Service is unavailable")
 }
 
-func (a *APIStore) GetTeamFromAPIKey(ctx context.Context, _ *gin.Context, apiKey string) (*types.Team, *api.APIError) {
+func (a *APIStore) GetTeamFromAPIKey(ctx context.Context, ginCtx *gin.Context, apiKey string) (*types.Team, *api.APIError) {
 	ctx, span := tracer.Start(ctx, "get team from api key")
 	defer span.End()
 
-	return a.authService.ValidateAPIKey(ctx, apiKey)
+	return a.authService.ValidateAPIKey(ctx, ginCtx, apiKey)
 }
 
-func (a *APIStore) GetUserFromAccessToken(ctx context.Context, _ *gin.Context, accessToken string) (uuid.UUID, *api.APIError) {
+func (a *APIStore) GetUserFromAccessToken(ctx context.Context, ginCtx *gin.Context, accessToken string) (uuid.UUID, *api.APIError) {
 	ctx, span := tracer.Start(ctx, "get user from access token")
 	defer span.End()
 
-	return a.authService.ValidateAccessToken(ctx, accessToken)
+	return a.authService.ValidateAccessToken(ctx, ginCtx, accessToken)
 }
 
 func (a *APIStore) GetUserIDFromSupabaseToken(ctx context.Context, _ *gin.Context, supabaseToken string) (uuid.UUID, *api.APIError) {
diff --git a/packages/auth/pkg/auth/service.go b/packages/auth/pkg/auth/service.go
index 98994c4515..be68574294 100644
--- a/packages/auth/pkg/auth/service.go
+++ b/packages/auth/pkg/auth/service.go
@@ -10,6 +10,7 @@ import (
 	"github.com/google/uuid"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/keys"
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
 
 // AuthStore abstracts the DB operations needed for auth validation.
@@ -36,7 +37,7 @@ func NewAuthService[T any](store AuthStore[T], teamCache *AuthCache[T], jwtSecre
 }
 
 // ValidateAPIKey verifies the API key format and fetches the associated team via cache + store.
-func (s *AuthService[T]) ValidateAPIKey(ctx context.Context, apiKey string) (T, *APIError) {
+func (s *AuthService[T]) ValidateAPIKey(ctx context.Context, ginCtx *gin.Context, apiKey string) (T, *APIError) {
 	hashedKey, err := keys.VerifyKey(keys.ApiKeyPrefix, apiKey)
 	if err != nil {
 		var zero T
@@ -79,11 +80,13 @@ func (s *AuthService[T]) ValidateAPIKey(ctx context.Context, apiKey string) (T,
 		}
 	}
 
+	telemetry.SetAttributes(ginCtx.Request.Context(), telemetry.WithAPIKey(keys.ApiKeyPrefix, apiKey))
+
 	return result, nil
 }
 
 // ValidateAccessToken verifies the access token format and fetches the associated user ID.
-func (s *AuthService[T]) ValidateAccessToken(ctx context.Context, accessToken string) (uuid.UUID, *APIError) {
+func (s *AuthService[T]) ValidateAccessToken(ctx context.Context, ginCtx *gin.Context, accessToken string) (uuid.UUID, *APIError) {
 	hashedToken, err := keys.VerifyKey(keys.AccessTokenPrefix, accessToken)
 	if err != nil {
 		return uuid.UUID{}, &APIError{
@@ -102,6 +105,9 @@ func (s *AuthService[T]) ValidateAccessToken(ctx context.Context, accessToken st
 		}
 	}
 
+	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
+	telemetry.SetAttributes(ginCtx.Request.Context(), telemetry.WithAccessToken(keys.AccessTokenPrefix, accessToken))
+
 	return userID, nil
 }
 
diff --git a/packages/client-proxy/internal/proxy/proxy.go b/packages/client-proxy/internal/proxy/proxy.go
index 80d38df291..451f14d5ef 100644
--- a/packages/client-proxy/internal/proxy/proxy.go
+++ b/packages/client-proxy/internal/proxy/proxy.go
@@ -121,16 +121,7 @@ func NewClientProxy(meterProvider metric.MeterProvider, serviceName string, port
 				return nil, err
 			}
 
-			l := logger.L().With(
-				zap.String("origin_host", r.Host),
-				logger.WithSandboxID(sandboxId),
-				zap.Uint64("sandbox_req_port", port),
-				zap.String("sandbox_req_path", r.URL.Path),
-				zap.String("sandbox_req_method", r.Method),
-				zap.String("sandbox_req_user_agent", r.UserAgent()),
-				zap.String("remote_addr", r.RemoteAddr),
-				zap.Int64("content_length", r.ContentLength),
-			)
+			l := logger.L().With(logger.ProxyRequestFields(r, sandboxId, port)...)
 
 			trafficAccessToken := r.Header.Get(proxygrpc.MetadataTrafficAccessToken)
 			envdAccessToken := r.Header.Get(proxygrpc.MetadataEnvdHTTPAccessToken)
diff --git a/packages/orchestrator/internal/proxy/proxy.go b/packages/orchestrator/internal/proxy/proxy.go
index ae6313eefe..2c4dbd930b 100644
--- a/packages/orchestrator/internal/proxy/proxy.go
+++ b/packages/orchestrator/internal/proxy/proxy.go
@@ -10,13 +10,13 @@ import (
 	"time"
 
 	"go.opentelemetry.io/otel/metric"
-	"go.uber.org/zap"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/shared/pkg/connlimit"
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/env"
 	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
+	"github.com/e2b-dev/infra/packages/shared/pkg/httputil"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	reverseproxy "github.com/e2b-dev/infra/packages/shared/pkg/proxy"
 	"github.com/e2b-dev/infra/packages/shared/pkg/proxy/pool"
@@ -103,16 +103,11 @@ func NewSandboxProxy(meterProvider metric.MeterProvider, port uint16, sandboxes
 			}
 
 			logger := logger.L().With(
-				zap.String("origin_host", r.Host),
-				logger.WithSandboxID(sbx.Runtime.SandboxID),
-				logger.WithTeamID(sbx.Runtime.TeamID),
-				logger.WithSandboxIP(sbx.Slot.HostIPString()),
-				zap.Uint64("sandbox_req_port", port),
-				zap.String("sandbox_req_path", r.URL.Path),
-				zap.String("sandbox_req_method", r.Method),
-				zap.String("sandbox_req_user_agent", r.UserAgent()),
-				zap.String("remote_addr", r.RemoteAddr),
-				zap.Int64("content_length", r.ContentLength),
+				append(
+					logger.ProxyRequestFields(r, sbx.Runtime.SandboxID, port),
+					logger.WithTeamID(sbx.Runtime.TeamID),
+					logger.WithSandboxIP(sbx.Slot.HostIPString()),
+				)...,
 			)
 
 			return &pool.Destination{
diff --git a/packages/shared/pkg/logger/fields.go b/packages/shared/pkg/logger/fields.go
index e79475f660..bc179c8418 100644
--- a/packages/shared/pkg/logger/fields.go
+++ b/packages/shared/pkg/logger/fields.go
@@ -1,6 +1,11 @@
 package logger
 
 import (
+	"fmt"
+	"net"
+	"net/http"
+	"strings"
+
 	"github.com/google/uuid"
 	"go.uber.org/zap"
 )
@@ -44,3 +49,58 @@ func WithSandboxIP(sandboxIP string) zap.Field {
 func WithEnvdVersion(envdVersion string) zap.Field {
 	return zap.String("envd.version", envdVersion)
 }
+
+func WithClientIP(clientIP string) zap.Field {
+	return zap.String("http.client_ip", clientIP)
+}
+
+func WithAPIKey(prefix, apiKey string) zap.Field {
+	return zap.String("auth.api_key", tokenHint(prefix, apiKey))
+}
+
+func WithAccessToken(prefix, accessToken string) zap.Field {
+	return zap.String("auth.access_token", tokenHint(prefix, accessToken))
+}
+
+// ProxyRequestFields returns the common logger fields for a proxied HTTP request.
+func ProxyRequestFields(r *http.Request, sandboxID string, sandboxPort uint64) []zap.Field {
+	return []zap.Field{
+		zap.String("origin_host", r.Host),
+		WithSandboxID(sandboxID),
+		zap.Uint64("sandbox_req_port", sandboxPort),
+		zap.String("sandbox_req_path", r.URL.Path),
+		zap.String("sandbox_req_method", r.Method),
+		zap.String("sandbox_req_user_agent", r.UserAgent()),
+		zap.String("remote_addr", r.RemoteAddr),
+		WithClientIP(clientIP(r)),
+		zap.Int64("content_length", r.ContentLength),
+	}
+}
+
+// tokenHint returns a masked version of the token value for debugging.
+// It strips the prefix and shows the first 2 and last 2 characters (e.g. "<prefix>ab...9f").
+func tokenHint(prefix string, token string) string {
+	value := strings.TrimPrefix(token, prefix)
+	if len(value) < 5 {
+		return fmt.Sprintf("%s...", prefix)
+	}
+
+	return fmt.Sprintf("%s%s...%s", prefix, value[:2], value[len(value)-2:])
+}
+
+// clientIP extracts the real client IP from the request.
+// It reads the first entry from X-Forwarded-For, falling back to RemoteAddr with the port stripped.
+func clientIP(r *http.Request) string {
+	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
+		if ip := strings.TrimSpace(strings.SplitN(xff, ",", 2)[0]); ip != "" {
+			return ip
+		}
+	}
+
+	host, _, err := net.SplitHostPort(r.RemoteAddr)
+	if err != nil {
+		return r.RemoteAddr
+	}
+
+	return host
+}
diff --git a/packages/shared/pkg/middleware/logging.go b/packages/shared/pkg/middleware/logging.go
index 97a1fcf0e3..4642f333c9 100644
--- a/packages/shared/pkg/middleware/logging.go
+++ b/packages/shared/pkg/middleware/logging.go
@@ -10,7 +10,7 @@ import (
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 
-	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	logpkg "github.com/e2b-dev/infra/packages/shared/pkg/logger"
 )
 
 // Based on https://github.com/gin-contrib/zap
@@ -32,7 +32,7 @@ type Config struct {
 	Skipper Skipper
 }
 
-func LoggingMiddleware(logger logger.Logger, conf Config) gin.HandlerFunc {
+func LoggingMiddleware(logger logpkg.Logger, conf Config) gin.HandlerFunc {
 	skipPaths := make(map[string]bool, len(conf.SkipPaths))
 	for _, path := range conf.SkipPaths {
 		skipPaths[path] = true
diff --git a/packages/shared/pkg/telemetry/fields.go b/packages/shared/pkg/telemetry/fields.go
index d4bdbcdb00..fc6fd6f1de 100644
--- a/packages/shared/pkg/telemetry/fields.go
+++ b/packages/shared/pkg/telemetry/fields.go
@@ -44,6 +44,14 @@ func WithEnvdVersion(envdVersion string) attribute.KeyValue {
 	return zapFieldToOTELAttribute(logger.WithEnvdVersion(envdVersion))
 }
 
+func WithAPIKey(prefix, apiKey string) attribute.KeyValue {
+	return zapFieldToOTELAttribute(logger.WithAPIKey(prefix, apiKey))
+}
+
+func WithAccessToken(prefix, accessToken string) attribute.KeyValue {
+	return zapFieldToOTELAttribute(logger.WithAccessToken(prefix, accessToken))
+}
+
 func zapFieldToOTELAttribute(f zap.Field) attribute.KeyValue {
 	e := &ZapFieldToOTELAttributeEncoder{}
 	f.AddTo(e)

From 5603579166a5f2537147851fb50d3426590730d9 Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Thu, 12 Mar 2026 13:57:34 +0100
Subject: [PATCH 2/9] fix: add missing nolnt

---
 packages/auth/pkg/auth/service.go         | 1 +
 packages/shared/pkg/middleware/logging.go | 4 ++--
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/packages/auth/pkg/auth/service.go b/packages/auth/pkg/auth/service.go
index be68574294..e90eb79ef8 100644
--- a/packages/auth/pkg/auth/service.go
+++ b/packages/auth/pkg/auth/service.go
@@ -80,6 +80,7 @@ func (s *AuthService[T]) ValidateAPIKey(ctx context.Context, ginCtx *gin.Context
 		}
 	}
 
+	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
 	telemetry.SetAttributes(ginCtx.Request.Context(), telemetry.WithAPIKey(keys.ApiKeyPrefix, apiKey))
 
 	return result, nil
diff --git a/packages/shared/pkg/middleware/logging.go b/packages/shared/pkg/middleware/logging.go
index 4642f333c9..97a1fcf0e3 100644
--- a/packages/shared/pkg/middleware/logging.go
+++ b/packages/shared/pkg/middleware/logging.go
@@ -10,7 +10,7 @@ import (
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 
-	logpkg "github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 )
 
 // Based on https://github.com/gin-contrib/zap
@@ -32,7 +32,7 @@ type Config struct {
 	Skipper Skipper
 }
 
-func LoggingMiddleware(logger logpkg.Logger, conf Config) gin.HandlerFunc {
+func LoggingMiddleware(logger logger.Logger, conf Config) gin.HandlerFunc {
 	skipPaths := make(map[string]bool, len(conf.SkipPaths))
 	for _, path := range conf.SkipPaths {
 		skipPaths[path] = true

From 303d0c0d74596a79f08e1f9c8d51ab1b673465e3 Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Thu, 12 Mar 2026 14:10:46 +0100
Subject: [PATCH 3/9] chore: reuse keys package

---
 packages/auth/pkg/auth/service.go       |  4 ++--
 packages/shared/pkg/logger/fields.go    | 22 +++++++++++-----------
 packages/shared/pkg/telemetry/fields.go |  8 ++++----
 3 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/packages/auth/pkg/auth/service.go b/packages/auth/pkg/auth/service.go
index e90eb79ef8..a371f1c65e 100644
--- a/packages/auth/pkg/auth/service.go
+++ b/packages/auth/pkg/auth/service.go
@@ -81,7 +81,7 @@ func (s *AuthService[T]) ValidateAPIKey(ctx context.Context, ginCtx *gin.Context
 	}
 
 	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
-	telemetry.SetAttributes(ginCtx.Request.Context(), telemetry.WithAPIKey(keys.ApiKeyPrefix, apiKey))
+	telemetry.SetAttributes(ginCtx.Request.Context(), telemetry.WithAPIKey(apiKey))
 
 	return result, nil
 }
@@ -107,7 +107,7 @@ func (s *AuthService[T]) ValidateAccessToken(ctx context.Context, ginCtx *gin.Co
 	}
 
 	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
-	telemetry.SetAttributes(ginCtx.Request.Context(), telemetry.WithAccessToken(keys.AccessTokenPrefix, accessToken))
+	telemetry.SetAttributes(ginCtx.Request.Context(), telemetry.WithAccessToken(accessToken))
 
 	return userID, nil
 }
diff --git a/packages/shared/pkg/logger/fields.go b/packages/shared/pkg/logger/fields.go
index bc179c8418..cc16ec575f 100644
--- a/packages/shared/pkg/logger/fields.go
+++ b/packages/shared/pkg/logger/fields.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"strings"
 
+	"github.com/e2b-dev/infra/packages/shared/pkg/keys"
 	"github.com/google/uuid"
 	"go.uber.org/zap"
 )
@@ -54,12 +55,12 @@ func WithClientIP(clientIP string) zap.Field {
 	return zap.String("http.client_ip", clientIP)
 }
 
-func WithAPIKey(prefix, apiKey string) zap.Field {
-	return zap.String("auth.api_key", tokenHint(prefix, apiKey))
+func WithAPIKey(apiKey string) zap.Field {
+	return zap.String("auth.api_key", maskedToken(keys.ApiKeyPrefix, apiKey))
 }
 
-func WithAccessToken(prefix, accessToken string) zap.Field {
-	return zap.String("auth.access_token", tokenHint(prefix, accessToken))
+func WithAccessToken(accessToken string) zap.Field {
+	return zap.String("auth.access_token", maskedToken(keys.AccessTokenPrefix, accessToken))
 }
 
 // ProxyRequestFields returns the common logger fields for a proxied HTTP request.
@@ -77,15 +78,14 @@ func ProxyRequestFields(r *http.Request, sandboxID string, sandboxPort uint64) [
 	}
 }
 
-// tokenHint returns a masked version of the token value for debugging.
-// It strips the prefix and shows the first 2 and last 2 characters (e.g. "<prefix>ab...9f").
-func tokenHint(prefix string, token string) string {
-	value := strings.TrimPrefix(token, prefix)
-	if len(value) < 5 {
-		return fmt.Sprintf("%s...", prefix)
+func maskedToken(prefix string, token string) string {
+	tokenWithoutPrefix := strings.TrimPrefix(token, prefix)
+	masked, err := keys.MaskKey(prefix, tokenWithoutPrefix)
+	if err != nil {
+		return "invalid_token_format"
 	}
 
-	return fmt.Sprintf("%s%s...%s", prefix, value[:2], value[len(value)-2:])
+	return fmt.Sprintf("%s%s...%s", masked.Prefix, masked.MaskedValuePrefix, masked.MaskedValueSuffix)
 }
 
 // clientIP extracts the real client IP from the request.
diff --git a/packages/shared/pkg/telemetry/fields.go b/packages/shared/pkg/telemetry/fields.go
index fc6fd6f1de..671bee07fc 100644
--- a/packages/shared/pkg/telemetry/fields.go
+++ b/packages/shared/pkg/telemetry/fields.go
@@ -44,12 +44,12 @@ func WithEnvdVersion(envdVersion string) attribute.KeyValue {
 	return zapFieldToOTELAttribute(logger.WithEnvdVersion(envdVersion))
 }
 
-func WithAPIKey(prefix, apiKey string) attribute.KeyValue {
-	return zapFieldToOTELAttribute(logger.WithAPIKey(prefix, apiKey))
+func WithAPIKey(apiKey string) attribute.KeyValue {
+	return zapFieldToOTELAttribute(logger.WithAPIKey(apiKey))
 }
 
-func WithAccessToken(prefix, accessToken string) attribute.KeyValue {
-	return zapFieldToOTELAttribute(logger.WithAccessToken(prefix, accessToken))
+func WithAccessToken(accessToken string) attribute.KeyValue {
+	return zapFieldToOTELAttribute(logger.WithAccessToken(accessToken))
 }
 
 func zapFieldToOTELAttribute(f zap.Field) attribute.KeyValue {

From 6550c801448ee804f9ade51b8393ca22d48d67ca Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Thu, 12 Mar 2026 14:13:41 +0100
Subject: [PATCH 4/9] chore: add comment

---
 packages/shared/pkg/logger/fields.go | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/packages/shared/pkg/logger/fields.go b/packages/shared/pkg/logger/fields.go
index cc16ec575f..ebb4b71198 100644
--- a/packages/shared/pkg/logger/fields.go
+++ b/packages/shared/pkg/logger/fields.go
@@ -90,6 +90,10 @@ func maskedToken(prefix string, token string) string {
 
 // clientIP extracts the real client IP from the request.
 // It reads the first entry from X-Forwarded-For, falling back to RemoteAddr with the port stripped.
+//
+// This assumes a trusted upstream proxy overwrites the
+// X-Forwarded-For header with the real client IP. The header value is NOT
+// client-controllable in this setup because the LB always replaces it.
 func clientIP(r *http.Request) string {
 	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
 		if ip := strings.TrimSpace(strings.SplitN(xff, ",", 2)[0]); ip != "" {

From d62fa6e75cffa2c08561e78a4043be515210d4b0 Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Thu, 12 Mar 2026 14:37:28 +0100
Subject: [PATCH 5/9] fix: extra import

---
 packages/orchestrator/internal/proxy/proxy.go | 1 -
 1 file changed, 1 deletion(-)

diff --git a/packages/orchestrator/internal/proxy/proxy.go b/packages/orchestrator/internal/proxy/proxy.go
index 2c4dbd930b..dbe1bea0a9 100644
--- a/packages/orchestrator/internal/proxy/proxy.go
+++ b/packages/orchestrator/internal/proxy/proxy.go
@@ -16,7 +16,6 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/env"
 	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
-	"github.com/e2b-dev/infra/packages/shared/pkg/httputil"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	reverseproxy "github.com/e2b-dev/infra/packages/shared/pkg/proxy"
 	"github.com/e2b-dev/infra/packages/shared/pkg/proxy/pool"

From ce89e6298402c8281750f7c2d3f87dc2dfea61e0 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]" <github-actions[bot]@users.noreply.github.com>
Date: Thu, 12 Mar 2026 13:39:37 +0000
Subject: [PATCH 6/9] chore: auto-commit generated changes

---
 packages/shared/pkg/logger/fields.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/packages/shared/pkg/logger/fields.go b/packages/shared/pkg/logger/fields.go
index ebb4b71198..baf1c1c113 100644
--- a/packages/shared/pkg/logger/fields.go
+++ b/packages/shared/pkg/logger/fields.go
@@ -6,9 +6,10 @@ import (
 	"net/http"
 	"strings"
 
-	"github.com/e2b-dev/infra/packages/shared/pkg/keys"
 	"github.com/google/uuid"
 	"go.uber.org/zap"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/keys"
 )
 
 func WithSandboxID(sandboxID string) zap.Field {

From 4bf4c51b639bb95e17a43e0cd65a0643a2adc505 Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Thu, 12 Mar 2026 20:35:43 +0100
Subject: [PATCH 7/9] chore: add more info

---
 packages/api/internal/handlers/store.go       |  4 +--
 packages/auth/pkg/auth/service.go             | 33 +++++++++++++++----
 packages/auth/pkg/types/teams.go              |  6 +++-
 .../dashboard-api/internal/handlers/store.go  |  4 +--
 packages/shared/pkg/logger/fields.go          |  4 +++
 packages/shared/pkg/telemetry/fields.go       |  4 +++
 6 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/packages/api/internal/handlers/store.go b/packages/api/internal/handlers/store.go
index a3f35d5437..bc1c109bda 100644
--- a/packages/api/internal/handlers/store.go
+++ b/packages/api/internal/handlers/store.go
@@ -286,11 +286,11 @@ func (a *APIStore) GetUserFromAccessToken(ctx context.Context, ginCtx *gin.Conte
 	return a.authService.ValidateAccessToken(ctx, ginCtx, accessToken)
 }
 
-func (a *APIStore) GetUserIDFromSupabaseToken(ctx context.Context, _ *gin.Context, supabaseToken string) (uuid.UUID, *api.APIError) {
+func (a *APIStore) GetUserIDFromSupabaseToken(ctx context.Context, ginCtx *gin.Context, supabaseToken string) (uuid.UUID, *api.APIError) {
 	ctx, span := tracer.Start(ctx, "get user id from supabase token")
 	defer span.End()
 
-	return a.authService.ValidateSupabaseToken(ctx, supabaseToken)
+	return a.authService.ValidateSupabaseToken(ctx, ginCtx, supabaseToken)
 }
 
 func (a *APIStore) GetTeamFromSupabaseToken(ctx context.Context, ginCtx *gin.Context, teamID string) (*types.Team, *api.APIError) {
diff --git a/packages/auth/pkg/auth/service.go b/packages/auth/pkg/auth/service.go
index a371f1c65e..43084a4fd3 100644
--- a/packages/auth/pkg/auth/service.go
+++ b/packages/auth/pkg/auth/service.go
@@ -13,22 +13,26 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
 
+type TeamItem interface {
+	TeamID() string
+}
+
 // AuthStore abstracts the DB operations needed for auth validation.
-type AuthStore[T any] interface {
+type AuthStore[T TeamItem] interface {
 	GetTeamByHashedAPIKey(ctx context.Context, hashedKey string) (T, error)
 	GetTeamByIDAndUserID(ctx context.Context, userID uuid.UUID, teamID string) (T, error)
 	GetUserIDByHashedAccessToken(ctx context.Context, hashedToken string) (uuid.UUID, error)
 }
 
 // AuthService encapsulates the cache, store, and JWT secrets for auth validation.
-type AuthService[T any] struct {
+type AuthService[T TeamItem] struct {
 	store      AuthStore[T]
 	teamCache  *AuthCache[T]
 	jwtSecrets []string
 }
 
 // NewAuthService creates an AuthService with the given store, cache, and JWT secrets.
-func NewAuthService[T any](store AuthStore[T], teamCache *AuthCache[T], jwtSecrets []string) *AuthService[T] {
+func NewAuthService[T TeamItem](store AuthStore[T], teamCache *AuthCache[T], jwtSecrets []string) *AuthService[T] {
 	return &AuthService[T]{
 		store:      store,
 		teamCache:  teamCache,
@@ -81,7 +85,10 @@ func (s *AuthService[T]) ValidateAPIKey(ctx context.Context, ginCtx *gin.Context
 	}
 
 	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
-	telemetry.SetAttributes(ginCtx.Request.Context(), telemetry.WithAPIKey(apiKey))
+	telemetry.SetAttributes(ginCtx.Request.Context(),
+		telemetry.WithAPIKey(apiKey),
+		telemetry.WithTeamID(result.TeamID()),
+	)
 
 	return result, nil
 }
@@ -107,13 +114,16 @@ func (s *AuthService[T]) ValidateAccessToken(ctx context.Context, ginCtx *gin.Co
 	}
 
 	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
-	telemetry.SetAttributes(ginCtx.Request.Context(), telemetry.WithAccessToken(accessToken))
+	telemetry.SetAttributes(ginCtx.Request.Context(),
+		telemetry.WithAccessToken(accessToken),
+		telemetry.WithUserID(userID.String()),
+	)
 
 	return userID, nil
 }
 
 // ValidateSupabaseToken parses a Supabase JWT and extracts the user ID.
-func (s *AuthService[T]) ValidateSupabaseToken(ctx context.Context, supabaseToken string) (uuid.UUID, *APIError) {
+func (s *AuthService[T]) ValidateSupabaseToken(ctx context.Context, ginCtx *gin.Context, supabaseToken string) (uuid.UUID, *APIError) {
 	userID, err := ParseUserIDFromToken(ctx, s.jwtSecrets, supabaseToken)
 	if err != nil {
 		return uuid.UUID{}, &APIError{
@@ -123,6 +133,11 @@ func (s *AuthService[T]) ValidateSupabaseToken(ctx context.Context, supabaseToke
 		}
 	}
 
+	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
+	telemetry.SetAttributes(ginCtx.Request.Context(),
+		telemetry.WithUserID(userID.String()),
+	)
+
 	return userID, nil
 }
 
@@ -172,6 +187,12 @@ func (s *AuthService[T]) ValidateSupabaseTeam(ctx context.Context, ginCtx *gin.C
 		}
 	}
 
+	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
+	telemetry.SetAttributes(ginCtx.Request.Context(),
+		telemetry.WithUserID(userID.String()),
+		telemetry.WithTeamID(result.TeamID()),
+	)
+
 	return result, nil
 }
 
diff --git a/packages/auth/pkg/types/teams.go b/packages/auth/pkg/types/teams.go
index 45b34ce9ea..74d8a5af4c 100644
--- a/packages/auth/pkg/types/teams.go
+++ b/packages/auth/pkg/types/teams.go
@@ -1,7 +1,7 @@
 package types
 
 import (
-	"github.com/e2b-dev/infra/packages/db/pkg/auth/queries"
+	authqueries "github.com/e2b-dev/infra/packages/db/pkg/auth/queries"
 )
 
 type Team struct {
@@ -10,6 +10,10 @@ type Team struct {
 	Limits *TeamLimits
 }
 
+func (t *Team) TeamID() string {
+	return t.Team.ID.String()
+}
+
 func newTeamLimits(
 	teamLimits *authqueries.TeamLimit,
 ) *TeamLimits {
diff --git a/packages/dashboard-api/internal/handlers/store.go b/packages/dashboard-api/internal/handlers/store.go
index f5aa255769..0cb45346a3 100644
--- a/packages/dashboard-api/internal/handlers/store.go
+++ b/packages/dashboard-api/internal/handlers/store.go
@@ -47,8 +47,8 @@ func (s *APIStore) GetHealth(c *gin.Context) {
 	})
 }
 
-func (s *APIStore) GetUserIDFromSupabaseToken(ctx context.Context, _ *gin.Context, supabaseToken string) (uuid.UUID, *sharedauth.APIError) {
-	return s.authService.ValidateSupabaseToken(ctx, supabaseToken)
+func (s *APIStore) GetUserIDFromSupabaseToken(ctx context.Context, ginCtx *gin.Context, supabaseToken string) (uuid.UUID, *sharedauth.APIError) {
+	return s.authService.ValidateSupabaseToken(ctx, ginCtx, supabaseToken)
 }
 
 func (s *APIStore) GetTeamFromSupabaseToken(ctx context.Context, ginCtx *gin.Context, teamID string) (*types.Team, *sharedauth.APIError) {
diff --git a/packages/shared/pkg/logger/fields.go b/packages/shared/pkg/logger/fields.go
index baf1c1c113..1858779203 100644
--- a/packages/shared/pkg/logger/fields.go
+++ b/packages/shared/pkg/logger/fields.go
@@ -28,6 +28,10 @@ func WithExecutionID(executionID string) zap.Field {
 	return zap.String("execution.id", executionID)
 }
 
+func WithUserID(userID string) zap.Field {
+	return zap.String("user.id", userID)
+}
+
 func WithTeamID(teamID string) zap.Field {
 	return zap.String("team.id", teamID)
 }
diff --git a/packages/shared/pkg/telemetry/fields.go b/packages/shared/pkg/telemetry/fields.go
index 671bee07fc..409a029454 100644
--- a/packages/shared/pkg/telemetry/fields.go
+++ b/packages/shared/pkg/telemetry/fields.go
@@ -36,6 +36,10 @@ func WithClusterID(clusterID uuid.UUID) attribute.KeyValue {
 	return zapFieldToOTELAttribute(logger.WithClusterID(clusterID))
 }
 
+func WithUserID(userID string) attribute.KeyValue {
+	return zapFieldToOTELAttribute(logger.WithUserID(userID))
+}
+
 func WithTeamID(teamID string) attribute.KeyValue {
 	return zapFieldToOTELAttribute(logger.WithTeamID(teamID))
 }

From e919be6502f9e67c860b3c01ba7e62eec802f562 Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Fri, 13 Mar 2026 12:56:44 +0100
Subject: [PATCH 8/9] chore: imporve naming

---
 packages/auth/pkg/auth/service.go       | 4 ++--
 packages/shared/pkg/logger/fields.go    | 4 ++--
 packages/shared/pkg/telemetry/fields.go | 8 ++++----
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/packages/auth/pkg/auth/service.go b/packages/auth/pkg/auth/service.go
index 43084a4fd3..4b656caff8 100644
--- a/packages/auth/pkg/auth/service.go
+++ b/packages/auth/pkg/auth/service.go
@@ -86,7 +86,7 @@ func (s *AuthService[T]) ValidateAPIKey(ctx context.Context, ginCtx *gin.Context
 
 	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
 	telemetry.SetAttributes(ginCtx.Request.Context(),
-		telemetry.WithAPIKey(apiKey),
+		telemetry.WithMaskedAPIKey(apiKey),
 		telemetry.WithTeamID(result.TeamID()),
 	)
 
@@ -115,7 +115,7 @@ func (s *AuthService[T]) ValidateAccessToken(ctx context.Context, ginCtx *gin.Co
 
 	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
 	telemetry.SetAttributes(ginCtx.Request.Context(),
-		telemetry.WithAccessToken(accessToken),
+		telemetry.WithMaskedAccessToken(accessToken),
 		telemetry.WithUserID(userID.String()),
 	)
 
diff --git a/packages/shared/pkg/logger/fields.go b/packages/shared/pkg/logger/fields.go
index 1858779203..0af6bf469f 100644
--- a/packages/shared/pkg/logger/fields.go
+++ b/packages/shared/pkg/logger/fields.go
@@ -60,11 +60,11 @@ func WithClientIP(clientIP string) zap.Field {
 	return zap.String("http.client_ip", clientIP)
 }
 
-func WithAPIKey(apiKey string) zap.Field {
+func WithMaskedAPIKey(apiKey string) zap.Field {
 	return zap.String("auth.api_key", maskedToken(keys.ApiKeyPrefix, apiKey))
 }
 
-func WithAccessToken(accessToken string) zap.Field {
+func WithMaskedAccessToken(accessToken string) zap.Field {
 	return zap.String("auth.access_token", maskedToken(keys.AccessTokenPrefix, accessToken))
 }
 
diff --git a/packages/shared/pkg/telemetry/fields.go b/packages/shared/pkg/telemetry/fields.go
index 409a029454..dea45ba5ba 100644
--- a/packages/shared/pkg/telemetry/fields.go
+++ b/packages/shared/pkg/telemetry/fields.go
@@ -48,12 +48,12 @@ func WithEnvdVersion(envdVersion string) attribute.KeyValue {
 	return zapFieldToOTELAttribute(logger.WithEnvdVersion(envdVersion))
 }
 
-func WithAPIKey(apiKey string) attribute.KeyValue {
-	return zapFieldToOTELAttribute(logger.WithAPIKey(apiKey))
+func WithMaskedAPIKey(apiKey string) attribute.KeyValue {
+	return zapFieldToOTELAttribute(logger.WithMaskedAPIKey(apiKey))
 }
 
-func WithAccessToken(accessToken string) attribute.KeyValue {
-	return zapFieldToOTELAttribute(logger.WithAccessToken(accessToken))
+func WithMaskedAccessToken(accessToken string) attribute.KeyValue {
+	return zapFieldToOTELAttribute(logger.WithMaskedAccessToken(accessToken))
 }
 
 func zapFieldToOTELAttribute(f zap.Field) attribute.KeyValue {

From 57e6e27c575f9d228534ef5791ba4dedb8e769ba Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Mon, 16 Mar 2026 21:30:46 +0100
Subject: [PATCH 9/9] chore: address pr comments

---
 packages/auth/pkg/auth/service.go       |  4 ++--
 packages/shared/pkg/keys/key.go         | 12 ++++++++++++
 packages/shared/pkg/logger/fields.go    | 21 ++++-----------------
 packages/shared/pkg/telemetry/fields.go |  8 ++++----
 4 files changed, 22 insertions(+), 23 deletions(-)

diff --git a/packages/auth/pkg/auth/service.go b/packages/auth/pkg/auth/service.go
index 4b656caff8..149dc1ef45 100644
--- a/packages/auth/pkg/auth/service.go
+++ b/packages/auth/pkg/auth/service.go
@@ -86,7 +86,7 @@ func (s *AuthService[T]) ValidateAPIKey(ctx context.Context, ginCtx *gin.Context
 
 	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
 	telemetry.SetAttributes(ginCtx.Request.Context(),
-		telemetry.WithMaskedAPIKey(apiKey),
+		telemetry.WithMaskedAPIKey(keys.MaskToken(keys.ApiKeyPrefix, apiKey)),
 		telemetry.WithTeamID(result.TeamID()),
 	)
 
@@ -115,7 +115,7 @@ func (s *AuthService[T]) ValidateAccessToken(ctx context.Context, ginCtx *gin.Co
 
 	//nolint:contextcheck // We use the gin request context to set attributes on the parent span.
 	telemetry.SetAttributes(ginCtx.Request.Context(),
-		telemetry.WithMaskedAccessToken(accessToken),
+		telemetry.WithMaskedAccessToken(keys.MaskToken(keys.AccessTokenPrefix, accessToken)),
 		telemetry.WithUserID(userID.String()),
 	)
 
diff --git a/packages/shared/pkg/keys/key.go b/packages/shared/pkg/keys/key.go
index 24379a0ff5..fe6da3a939 100644
--- a/packages/shared/pkg/keys/key.go
+++ b/packages/shared/pkg/keys/key.go
@@ -84,6 +84,18 @@ func GenerateKey(prefix string) (Key, error) {
 	}, nil
 }
 
+// MaskToken masks a prefixed token (e.g. API key or access token) for safe logging.
+// The raw token must never be passed to logging or telemetry libraries directly.
+func MaskToken(prefix, token string) string {
+	tokenWithoutPrefix := strings.TrimPrefix(token, prefix)
+	masked, err := MaskKey(prefix, tokenWithoutPrefix)
+	if err != nil {
+		return "invalid_token_format"
+	}
+
+	return fmt.Sprintf("%s%s...%s", masked.Prefix, masked.MaskedValuePrefix, masked.MaskedValueSuffix)
+}
+
 func VerifyKey(prefix string, key string) (string, error) {
 	if !strings.HasPrefix(key, prefix) {
 		return "", fmt.Errorf("invalid key prefix")
diff --git a/packages/shared/pkg/logger/fields.go b/packages/shared/pkg/logger/fields.go
index 0af6bf469f..a48c09b868 100644
--- a/packages/shared/pkg/logger/fields.go
+++ b/packages/shared/pkg/logger/fields.go
@@ -1,15 +1,12 @@
 package logger
 
 import (
-	"fmt"
 	"net"
 	"net/http"
 	"strings"
 
 	"github.com/google/uuid"
 	"go.uber.org/zap"
-
-	"github.com/e2b-dev/infra/packages/shared/pkg/keys"
 )
 
 func WithSandboxID(sandboxID string) zap.Field {
@@ -60,12 +57,12 @@ func WithClientIP(clientIP string) zap.Field {
 	return zap.String("http.client_ip", clientIP)
 }
 
-func WithMaskedAPIKey(apiKey string) zap.Field {
-	return zap.String("auth.api_key", maskedToken(keys.ApiKeyPrefix, apiKey))
+func WithMaskedAPIKey(maskedAPIKey string) zap.Field {
+	return zap.String("auth.api_key", maskedAPIKey)
 }
 
-func WithMaskedAccessToken(accessToken string) zap.Field {
-	return zap.String("auth.access_token", maskedToken(keys.AccessTokenPrefix, accessToken))
+func WithMaskedAccessToken(maskedAccessToken string) zap.Field {
+	return zap.String("auth.access_token", maskedAccessToken)
 }
 
 // ProxyRequestFields returns the common logger fields for a proxied HTTP request.
@@ -83,16 +80,6 @@ func ProxyRequestFields(r *http.Request, sandboxID string, sandboxPort uint64) [
 	}
 }
 
-func maskedToken(prefix string, token string) string {
-	tokenWithoutPrefix := strings.TrimPrefix(token, prefix)
-	masked, err := keys.MaskKey(prefix, tokenWithoutPrefix)
-	if err != nil {
-		return "invalid_token_format"
-	}
-
-	return fmt.Sprintf("%s%s...%s", masked.Prefix, masked.MaskedValuePrefix, masked.MaskedValueSuffix)
-}
-
 // clientIP extracts the real client IP from the request.
 // It reads the first entry from X-Forwarded-For, falling back to RemoteAddr with the port stripped.
 //
diff --git a/packages/shared/pkg/telemetry/fields.go b/packages/shared/pkg/telemetry/fields.go
index dea45ba5ba..0baa5f859d 100644
--- a/packages/shared/pkg/telemetry/fields.go
+++ b/packages/shared/pkg/telemetry/fields.go
@@ -48,12 +48,12 @@ func WithEnvdVersion(envdVersion string) attribute.KeyValue {
 	return zapFieldToOTELAttribute(logger.WithEnvdVersion(envdVersion))
 }
 
-func WithMaskedAPIKey(apiKey string) attribute.KeyValue {
-	return zapFieldToOTELAttribute(logger.WithMaskedAPIKey(apiKey))
+func WithMaskedAPIKey(maskedAPIKey string) attribute.KeyValue {
+	return zapFieldToOTELAttribute(logger.WithMaskedAPIKey(maskedAPIKey))
 }
 
-func WithMaskedAccessToken(accessToken string) attribute.KeyValue {
-	return zapFieldToOTELAttribute(logger.WithMaskedAccessToken(accessToken))
+func WithMaskedAccessToken(maskedAccessToken string) attribute.KeyValue {
+	return zapFieldToOTELAttribute(logger.WithMaskedAccessToken(maskedAccessToken))
 }
 
 func zapFieldToOTELAttribute(f zap.Field) attribute.KeyValue {