feat(storage): sign Azure layer-file uploads with a SAS

packages/api/internal/handlers/template_layer_files_upload.go

@@ -68,8 +68,13 @@ func (a *APIStore) GetTemplatesTemplateIDFilesHash(c *gin.Context, templateID ap
 		return
 	}
 
-	c.JSON(http.StatusCreated, &api.TemplateBuildFileUpload{
+	upload := api.TemplateBuildFileUpload{
 		Present: resp.GetPresent(),
 		Url:     resp.Url,
-	})
+	}
+	if headers := resp.GetUploadHeaders(); len(headers) > 0 {
+		upload.Headers = &headers
+	}
+
+	c.JSON(http.StatusCreated, &upload)
 }

packages/api/internal/handlers/template_layer_files_upload_test.go

@@ -0,0 +1,41 @@
+package handlers
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/api/internal/api"
+)
+
+// A client written against the pre-headers response must keep parsing the new one, and the
+// providers that need no request headers must keep producing the old bytes exactly.
+func TestTemplateBuildFileUploadHeadersAreAdditive(t *testing.T) {
+	t.Parallel()
+
+	t.Run("no headers keeps the response bytes unchanged", func(t *testing.T) {
+		t.Parallel()
+
+		body, err := json.Marshal(api.TemplateBuildFileUpload{Present: false, Url: new("https://bucket.example/signed")})
+		require.NoError(t, err)
+		assert.JSONEq(t, `{"present":false,"url":"https://bucket.example/signed"}`, string(body))
+	})
+
+	t.Run("a client ignoring headers still parses the response", func(t *testing.T) {
+		t.Parallel()
+
+		var legacy struct {
+			Present bool    `json:"present"`
+			Url     *string `json:"url"`
+		}
+		require.NoError(t, json.Unmarshal(
+			[]byte(`{"present":false,"url":"https://account.example/signed","headers":{"x-ms-blob-type":"BlockBlob"}}`),
+			&legacy))
+
+		assert.False(t, legacy.Present)
+		require.NotNil(t, legacy.Url)
+		assert.Equal(t, "https://account.example/signed", *legacy.Url)
+	})
+}

packages/orchestrator/pkg/sandbox/template/peerclient/storage.go

@@ -94,7 +94,7 @@ func (p *routingProvider) DeleteObjectsWithPrefix(ctx context.Context, prefix st
 	return p.base.DeleteObjectsWithPrefix(ctx, prefix)
 }
 
-func (p *routingProvider) UploadSignedURL(ctx context.Context, path string, ttl time.Duration) (string, error) {
+func (p *routingProvider) UploadSignedURL(ctx context.Context, path string, ttl time.Duration) (storage.UploadURL, error) {
 	return p.base.UploadSignedURL(ctx, path, ttl)
 }
 
@@ -168,7 +168,7 @@ func (p *peerStorageProvider) DeleteObjectsWithPrefix(ctx context.Context, prefi
 	return p.base.DeleteObjectsWithPrefix(ctx, prefix)
 }
 
-func (p *peerStorageProvider) UploadSignedURL(ctx context.Context, path string, ttl time.Duration) (string, error) {
+func (p *peerStorageProvider) UploadSignedURL(ctx context.Context, path string, ttl time.Duration) (storage.UploadURL, error) {
 	return p.base.UploadSignedURL(ctx, path, ttl)
 }
 

packages/orchestrator/pkg/template/server/upload_layer_files_template.go

@@ -36,7 +36,7 @@ func (s *ServerStore) InitLayerFileUpload(ctx context.Context, in *templatemanag
 		return nil, fmt.Errorf("failed to check if layer files exists: %w", err)
 	}
 
-	signedUrl, err := s.buildStorage.UploadSignedURL(ctx, path, signedUrlExpiration)
+	upload, err := s.buildStorage.UploadSignedURL(ctx, path, signedUrlExpiration)
 	if err != nil {
 		// A cache hit needs no upload URL, so a provider that cannot sign one is fatal
 		// only on a miss.
@@ -48,7 +48,8 @@ func (s *ServerStore) InitLayerFileUpload(ctx context.Context, in *templatemanag
 	}
 
 	return &templatemanager.InitLayerFileUploadResponse{
-		Present: exists,
-		Url:     &signedUrl,
+		Present:       exists,
+		Url:           &upload.URL,
+		UploadHeaders: upload.Headers,
 	}, nil
 }

packages/orchestrator/pkg/template/server/upload_layer_files_template_test.go

@@ -21,7 +21,7 @@ const (
 	testFilesHash  = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
 )
 
-func newInitLayerFileUploadServer(t *testing.T, exists bool, signedURL string, signErr error) *ServerStore {
+func newInitLayerFileUploadServer(t *testing.T, exists bool, upload storage.UploadURL, signErr error) *ServerStore {
 	t.Helper()
 
 	blob := storage.NewMockBlob(t)
@@ -30,7 +30,7 @@ func newInitLayerFileUploadServer(t *testing.T, exists bool, signedURL string, s
 	provider := storage.NewMockStorageProvider(t)
 	path := paths.GetLayerFilesCachePath(testTemplateID, testFilesHash)
 	provider.EXPECT().OpenBlob(mock.Anything, path).Return(blob, nil)
-	provider.EXPECT().UploadSignedURL(mock.Anything, path, signedUrlExpiration).Return(signedURL, signErr)
+	provider.EXPECT().UploadSignedURL(mock.Anything, path, signedUrlExpiration).Return(upload, signErr)
 
 	return &ServerStore{buildStorage: provider}
 }
@@ -50,7 +50,7 @@ func TestInitLayerFileUploadUnsignableProvider(t *testing.T) {
 	t.Run("cache hit reports present without a url", func(t *testing.T) {
 		t.Parallel()
 
-		s := newInitLayerFileUploadServer(t, true, "", unsupported)
+		s := newInitLayerFileUploadServer(t, true, storage.UploadURL{}, unsupported)
 
 		resp, err := s.InitLayerFileUpload(t.Context(), initLayerFileUploadRequest())
 		require.NoError(t, err)
@@ -61,7 +61,7 @@ func TestInitLayerFileUploadUnsignableProvider(t *testing.T) {
 	t.Run("cache miss still fails", func(t *testing.T) {
 		t.Parallel()
 
-		s := newInitLayerFileUploadServer(t, false, "", unsupported)
+		s := newInitLayerFileUploadServer(t, false, storage.UploadURL{}, unsupported)
 
 		_, err := s.InitLayerFileUpload(t.Context(), initLayerFileUploadRequest())
 		require.ErrorIs(t, err, storage.ErrSignedUploadURLUnsupported)
@@ -74,7 +74,7 @@ func TestInitLayerFileUploadSigningErrorOnCacheHit(t *testing.T) {
 	t.Parallel()
 
 	signErr := errors.New("failed to parse GCP service account")
-	s := newInitLayerFileUploadServer(t, true, "", signErr)
+	s := newInitLayerFileUploadServer(t, true, storage.UploadURL{}, signErr)
 
 	_, err := s.InitLayerFileUpload(t.Context(), initLayerFileUploadRequest())
 	require.ErrorIs(t, err, signErr)
@@ -89,12 +89,24 @@ func TestInitLayerFileUploadKeepsURLOnCacheHit(t *testing.T) {
 		t.Run(fmt.Sprintf("exists=%v", exists), func(t *testing.T) {
 			t.Parallel()
 
-			s := newInitLayerFileUploadServer(t, exists, "https://bucket.example/signed", nil)
+			s := newInitLayerFileUploadServer(t, exists, storage.UploadURL{URL: "https://bucket.example/signed"}, nil)
 
 			resp, err := s.InitLayerFileUpload(t.Context(), initLayerFileUploadRequest())
 			require.NoError(t, err)
 			assert.Equal(t, exists, resp.GetPresent())
 			assert.Equal(t, "https://bucket.example/signed", resp.GetUrl())
+			assert.Empty(t, resp.GetUploadHeaders(), "providers that need no request headers must not send any")
 		})
 	}
 }
+
+func TestInitLayerFileUploadForwardsUploadHeaders(t *testing.T) {
+	t.Parallel()
+
+	headers := map[string]string{"x-ms-blob-type": "BlockBlob"}
+	s := newInitLayerFileUploadServer(t, false, storage.UploadURL{URL: "https://account.example/signed", Headers: headers}, nil)
+
+	resp, err := s.InitLayerFileUpload(t.Context(), initLayerFileUploadRequest())
+	require.NoError(t, err)
+	assert.Equal(t, headers, resp.GetUploadHeaders())
+}

packages/orchestrator/template-manager.proto

@@ -16,6 +16,8 @@ message InitLayerFileUploadRequest {
 message InitLayerFileUploadResponse{
   bool present = 1;
   optional string url = 2;
+  // Request headers the upload client must send on the PUT to `url`.
+  map<string, string> uploadHeaders = 3;
 }
 
 message TemplateStep {

packages/shared/pkg/storage/storage.go

@@ -65,9 +65,17 @@ func (t SeekableObjectType) String() string {
 	}
 }
 
+// UploadURL is a signed upload target for an external client. Headers must be sent
+// verbatim on the PUT: Azure's Put Blob rejects the request without x-ms-blob-type, and a
+// SAS can only pin response headers, never require a request one.
+type UploadURL struct {
+	URL     string
+	Headers map[string]string
+}
+
 type StorageProvider interface {
 	DeleteObjectsWithPrefix(ctx context.Context, prefix string) error
-	UploadSignedURL(ctx context.Context, path string, ttl time.Duration) (string, error)
+	UploadSignedURL(ctx context.Context, path string, ttl time.Duration) (UploadURL, error)
 	OpenBlob(ctx context.Context, path string) (Blob, error)
 	OpenSeekable(ctx context.Context, path string) (Seekable, error)
 	GetDetails() string