feat: keep user and workdir from the template in the sandbox (#1349)

Author: dobrac

packages/api/internal/handlers/template_start_build.go

@@ -18,6 +18,7 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/models"
 	"github.com/e2b-dev/infra/packages/shared/pkg/models/envbuild"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+	"github.com/e2b-dev/infra/packages/shared/pkg/templates"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
@@ -141,7 +142,7 @@ func (a *APIStore) PostTemplatesTemplateIDBuildsBuildID(c *gin.Context, template
 	// Call the Template Manager to build the environment
 	forceRebuild := true
 	fromImage := ""
-	buildErr := a.templateManager.CreateTemplate(
+	err = a.templateManager.CreateTemplate(
 		ctx,
 		team.ID,
 		templateID,
@@ -160,20 +161,22 @@ func (a *APIStore) PostTemplatesTemplateIDBuildsBuildID(c *gin.Context, template
 		nil,
 		apiutils.WithClusterFallback(team.ClusterID),
 		build.ClusterNodeID,
+		templates.TemplateV1Version,
 	)
-	if buildErr != nil {
-		telemetry.ReportCriticalError(ctx, "build failed", buildErr, telemetry.WithTemplateID(templateID))
-		a.sendAPIStoreError(c, http.StatusInternalServerError, fmt.Sprintf("Error when starting template build: %s", buildErr))
-		return
-	}
 
 	a.posthog.CreateAnalyticsUserEvent(userID.String(), team.ID.String(), "built environment", posthog.NewProperties().
 		Set("user_id", userID).
 		Set("environment", templateID).
 		Set("build_id", buildID).
 		Set("duration", time.Since(startTime).String()).
-		Set("success", err != nil),
+		Set("success", err == nil),
 	)
 
+	if err != nil {
+		telemetry.ReportCriticalError(ctx, "build failed", err, telemetry.WithTemplateID(templateID))
+		a.sendAPIStoreError(c, http.StatusInternalServerError, fmt.Sprintf("Error when starting template build: %s", err))
+		return
+	}
+
 	c.Status(http.StatusAccepted)
 }

packages/api/internal/handlers/template_start_build_v2.go

@@ -4,17 +4,27 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 	"github.com/posthog/posthog-go"
+	"go.uber.org/zap"
 
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	apiutils "github.com/e2b-dev/infra/packages/api/internal/utils"
 	"github.com/e2b-dev/infra/packages/db/queries"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/models/envbuild"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+	"github.com/e2b-dev/infra/packages/shared/pkg/templates"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
+)
+
+const (
+	jsSDKPrefix     = "e2b-js-sdk/"
+	pythonSDKPrefix = "e2b-python-sdk/"
 )
 
 type dockerfileStore struct {
@@ -117,8 +127,15 @@ func (a *APIStore) PostV2TemplatesTemplateIDBuildsBuildID(c *gin.Context, templa
 		return
 	}
 
+	version, err := userAgentToTemplateVersion(zap.L().With(logger.WithTemplateID(templateID), logger.WithBuildID(buildID)), c.Request.UserAgent())
+	if err != nil {
+		a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Error when parsing user agent: %s", err))
+		telemetry.ReportCriticalError(ctx, "error when parsing user agent", err, telemetry.WithTemplateID(templateID))
+		return
+	}
+
 	// Call the Template Manager to build the environment
-	buildErr := a.templateManager.CreateTemplate(
+	err = a.templateManager.CreateTemplate(
 		ctx,
 		team.ID,
 		templateID,
@@ -137,19 +154,56 @@ func (a *APIStore) PostV2TemplatesTemplateIDBuildsBuildID(c *gin.Context, templa
 		body.Steps,
 		apiutils.WithClusterFallback(team.ClusterID),
 		build.ClusterNodeID,
+		version,
 	)
-	if buildErr != nil {
-		telemetry.ReportCriticalError(ctx, "build failed", buildErr, telemetry.WithTemplateID(templateID))
-		a.sendAPIStoreError(c, http.StatusInternalServerError, fmt.Sprintf("Error when starting template build: %s", buildErr))
-		return
-	}
 
 	a.posthog.CreateAnalyticsTeamEvent(team.ID.String(), "built environment", posthog.NewProperties().
 		Set("environment", templateID).
 		Set("build_id", buildID).
 		Set("duration", time.Since(startTime).String()).
-		Set("success", err != nil),
+		Set("success", err == nil),
 	)
 
+	if err != nil {
+		telemetry.ReportCriticalError(ctx, "build failed", err, telemetry.WithTemplateID(templateID))
+		a.sendAPIStoreError(c, http.StatusInternalServerError, fmt.Sprintf("Error when starting template build: %s", err))
+		return
+	}
+
 	c.Status(http.StatusAccepted)
 }
+
+// userAgentToTemplateVersion returns the template semver version based on the user agent string.
+// If the user agent is not recognized, it defaults to the latest stable version.
+func userAgentToTemplateVersion(logger *zap.Logger, userAgent string) (string, error) {
+	version := templates.TemplateV2LatestVersion
+
+	switch {
+	case strings.HasPrefix(userAgent, jsSDKPrefix):
+		sdk := strings.TrimPrefix(userAgent, jsSDKPrefix)
+
+		// Check if the SDK version supports the latest template version
+		ok, err := utils.IsGTEVersion(sdk, templates.SDKTemplateReleaseVersion)
+		if err != nil {
+			return "", fmt.Errorf("parsing JS SDK version: %w", err)
+		}
+		if !ok {
+			version = templates.TemplateV2BetaVersion
+		}
+	case strings.HasPrefix(userAgent, pythonSDKPrefix):
+		sdk := strings.TrimPrefix(userAgent, pythonSDKPrefix)
+
+		// Check if the SDK version supports the latest template version
+		ok, err := utils.IsGTEVersion(sdk, templates.SDKTemplateReleaseVersion)
+		if err != nil {
+			return "", fmt.Errorf("parsing Python SDK version: %w", err)
+		}
+		if !ok {
+			version = templates.TemplateV2BetaVersion
+		}
+	default:
+		logger.Debug("Unrecognized user agent, defaulting to the latest template version", zap.String("user_agent", userAgent), zap.String("version", version))
+	}
+
+	return version, nil
+}

packages/api/internal/template-manager/create_template.go

@@ -52,6 +52,7 @@ func (tm *TemplateManager) CreateTemplate(
 	steps *[]api.TemplateStep,
 	clusterID uuid.UUID,
 	nodeID string,
+	version string,
 ) (e error) {
 	ctx, span := tracer.Start(ctx, "create-template",
 		trace.WithAttributes(
@@ -152,6 +153,7 @@ func (tm *TemplateManager) CreateTemplate(
 		reqCtx, &templatemanagergrpc.TemplateCreateRequest{
 			Template:   template,
 			CacheScope: ut.ToPtr(teamID.String()),
+			Version:    &version,
 		},
 	)
 

packages/envd/internal/api/api.gen.go

@@ -63,7 +63,7 @@ func (a *API) generateSignature(path string, username string, operation string,
 	return fmt.Sprintf("v1_%s", hasher.HashWithoutPrefix([]byte(signature))), nil
 }
 
-func (a *API) validateSigning(r *http.Request, signature *string, signatureExpiration *int, username string, path string, operation string) (err error) {
+func (a *API) validateSigning(r *http.Request, signature *string, signatureExpiration *int, username *string, path string, operation string) (err error) {
 	var expectedSignature string
 
 	// no need to validate signing key if access token is not set
@@ -85,11 +85,17 @@ func (a *API) validateSigning(r *http.Request, signature *string, signatureExpir
 		return fmt.Errorf("missing signature query parameter")
 	}
 
+	// Empty string is used when no username is provided and the default user should be used
+	signatureUsername := ""
+	if username != nil {
+		signatureUsername = *username
+	}
+
 	if signatureExpiration == nil {
-		expectedSignature, err = a.generateSignature(path, username, operation, nil)
+		expectedSignature, err = a.generateSignature(path, signatureUsername, operation, nil)
 	} else {
 		exp := int64(*signatureExpiration)
-		expectedSignature, err = a.generateSignature(path, username, operation, &exp)
+		expectedSignature, err = a.generateSignature(path, signatureUsername, operation, &exp)
 	}
 
 	if err != nil {

packages/envd/internal/api/auth.go

@@ -8,6 +8,7 @@ import (
 	"os/user"
 	"time"
 
+	"github.com/e2b-dev/infra/packages/envd/internal/execcontext"
 	"github.com/e2b-dev/infra/packages/envd/internal/logs"
 	"github.com/e2b-dev/infra/packages/envd/internal/permissions"
 )
@@ -33,13 +34,20 @@ func (a *API) GetFiles(w http.ResponseWriter, r *http.Request, params GetFilesPa
 		return
 	}
 
+	username, err := execcontext.ResolveDefaultUsername(params.Username, a.defaults.User)
+	if err != nil {
+		a.logger.Error().Err(err).Str(string(logs.OperationIDKey), operationID).Msg("no user specified")
+		jsonError(w, http.StatusBadRequest, err)
+		return
+	}
+
 	defer func() {
 		l := a.logger.
 			Err(errMsg).
 			Str("method", r.Method+" "+r.URL.Path).
 			Str(string(logs.OperationIDKey), operationID).
 			Str("path", path).
-			Str("username", params.Username)
+			Str("username", username)
 
 		if errMsg != nil {
 			l = l.Int("error_code", errorCode)
@@ -48,16 +56,16 @@ func (a *API) GetFiles(w http.ResponseWriter, r *http.Request, params GetFilesPa
 		l.Msg("File read")
 	}()
 
-	u, err := user.Lookup(params.Username)
+	u, err := user.Lookup(username)
 	if err != nil {
-		errMsg = fmt.Errorf("error looking up user '%s': %w", params.Username, err)
+		errMsg = fmt.Errorf("error looking up user '%s': %w", username, err)
 		errorCode = http.StatusUnauthorized
 		jsonError(w, errorCode, errMsg)
 
 		return
 	}
 
-	resolvedPath, err := permissions.ExpandAndResolve(path, u)
+	resolvedPath, err := permissions.ExpandAndResolve(path, u, a.defaults.Workdir)
 	if err != nil {
 		errMsg = fmt.Errorf("error expanding and resolving path '%s': %w", path, err)
 		errorCode = http.StatusBadRequest

packages/envd/internal/api/download.go

@@ -13,7 +13,7 @@ func (a *API) GetEnvs(w http.ResponseWriter, _ *http.Request) {
 	a.logger.Debug().Str(string(logs.OperationIDKey), operationID).Msg("Getting env vars")
 
 	envs := make(EnvVars)
-	a.envVars.Range(func(key, value string) bool {
+	a.defaults.EnvVars.Range(func(key, value string) bool {
 		envs[key] = value
 
 		return true

packages/envd/internal/api/envs.go

@@ -69,7 +69,7 @@ func (a *API) PostInit(w http.ResponseWriter, r *http.Request) {
 	go func() { //nolint:contextcheck // TODO: fix this later
 		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
 		defer cancel()
-		host.PollForMMDSOpts(ctx, a.mmdsChan, a.envVars)
+		host.PollForMMDSOpts(ctx, a.mmdsChan, a.defaults.EnvVars)
 	}()
 
 	w.Header().Set("Cache-Control", "no-store")
@@ -99,7 +99,7 @@ func (a *API) SetData(logger zerolog.Logger, data PostInitJSONBody) error {
 
 		for key, value := range *data.EnvVars {
 			logger.Debug().Msgf("Setting env var for %s", key)
-			a.envVars.Store(key, value)
+			a.defaults.EnvVars.Store(key, value)
 		}
 	}
 
@@ -117,6 +117,16 @@ func (a *API) SetData(logger zerolog.Logger, data PostInitJSONBody) error {
 		go a.SetupHyperloop(*data.HyperloopIP)
 	}
 
+	if data.DefaultUser != nil && *data.DefaultUser != "" {
+		logger.Debug().Msgf("Setting default user to: %s", *data.DefaultUser)
+		a.defaults.User = *data.DefaultUser
+	}
+
+	if data.DefaultWorkdir != nil && *data.DefaultWorkdir != "" {
+		logger.Debug().Msgf("Setting default workdir to: %s", *data.DefaultWorkdir)
+		a.defaults.Workdir = data.DefaultWorkdir
+	}
+
 	return nil
 }
 
@@ -127,7 +137,7 @@ func (a *API) SetupHyperloop(address string) {
 	if err := rewriteHostsFile(address, "/etc/hosts"); err != nil {
 		a.logger.Error().Err(err).Msg("failed to modify hosts file")
 	} else {
-		a.envVars.Store("E2B_EVENTS_ADDRESS", fmt.Sprintf("http://%s", address))
+		a.defaults.EnvVars.Store("E2B_EVENTS_ADDRESS", fmt.Sprintf("http://%s", address))
 	}
 }