fix(orchestrator): tolerate unsignable uploads on a layer cache hit

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

@@ -4,11 +4,13 @@ package server
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"time"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/storage/paths"
 	templatemanager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
 )
 
 const signedUrlExpiration = time.Minute * 30
@@ -29,14 +31,20 @@ func (s *ServerStore) InitLayerFileUpload(ctx context.Context, in *templatemanag
 		return nil, fmt.Errorf("failed to open layer files cache: %w", err)
 	}
 
-	signedUrl, err := s.buildStorage.UploadSignedURL(ctx, path, signedUrlExpiration)
+	exists, err := obj.Exists(ctx)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get signed url: %w", err)
+		return nil, fmt.Errorf("failed to check if layer files exists: %w", err)
 	}
 
-	exists, err := obj.Exists(ctx)
+	signedUrl, err := s.buildStorage.UploadSignedURL(ctx, path, signedUrlExpiration)
 	if err != nil {
-		return nil, fmt.Errorf("failed to check if layer files exists: %w", err)
+		// A cache hit needs no upload URL, so a provider that cannot sign one is fatal
+		// only on a miss.
+		if exists && errors.Is(err, storage.ErrSignedUploadURLUnsupported) {
+			return &templatemanager.InitLayerFileUploadResponse{Present: true}, nil
+		}
+
+		return nil, fmt.Errorf("failed to get signed url: %w", err)
 	}
 
 	return &templatemanager.InitLayerFileUploadResponse{

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

@@ -0,0 +1,100 @@
+//go:build linux
+
+package server
+
+import (
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/storage/paths"
+	templatemanager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
+)
+
+const (
+	testTemplateID = "tmpl-init-layer-upload"
+	testFilesHash  = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
+)
+
+func newInitLayerFileUploadServer(t *testing.T, exists bool, signedURL string, signErr error) *ServerStore {
+	t.Helper()
+
+	blob := storage.NewMockBlob(t)
+	blob.EXPECT().Exists(mock.Anything).Return(exists, nil)
+
+	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)
+
+	return &ServerStore{buildStorage: provider}
+}
+
+func initLayerFileUploadRequest() *templatemanager.InitLayerFileUploadRequest {
+	return &templatemanager.InitLayerFileUploadRequest{
+		TemplateID: testTemplateID,
+		Hash:       testFilesHash,
+	}
+}
+
+func TestInitLayerFileUploadUnsignableProvider(t *testing.T) {
+	t.Parallel()
+
+	unsupported := fmt.Errorf("%w: test provider", storage.ErrSignedUploadURLUnsupported)
+
+	t.Run("cache hit reports present without a url", func(t *testing.T) {
+		t.Parallel()
+
+		s := newInitLayerFileUploadServer(t, true, "", unsupported)
+
+		resp, err := s.InitLayerFileUpload(t.Context(), initLayerFileUploadRequest())
+		require.NoError(t, err)
+		assert.True(t, resp.GetPresent())
+		assert.Nil(t, resp.Url, "a cache hit must not hand back an unusable url")
+	})
+
+	t.Run("cache miss still fails", func(t *testing.T) {
+		t.Parallel()
+
+		s := newInitLayerFileUploadServer(t, false, "", unsupported)
+
+		_, err := s.InitLayerFileUpload(t.Context(), initLayerFileUploadRequest())
+		require.ErrorIs(t, err, storage.ErrSignedUploadURLUnsupported)
+	})
+}
+
+// Signing failures that are not the "unsupported" sentinel stay fatal even on a cache hit:
+// tolerating them would hide a broken credential behind a green response.
+func TestInitLayerFileUploadSigningErrorOnCacheHit(t *testing.T) {
+	t.Parallel()
+
+	signErr := errors.New("failed to parse GCP service account")
+	s := newInitLayerFileUploadServer(t, true, "", signErr)
+
+	_, err := s.InitLayerFileUpload(t.Context(), initLayerFileUploadRequest())
+	require.ErrorIs(t, err, signErr)
+}
+
+// The url is returned on a cache hit as well: the SDK gates forceUpload re-uploads on
+// url != nil, so dropping it there would silently turn forceUpload into a no-op.
+func TestInitLayerFileUploadKeepsURLOnCacheHit(t *testing.T) {
+	t.Parallel()
+
+	for _, exists := range []bool{true, false} {
+		t.Run(fmt.Sprintf("exists=%v", exists), func(t *testing.T) {
+			t.Parallel()
+
+			s := newInitLayerFileUploadServer(t, exists, "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())
+		})
+	}
+}

packages/shared/pkg/storage/storage.go

@@ -21,6 +21,10 @@ var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/shared/pkg/storage")
 
 var ErrObjectNotExist = errors.New("object does not exist")
 
+// ErrSignedUploadURLUnsupported means the provider cannot sign an upload URL usable by an
+// external client. Callers that only need the URL on a cache miss may tolerate it.
+var ErrSignedUploadURLUnsupported = errors.New("signed upload URLs are not supported by this storage provider")
+
 // ErrObjectRateLimited means per-object mutation rate limiting —
 // multiple concurrent writers racing to write the same content-addressed object.
 var ErrObjectRateLimited = errors.New("object access rate limited")

packages/shared/pkg/storage/storage_azure.go

@@ -190,7 +190,7 @@ func (s *azureStorage) GetDetails() string {
 // shape as the filesystem provider's local upload path, and no proto or public-API
 // change); SAS PUT URLs are deliberately never issued.
 func (s *azureStorage) UploadSignedURL(_ context.Context, path string, _ time.Duration) (string, error) {
-	return "", fmt.Errorf("signed upload URLs are not supported on Azure (%q): Put Blob requires the x-ms-blob-type request header, which a SAS cannot carry and external upload clients do not send", path)
+	return "", fmt.Errorf("%w: Azure (%q) Put Blob requires the x-ms-blob-type request header, which a SAS cannot carry and external upload clients do not send", ErrSignedUploadURLUnsupported, path)
 }
 
 func (s *azureStorage) OpenSeekable(_ context.Context, path string) (Seekable, error) {

packages/shared/pkg/storage/storage_azure_integration_test.go

@@ -339,7 +339,7 @@ func TestAzureIntegration(t *testing.T) {
 		// The interface method must fail loudly rather than return a URL an external
 		// client cannot use; see its doc comment.
 		_, err := provider.UploadSignedURL(ctx, "signed/refused.bin", time.Hour)
-		require.Error(t, err)
+		require.ErrorIs(t, err, ErrSignedUploadURLUnsupported)
 		assert.Contains(t, err.Error(), "x-ms-blob-type")
 	})