fix(sdk): stop GetManifest from splitting the key under the lock

sdk/chunked_test.go

@@ -1241,3 +1241,167 @@ func TestChunkedConcurrentDuplicateIndex(t *testing.T) {
 	}
 	assert.Equal(t, 1, won, "exactly one writer may claim an index")
 }
+
+// blockingSplitter signals when Split is entered and stays there until
+// released, so a test can observe what the writer holds while a split
+// is in flight. Only the first Split blocks; later ones pass straight
+// through.
+type blockingSplitter struct {
+	inner   KeySplitter
+	entered chan struct{}
+	release chan struct{}
+	once    sync.Once
+}
+
+func (s *blockingSplitter) Split(ctx context.Context, attrs []*policy.Value, dek []byte, defaultKAS *policy.SimpleKasKey) (*SplitResult, error) {
+	first := false
+	s.once.Do(func() {
+		first = true
+		close(s.entered)
+	})
+	if first {
+		<-s.release
+	}
+	return s.inner.Split(ctx, attrs, dek, defaultKAS)
+}
+
+// TestChunkedGetManifestDoesNotBlockWriteSegment pins the reason
+// GetManifest snapshots segment state and releases the lock before
+// splitting. A real splitter resolves KAS keys over the network; while
+// GetManifest held RLock across that call, every WriteSegment queued on
+// the write lock for its duration. This test deadlocks on the old
+// shape and passes on the new one.
+func TestChunkedGetManifestDoesNotBlockWriteSegment(t *testing.T) {
+	ctx := context.Background()
+	kasBundle := newChunkedFakeKAS(t)
+	defer kasBundle.server.Close()
+
+	splitter := &blockingSplitter{
+		inner:   DefaultKeySplitter(),
+		entered: make(chan struct{}),
+		release: make(chan struct{}),
+	}
+	w, err := NewChunkedWriter(ctx,
+		WithChunkedDefaultKAS(kasBundle.simpleKey()),
+		WithChunkedKeySplitter(splitter),
+	)
+	require.NoError(t, err)
+
+	_, err = w.WriteSegment(ctx, 0, []byte("first"))
+	require.NoError(t, err)
+
+	type manifestResult struct {
+		manifest *Manifest
+		err      error
+	}
+	manifests := make(chan manifestResult, 1)
+	go func() {
+		m, err := w.GetManifest(ctx)
+		manifests <- manifestResult{manifest: m, err: err}
+	}()
+
+	select {
+	case <-splitter.entered:
+	case <-time.After(10 * time.Second):
+		t.Fatal("GetManifest never reached the splitter")
+	}
+
+	wrote := make(chan error, 1)
+	go func() {
+		_, err := w.WriteSegment(ctx, 1, []byte("second"))
+		wrote <- err
+	}()
+	select {
+	case err := <-wrote:
+		require.NoError(t, err)
+	case <-time.After(10 * time.Second):
+		close(splitter.release)
+		t.Fatal("WriteSegment blocked while GetManifest was splitting")
+	}
+
+	close(splitter.release)
+	got := <-manifests
+	require.NoError(t, got.err)
+
+	// The manifest describes the writer as of the snapshot, not as of
+	// the return. Segment 1 landed after the lock was released, so it
+	// is deliberately absent -- GetManifest is a point-in-time view.
+	assert.Len(t, got.manifest.Segments, 1)
+
+	// The segment written during the split is still committed and shows
+	// up in the next call.
+	later, err := w.GetManifest(ctx)
+	require.NoError(t, err)
+	assert.Len(t, later.Segments, 2)
+}
+
+// blockingCipher blocks the first EncryptInPlace call whose input
+// matches blockOn, then delegates. Used to hold a WriteSegment call
+// inside its unlocked encrypt/sign/archive-write window so a
+// concurrent GetManifest can be exercised against an in-flight
+// reservation.
+type blockingCipher struct {
+	inner   segmentCipher
+	blockOn []byte
+	entered chan struct{}
+	release chan struct{}
+	once    sync.Once
+}
+
+func (c *blockingCipher) EncryptInPlace(data []byte) ([]byte, []byte, error) {
+	if bytes.Equal(data, c.blockOn) {
+		c.once.Do(func() {
+			close(c.entered)
+			<-c.release
+		})
+	}
+	return c.inner.EncryptInPlace(data)
+}
+
+// TestChunkedGetManifestExcludesInFlightReservation pins
+// segmentOrderLocked's handling of a reservation placeholder: a
+// concurrent GetManifest must snapshot only segments that have
+// actually landed, not error out because another goroutine's
+// WriteSegment has reserved-but-not-yet-committed an index.
+func TestChunkedGetManifestExcludesInFlightReservation(t *testing.T) {
+	ctx := context.Background()
+	blockOn := []byte("second")
+	cipher := &blockingCipher{blockOn: blockOn, entered: make(chan struct{}), release: make(chan struct{})}
+
+	writer, _ := newChunkedWriterForTest(ctx, t, withChunkedCipherFactory(func(dek []byte) (segmentCipher, error) {
+		inner, err := defaultSegmentCipherFactory(dek)
+		if err != nil {
+			return nil, err
+		}
+		cipher.inner = inner
+		return cipher, nil
+	}))
+
+	writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("first")})
+
+	wrote := make(chan error, 1)
+	go func() {
+		_, err := writer.WriteSegment(ctx, 1, blockOn)
+		wrote <- err
+	}()
+
+	select {
+	case <-cipher.entered:
+	case <-time.After(10 * time.Second):
+		t.Fatal("WriteSegment never reached the cipher")
+	}
+
+	// Segment 1's reservation is in flight but not yet committed; the
+	// snapshot must describe only segment 0, and must not error out
+	// treating the reservation as an unwritten-but-named segment.
+	manifest, err := writer.GetManifest(ctx)
+	require.NoError(t, err)
+	assert.Len(t, manifest.Segments, 1)
+
+	close(cipher.release)
+	require.NoError(t, <-wrote)
+
+	later, err := writer.GetManifest(ctx)
+	require.NoError(t, err)
+	assert.Len(t, later.Segments, 2)
+}

sdk/chunked_writer.go

@@ -453,7 +453,17 @@ func (w *chunkedWriter) Finalize(ctx context.Context, opts ...ChunkedFinalizeOpt
 		return nil, err
 	}
 
-	manifest, totalPlaintext, totalEncrypted, err := w.buildManifest(ctx, cfg)
+	// Finalize keeps the write lock across the split. It is terminal --
+	// no further WriteSegment may succeed after it returns -- so there is
+	// no concurrency to preserve, and releasing the lock to split would
+	// open a window where a segment lands in the archive after the
+	// snapshot that determines the manifest.
+	snap, err := w.snapshotLocked(cfg.keepSegments)
+	if err != nil {
+		return nil, err
+	}
+
+	manifest, totalPlaintext, totalEncrypted, err := w.buildManifest(ctx, cfg, snap)
 	if err != nil {
 		return nil, err
 	}
@@ -486,20 +496,36 @@ func (w *chunkedWriter) Finalize(ctx context.Context, opts ...ChunkedFinalizeOpt
 }
 
 // GetManifest returns the manifest snapshot.
+//
+// The lock is held only long enough to copy segment metadata; the key
+// split -- which may make network calls to resolve KAS keys -- runs
+// unlocked. Holding RLock across it would block every WriteSegment
+// waiting on the write lock for the duration of those calls, and
+// RWMutex bars new readers once a writer is queued, so concurrent
+// GetManifest calls would serialize behind it too.
 func (w *chunkedWriter) GetManifest(ctx context.Context, opts ...ChunkedFinalizeOption) (*Manifest, error) {
 	w.mu.RLock()
-	defer w.mu.RUnlock()
 	if w.closeFailed {
+		w.mu.RUnlock()
 		return nil, ErrChunkedCloseFailed
 	}
 	if w.finalized && w.manifest != nil {
-		return cloneChunkedManifest(w.manifest), nil
+		manifest := cloneChunkedManifest(w.manifest)
+		w.mu.RUnlock()
+		return manifest, nil
 	}
 	cfg, err := w.applyFinalizeOptions(opts)
 	if err != nil {
+		w.mu.RUnlock()
 		return nil, err
 	}
-	manifest, _, _, err := w.buildManifest(ctx, cfg)
+	snap, err := w.snapshotLocked(cfg.keepSegments)
+	w.mu.RUnlock()
+	if err != nil {
+		return nil, err
+	}
+
+	manifest, _, _, err := w.buildManifest(ctx, cfg, snap)
 	if err != nil {
 		return nil, err
 	}
@@ -657,12 +683,42 @@ func (w *chunkedWriter) applyFinalizeOptions(opts []ChunkedFinalizeOption) (*Chu
 
 // buildManifest composes the manifest from writer state, splits the
 // DEK, wraps splits into KAOs, and computes the root signature.
-func (w *chunkedWriter) buildManifest(ctx context.Context, cfg *ChunkedFinalizeConfig) (*Manifest, int64, int64, error) {
-	order, err := w.segmentOrderLocked(cfg.keepSegments)
+// chunkedSnapshot is the mutable writer state buildManifest needs,
+// copied out from under the lock. Segment values rather than the
+// pointers held in w.segments: WriteSegment mutates those in place when
+// the archive accepts a write, so reading them after the lock is
+// released would race.
+type chunkedSnapshot struct {
+	// segments are the per-segment metadata records in emission order.
+	segments []Segment
+}
+
+// snapshotLocked resolves the emission order and copies each segment's
+// metadata out of w.segments. The caller must hold w.mu for reading.
+func (w *chunkedWriter) snapshotLocked(keep []int) (*chunkedSnapshot, error) {
+	order, err := w.segmentOrderLocked(keep)
 	if err != nil {
-		return nil, 0, 0, err
+		return nil, err
+	}
+	segments := make([]Segment, len(order))
+	for i, idx := range order {
+		seg, ok := w.segments[idx]
+		if !ok || seg.Size < 0 {
+			return nil, fmt.Errorf("segment %d not written; cannot finalize", idx)
+		}
+		if seg.Hash == "" {
+			return nil, fmt.Errorf("segment %d has empty hash", idx)
+		}
+		segments[i] = *seg
 	}
+	return &chunkedSnapshot{segments: segments}, nil
+}
 
+// buildManifest assembles the manifest from a snapshot. It reads no
+// mutable writer state and holds no lock: every other field it touches
+// (dek, splitter, the integrity algorithms, useHex) is fixed at
+// construction.
+func (w *chunkedWriter) buildManifest(ctx context.Context, cfg *ChunkedFinalizeConfig, snap *chunkedSnapshot) (*Manifest, int64, int64, error) {
 	splits, err := w.splitter.Split(ctx, cfg.attributes, w.dek, cfg.defaultKAS)
 	if err != nil {
 		return nil, 0, 0, err
@@ -686,34 +742,25 @@ func (w *chunkedWriter) buildManifest(ctx context.Context, cfg *ChunkedFinalizeC
 		},
 		IntegrityInformation: IntegrityInformation{
 			SegmentHashAlgorithm: integrityAlgorithmString(w.segmentIntegrityAlgorithm),
-			Segments:             make([]Segment, len(order)),
+			Segments:             make([]Segment, len(snap.segments)),
 		},
 	}
 
 	var aggregate bytes.Buffer
 	var totalPlaintext, totalEncrypted int64
-	for i, idx := range order {
-		seg, ok := w.segments[idx]
-		if !ok || seg.Size < 0 {
-			return nil, 0, 0, fmt.Errorf("segment %d not written; cannot finalize", idx)
-		}
-		if seg.Hash == "" {
-			return nil, 0, 0, fmt.Errorf("segment %d has empty hash", idx)
-		}
-		encInfo.Segments[i] = *seg
+	for i, seg := range snap.segments {
+		encInfo.Segments[i] = seg
 		totalPlaintext += seg.Size
 		totalEncrypted += seg.EncryptedSize
 		decoded, err := ocrypto.Base64Decode([]byte(seg.Hash))
 		if err != nil {
-			return nil, 0, 0, fmt.Errorf("decode segment %d hash: %w", idx, err)
+			return nil, 0, 0, fmt.Errorf("decode segment %d hash: %w", i, err)
 		}
 		aggregate.Write(decoded)
 	}
-	if len(order) > 0 {
-		if first, ok := w.segments[order[0]]; ok {
-			encInfo.DefaultEncryptedSegSize = first.EncryptedSize
-			encInfo.DefaultSegmentSize = first.Size
-		}
+	if len(snap.segments) > 0 {
+		encInfo.DefaultEncryptedSegSize = snap.segments[0].EncryptedSize
+		encInfo.DefaultSegmentSize = snap.segments[0].Size
 	}
 
 	rootSig, err := calculateSignature(aggregate.Bytes(), w.dek, w.integrityAlgorithm, w.useHex)
@@ -767,9 +814,18 @@ func (w *chunkedWriter) buildManifest(ctx context.Context, cfg *ChunkedFinalizeC
 // stores segments sorted by index. Reordering would make the manifest
 // disagree with the bytes on disk; dropping a segment that has bytes
 // after it would shift every later segment's offset.
+//
+// An index with a reservation still in flight (Size < 0; see
+// WriteSegment) is treated as not written rather than as an error, so
+// a concurrent GetManifest/Finalize snapshots only the segments that
+// have actually landed instead of failing on one still being written
+// by another goroutine.
 func (w *chunkedWriter) segmentOrderLocked(keep []int) ([]int, error) {
 	written := make([]int, 0, len(w.segments))
-	for idx := range w.segments {
+	for idx, seg := range w.segments {
+		if seg.Size < 0 {
+			continue // reservation in flight; not yet written
+		}
 		written = append(written, idx)
 	}
 	sort.Ints(written)
@@ -783,7 +839,7 @@ func (w *chunkedWriter) segmentOrderLocked(keep []int) ([]int, error) {
 		if idx == written[i] {
 			continue
 		}
-		if _, ok := w.segments[idx]; !ok {
+		if seg, ok := w.segments[idx]; !ok || seg.Size < 0 {
 			return nil, fmt.Errorf("WithChunkedSegments references segment %d which was not written", idx)
 		}
 		return nil, fmt.Errorf(