azure: fix race condition in PutContent causing duplicate append writes (#4888)

2026-07-17 02:00:41 +00:00 · 2026-06-07 09:20:23 -07:00
parent 915084f4bc 927250bcaf
commit af1b4b4e9a
2 changed files with 200 additions and 90 deletions
--- a/registry/storage/driver/azure/azure.go
+++ b/registry/storage/driver/azure/azure.go
@@ -119,55 +119,65 @@ func (d *driver) GetContent(ctx context.Context, path string) ([]byte, error) {

 // PutContent stores the []byte content at a location designated by "path".
 func (d *driver) PutContent(ctx context.Context, path string, contents []byte) error {
-	// TODO(milosgajdos): this check is not needed as UploadBuffer will return error if we exceed the max blockbytes limit
 	// max size for block blobs uploaded via single "Put Blob" for version after "2016-05-31"
 	// https://docs.microsoft.com/en-us/rest/api/storageservices/put-blob#remarks
 	if len(contents) > blockblob.MaxUploadBlobBytes {
 		return fmt.Errorf("content size exceeds max allowed limit (%d): %d", blockblob.MaxUploadBlobBytes, len(contents))
 	}

-	// PutContent creates blobs as AppendBlob. Blobs written by older versions
-	// of the driver may be BlockBlob (or other types), and Azure returns
-	// 409 InvalidBlobType if we try to create an AppendBlob over an existing
-	// blob of a different type. Delete any such legacy blob first so the
-	// subsequent Create succeeds.
-	//
-	// While we delete the blob and create a new one, there will be a small
-	// window of inconsistency and if the Create fails, we may end up losing
-	// the existing data. However, the expectation is that clients pushing
-	// will retry when they get an error response.
 	blobName := d.blobName(path)
+
+	// Azure cannot atomically replace an existing AppendBlob with a BlockBlob.
+	// Delete the old PutContent representation before the block-blob upload.
+	if err := d.deleteLegacyPutContentBlob(ctx, blobName); err != nil {
+		return err
+	}
+
+	if err := d.uploadBlockBlob(ctx, blobName, contents); err != nil {
+		return fmt.Errorf("failed to upload block blob: %w", err)
+	}
+
+	return nil
+}
+
+func (d *driver) uploadBlockBlob(ctx context.Context, blobName string, contents []byte) error {
+	blockBlobRef := d.client.NewBlockBlobClient(blobName)
+	_, err := blockBlobRef.Upload(ctx, streaming.NopCloser(bytes.NewReader(contents)), nil)
+	return err
+}
+
+func (d *driver) deleteLegacyPutContentBlob(ctx context.Context, blobName string) error {
 	blobRef := d.client.NewBlobClient(blobName)
 	props, err := blobRef.GetProperties(ctx, nil)
-	if err != nil && !is404(err) {
+	if err != nil {
+		if is404(err) {
+			return nil
+		}
 		return fmt.Errorf("failed to get blob properties: %v", err)
 	}
-	if err == nil && props.BlobType != nil && *props.BlobType != blob.BlobTypeAppendBlob {
-		if _, err := blobRef.Delete(ctx, nil); err != nil {
-			return fmt.Errorf("failed to delete legacy blob (%v): %v", *props.BlobType, err)
-		}
+
+	if props.BlobType == nil {
+		return fmt.Errorf("missing BlobType: %s", blobName)
+	}
+	if *props.BlobType == blob.BlobTypeBlockBlob {
+		return nil
+	}
+	if props.ETag == nil {
+		return fmt.Errorf("missing ETag for legacy blob: %s", blobName)
 	}

-	// Always create as AppendBlob
-	appendBlobRef := d.client.NewAppendBlobClient(blobName)
-	if _, err := appendBlobRef.Create(ctx, nil); err != nil {
-		return fmt.Errorf("failed to create append blob: %v", err)
-	}
-
-	// If we have content, append it
-	if len(contents) > 0 {
-		// Write in chunks of maxChunkSize otherwise Azure can barf
-		// when writing large piece of data in one sot:
-		// RESPONSE 413: 413 The uploaded entity blob is too large.
-		for offset := 0; offset < len(contents); offset += maxChunkSize {
-			end := min(offset+maxChunkSize, len(contents))
-
-			chunk := contents[offset:end]
-			_, err := appendBlobRef.AppendBlock(ctx, streaming.NopCloser(bytes.NewReader(chunk)), nil)
-			if err != nil {
-				return fmt.Errorf("failed to append content: %v", err)
-			}
+	_, err = blobRef.Delete(ctx, &blob.DeleteOptions{
+		AccessConditions: &blob.AccessConditions{
+			ModifiedAccessConditions: &blob.ModifiedAccessConditions{
+				IfMatch: props.ETag,
+			},
+		},
+	})
+	if err != nil {
+		if is404(err) || bloberror.HasCode(err, bloberror.ConditionNotMet) {
+			return nil
 		}
+		return fmt.Errorf("failed to delete legacy blob (%v): %v", *props.BlobType, err)
 	}

 	return nil
@@ -230,7 +240,32 @@ func (d *driver) Writer(ctx context.Context, path string, appendMode bool) (stor
 			if props.ContentLength == nil {
 				return nil, fmt.Errorf("missing ContentLength: %s", blobName)
 			}
+			if props.BlobType == nil {
+				return nil, fmt.Errorf("missing BlobType: %s", blobName)
+			}
+
 			size = *props.ContentLength
+			if *props.BlobType == blob.BlobTypeBlockBlob && size == 0 {
+				if props.ETag == nil {
+					return nil, fmt.Errorf("missing ETag for zero-length block blob: %s", blobName)
+				}
+				if _, err := blobRef.Delete(ctx, &blob.DeleteOptions{
+					AccessConditions: &blob.AccessConditions{
+						ModifiedAccessConditions: &blob.ModifiedAccessConditions{
+							IfMatch: props.ETag,
+						},
+					},
+				}); err != nil && !is404(err) {
+					return nil, fmt.Errorf("deleting zero-length block blob before append: %w", err)
+				}
+				res, err := d.client.NewAppendBlobClient(blobName).Create(ctx, nil)
+				if err != nil {
+					return nil, fmt.Errorf("creating new append blob: %w", err)
+				}
+				eTag = res.ETag
+			} else if *props.BlobType != blob.BlobTypeAppendBlob {
+				return nil, fmt.Errorf("cannot append to existing %s blob: %s", *props.BlobType, blobName)
+			}
 		} else {
 			if _, err := blobRef.Delete(ctx, nil); err != nil && !is404(err) {
 				return nil, fmt.Errorf("deleting existing blob before write: %w", err)
--- a/registry/storage/driver/azure/azure_test.go
+++ b/registry/storage/driver/azure/azure_test.go
@@ -1,13 +1,17 @@
 package azure

 import (
+	"bytes"
 	"context"
 	"math/rand"
 	"os"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"

+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming"
+	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/blob"
 	storagedriver "github.com/distribution/distribution/v3/registry/storage/driver"
 	"github.com/distribution/distribution/v3/registry/storage/driver/testsuites"
 )
@@ -174,72 +178,143 @@ func TestCommitAfterMove(t *testing.T) {
 	}
 }

-// TestPutContentOverExistingBlob regression-tests the migration check above
-// the AppendBlob Create call in PutContent. PutContent creates AppendBlobs,
-// and Azure returns 409 InvalidBlobType if Create is called over an existing
-// blob of a different type. The check must delete only blobs whose type does
-// not match what we're about to create.
-//
-// Two cases:
-//   - BlockBlob: written by an older version of the driver. Without the fix
-//     this failed with 409 InvalidBlobType.
-//   - AppendBlob: written by the current driver. Must still overwrite cleanly
-//     (one Put Blob, no delete) — the previous code reached this by deleting
-//     and recreating.
-func TestPutContentOverExistingBlob(t *testing.T) {
+func azureDriverForTest(t *testing.T) (storagedriver.StorageDriver, *driver) {
+	t.Helper()
 	skipCheck(t)

-	sd, err := azureDriverConstructor()
+	storageDriver, err := azureDriverConstructor()
 	if err != nil {
 		t.Fatalf("unexpected error creating azure driver: %v", err)
 	}
-	d := sd.(*Driver).Base.StorageDriver.(*driver)
-	ctx := context.Background()

-	tests := []struct {
-		name string
-		seed func(t *testing.T, blobName string)
-	}{
-		{
-			name: "BlockBlob",
-			seed: func(t *testing.T, blobName string) {
-				if _, err := d.client.NewBlockBlobClient(blobName).UploadBuffer(ctx, []byte("legacy"), nil); err != nil {
-					t.Fatalf("seeding BlockBlob: %v", err)
-				}
-			},
-		},
-		{
-			name: "AppendBlob",
-			seed: func(t *testing.T, blobName string) {
-				if err := sd.PutContent(ctx, "/regression/"+t.Name(), []byte("legacy")); err != nil {
-					t.Fatalf("seeding AppendBlob via PutContent: %v", err)
-				}
-			},
-		},
+	azureDriver, ok := storageDriver.(*Driver)
+	if !ok {
+		t.Fatalf("unexpected azure driver type: %T", storageDriver)
+	}
+	internalDriver, ok := azureDriver.baseEmbed.Base.StorageDriver.(*driver)
+	if !ok {
+		t.Fatalf("unexpected internal azure driver type: %T", azureDriver.baseEmbed.Base.StorageDriver)
 	}

-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			path := "/regression/" + t.Name()
-			blobName := d.blobName(path)
-			// nolint:errcheck
-			defer sd.Delete(ctx, path)
+	return storageDriver, internalDriver
+}

-			tc.seed(t, blobName)
+func assertBlobType(t *testing.T, internalDriver *driver, path string, expected blob.BlobType) {
+	t.Helper()

-			newContents := []byte("new content")
-			if err := sd.PutContent(ctx, path, newContents); err != nil {
-				t.Fatalf("PutContent over existing %s: %v", tc.name, err)
-			}
+	props, err := internalDriver.client.NewBlobClient(internalDriver.blobName(path)).GetProperties(context.Background(), nil)
+	if err != nil {
+		t.Fatalf("GetProperties(%s): unexpected error: %v", path, err)
+	}
+	if props.BlobType == nil {
+		t.Fatalf("GetProperties(%s): missing BlobType", path)
+	}
+	if *props.BlobType != expected {
+		t.Fatalf("GetProperties(%s): expected blob type %s, got %s", path, expected, *props.BlobType)
+	}
+}

-			got, err := sd.GetContent(ctx, path)
-			if err != nil {
-				t.Fatalf("GetContent after PutContent: %v", err)
-			}
-			if string(got) != string(newContents) {
-				t.Fatalf("GetContent: got %q, want %q", got, newContents)
-			}
-		})
+func TestPutContentCreatesBlockBlob(t *testing.T) {
+	storageDriver, internalDriver := azureDriverForTest(t)
+
+	ctx := context.Background()
+	path := "/putcontent-blockblob-" + randStringRunes(16)
+	contents := []byte("sha256:b708172ade643851b2b21c7b543bee28f732fd76d68e550b9dab16eb007ae4c0")
+
+	defer func() {
+		_ = storageDriver.Delete(ctx, path)
+	}()
+
+	if err := storageDriver.PutContent(ctx, path, contents); err != nil {
+		t.Fatalf("PutContent: unexpected error: %v", err)
+	}
+
+	assertBlobType(t, internalDriver, path, blob.BlobTypeBlockBlob)
+
+	got, err := storageDriver.GetContent(ctx, path)
+	if err != nil {
+		t.Fatalf("GetContent: unexpected error: %v", err)
+	}
+	if !bytes.Equal(got, contents) {
+		t.Fatalf("GetContent: expected %q, got %q", contents, got)
+	}
+}
+
+func TestPutContentMigratesLegacyAppendBlob(t *testing.T) {
+	storageDriver, internalDriver := azureDriverForTest(t)
+
+	ctx := context.Background()
+	path := "/putcontent-legacy-appendblob-" + randStringRunes(16)
+	contents := []byte("sha256:b708172ade643851b2b21c7b543bee28f732fd76d68e550b9dab16eb007ae4c0")
+	legacyContents := append(append([]byte{}, contents...), contents...)
+	appendBlobRef := internalDriver.client.NewAppendBlobClient(internalDriver.blobName(path))
+
+	defer func() {
+		_ = storageDriver.Delete(ctx, path)
+	}()
+
+	_ = storageDriver.Delete(ctx, path)
+	if _, err := appendBlobRef.Create(ctx, nil); err != nil {
+		t.Fatalf("AppendBlob.Create: unexpected error: %v", err)
+	}
+	if _, err := appendBlobRef.AppendBlock(ctx, streaming.NopCloser(bytes.NewReader(legacyContents)), nil); err != nil {
+		t.Fatalf("AppendBlob.AppendBlock: unexpected error: %v", err)
+	}
+
+	if err := storageDriver.PutContent(ctx, path, contents); err != nil {
+		t.Fatalf("PutContent: unexpected error: %v", err)
+	}
+
+	assertBlobType(t, internalDriver, path, blob.BlobTypeBlockBlob)
+
+	got, err := storageDriver.GetContent(ctx, path)
+	if err != nil {
+		t.Fatalf("GetContent: unexpected error: %v", err)
+	}
+	if !bytes.Equal(got, contents) {
+		t.Fatalf("GetContent: expected %q, got %q", contents, got)
+	}
+}
+
+func TestConcurrentPutContentSamePath(t *testing.T) {
+	storageDriver, internalDriver := azureDriverForTest(t)
+
+	ctx := context.Background()
+	path := "/putcontent-concurrent-" + randStringRunes(16)
+	contents := []byte("sha256:b708172ade643851b2b21c7b543bee28f732fd76d68e550b9dab16eb007ae4c0")
+
+	defer func() {
+		_ = storageDriver.Delete(ctx, path)
+	}()
+
+	const writers = 16
+	var wg sync.WaitGroup
+	errs := make(chan error, writers)
+
+	wg.Add(writers)
+	for i := 0; i < writers; i++ {
+		go func() {
+			defer wg.Done()
+			errs <- storageDriver.PutContent(ctx, path, contents)
+		}()
+	}
+	wg.Wait()
+	close(errs)
+
+	for err := range errs {
+		if err != nil {
+			t.Fatalf("PutContent: unexpected error: %v", err)
+		}
+	}
+
+	assertBlobType(t, internalDriver, path, blob.BlobTypeBlockBlob)
+
+	got, err := storageDriver.GetContent(ctx, path)
+	if err != nil {
+		t.Fatalf("GetContent: unexpected error: %v", err)
+	}
+	if !bytes.Equal(got, contents) {
+		t.Fatalf("GetContent: expected %q, got %q", contents, got)
 	}
 }