From 74243e84b666062782ede258187758e866e3ec9e Mon Sep 17 00:00:00 2001 From: Dave Mihalcik Date: Mon, 31 Aug 2026 21:50:14 -0400 Subject: [PATCH 1/2] feat(sdk): add a chunked segment writer (experimental) Adds ChunkedWriter, a TDF creation path that accepts segments in any order: callers encrypt and upload each segment independently, then call Finalize to get the ZIP closing bytes. Also adds KeySplitter, the pluggable attribute-to-KAS-split seam Finalize uses. Review fixes folded in on top of the original implementation: - WriteSegment's rollback is now panic-safe (a committed flag plus a single deferred release, instead of an explicit release() call on each error path), so a panic unwinding through an injected cipher or archive-writer seam can no longer leak a segment-index reservation. - release() now also calls archiveWriter.CleanupSegment when the archive itself accepted part of a write before the segment failed, completing the rollback contract zipstream.SegmentWriter documents. - Finalize returns the new sentinel ErrChunkedCloseFailed if archiveWriter.Close fails after archiveWriter.Finalize already succeeded -- the archive is terminally finalized internally at that point regardless, so the writer must refuse further calls rather than look retryable. - WithChunkedIntegrityAlgorithm and WithChunkedSegmentIntegrityAlgorithm now reject any value other than HS256/GMAC instead of letting calculateSignature silently treat an unrecognized value as GMAC. - The default KeySplitter rejects an empty default-KAS URI, which previously propagated silently into KeyAccess.KasURL and produced a TDF with no rewrap endpoint. - Removed the dead maxSegmentIndex field (written, never read). - Documented that WithChunkedSegments trimming does not shrink the archive: CleanupSegment is never called for a dropped index, so the caller must still append its bytes when assembling the file. Second review round -- the data-loss and contradicted-doc findings: - A cancelled context at Finalize no longer bricks the writer. zipstream.Error gains a documented Mutated flag, set per return site in segmentWriter.Finalize against the SetOrder boundary via three constructors (finalizeErr, finalizeErrMutated, writeSegmentErr). chunkedWriter.Finalize pre-checks ctx.Err() before touching the archive and now fences only on a mutation -- or on an error it cannot classify, which fails safe. Context cancellation was the only production-reachable path into ErrChunkedFinalizeFailed, and it is exactly the case where discarding an already-uploaded payload is wrong. - Finalize refuses a racing WriteSegment with the new sentinel ErrChunkedWriteInFlight instead of emitting an archive whose correctness depends on the interleaving. zipstream has already counted the in-flight segment's bytes into the payload entry by the time Finalize can see it, so the trailer's offsets overshoot unless the caller appends bytes the manifest never mentions. The in-flight set is derived from the segment map rather than counted, so it cannot drift. GetManifest deliberately still allows it: a snapshot may be one segment short, an archive may not. - ChunkedFinalizeResult gains ArchiveEncryptedSize and ArchiveTotalSize. EncryptedSize/TotalSize describe the manifest after WithChunkedSegments trimming; the archive totals describe every written segment, which is what the caller must actually append. Sizing an upload from the former silently truncates. - SplitResult.VerifyReconstruction(dek) checks that the shares XOR back to the DEK and that each is the DEK's length -- the half of the splitter contract Validate cannot see. buildChunkedKeyAccessObjects calls it after Validate, before the archive is touched. Without it a splitter returning short or wrong shares produces a TDF that wraps, uploads and rewraps cleanly, then fails the reader's root signature check as if tampered with. - The DEK handed to KeySplitter.Split is now a clone, and the four slice-valued options clone what they are given. A splitter that reuses its argument as scratch space desynchronized the segment signatures from the root signature; a caller reusing an attributes slice could loosen the written policy after the fact. - Removed WithChunkedExcludeVersion and ErrChunkedVersionHexMismatch. The option was dead in both directions: a no-op in legacy mode and an unconditional error outside it, because omitting schemaVersion and hex-encoding signatures can only be chosen together, at construction, which is what WithChunkedTargetMode does. - Doc corrections where the comment contradicted the code: Finalize's duplicated error clause, the "nothing usable was produced" claim on the Close branch (a complete trailer was produced and is being discarded), the unreachable injection seams advertised on NewChunkedWriter, TDFData's single-use reader, GetManifest's overstated concurrency, "must carry a SigningKey", Split.Data and Split.ID, and Experimental banners on unexported types. The lock inversion between WriteSegment's defer and Finalize is now named at both sites. New tests: context cancellation is retryable at both Finalize and WriteSegment, Finalize rejects in-flight writes while GetManifest still succeeds in the same window, a non-mutating trailer failure leaves the writer usable (the existing fencing test is unchanged as the fail-safe guard), archive-vs-manifest totals under trimming, TestSplitResultVerifyReconstruction, DEK and attribute cloning, WithChunkedMimeType, zero-length payload and interior segments, and the manifest deep clone. Signed-off-by: Dave Mihalcik --- sdk/chunked_options.go | 247 ++ sdk/chunked_test.go | 2543 +++++++++++++++++ sdk/chunked_writer.go | 1193 ++++++++ sdk/internal/zipstream/segment_writer.go | 97 +- sdk/internal/zipstream/segment_writer_test.go | 213 ++ sdk/internal/zipstream/writer.go | 14 + sdk/key_splitter.go | 308 ++ sdk/key_splitter_test.go | 435 +++ 8 files changed, 5028 insertions(+), 22 deletions(-) create mode 100644 sdk/chunked_options.go create mode 100644 sdk/chunked_test.go create mode 100644 sdk/chunked_writer.go create mode 100644 sdk/key_splitter.go create mode 100644 sdk/key_splitter_test.go diff --git a/sdk/chunked_options.go b/sdk/chunked_options.go new file mode 100644 index 0000000000..b7272447ec --- /dev/null +++ b/sdk/chunked_options.go @@ -0,0 +1,247 @@ +package sdk + +import ( + "errors" + "fmt" + "io" + "slices" + + "github.com/opentdf/platform/protocol/go/policy" +) + +// Each injection-seam option below rejects nil rather than storing it. +// A nil seam is not detectable later: the config field is +// indistinguishable from "not set", so NewChunkedWriter installs no +// default and the nil is dereferenced during writing -- for the +// splitter, not until Finalize, long after the caller has encrypted +// every segment. + +// The slice-valued options below clone what they are given. Each is retained +// for the lifetime of the writer or of one Finalize call, and each determines +// something the caller cannot inspect afterwards -- the written policy, the +// signed assertions, the segments the manifest names. Keeping the caller's +// backing array would let a mutation after the call silently change any of +// them, including loosening the policy on data already encrypted. +// +// The clone stops at the slice, though, so it only closes the half of that +// hole reachable by reslicing or reassigning an element. What the elements +// point at stays shared: buildChunkedPolicy reads each *policy.Value at +// Finalize, and signAssertions signs with whatever AssertionKey.Key holds. +// Mutating a retained policy value still rewrites the bound policy, which is +// the loosening described above -- cloning the slice does not prevent it. +// Callers must treat everything they hand these options as frozen until +// Finalize returns. + +// withChunkedArchiveWriterFactory overrides the ZIP archive writer +// factory used by the [ChunkedWriter]. The factory must not be nil. +func withChunkedArchiveWriterFactory(f archiveWriterFactory) ChunkedWriterOption { + return func(c *chunkedWriterConfig) error { + if f == nil { + return errors.New("chunked: archive writer factory must not be nil") + } + c.archiveFactory = f + return nil + } +} + +// withChunkedCipherFactory overrides the segment cipher factory used +// by the [ChunkedWriter]. The factory must not be nil. +func withChunkedCipherFactory(f segmentCipherFactory) ChunkedWriterOption { + return func(c *chunkedWriterConfig) error { + if f == nil { + return errors.New("chunked: cipher factory must not be nil") + } + c.cipherFactory = f + return nil + } +} + +// withChunkedClock overrides the time source used by the +// [ChunkedWriter] and, through it, by the zipstream layer that stamps +// ZIP header timestamps. Tests inject fixedClock for deterministic +// output. The clock must not be nil. +func withChunkedClock(clock clock) ChunkedWriterOption { + return func(c *chunkedWriterConfig) error { + if clock == nil { + return errors.New("chunked: clock must not be nil") + } + c.clock = clock + return nil + } +} + +// WithChunkedInitialAttributes sets attribute values used by Finalize +// when the Finalize call does not supply its own. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedInitialAttributes(values []*policy.Value) ChunkedWriterOption { + return func(c *chunkedWriterConfig) error { + c.initialAttributes = slices.Clone(values) + return nil + } +} + +// WithChunkedDefaultKAS sets the default KAS used by Finalize when +// the Finalize call does not supply its own. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedDefaultKAS(kas *policy.SimpleKasKey) ChunkedWriterOption { + return func(c *chunkedWriterConfig) error { + c.initialDefaultKAS = kas + return nil + } +} + +// WithChunkedKeySplitter overrides the key splitter used by the +// [ChunkedWriter]. Callers with multi-KAS attribute grants should +// inject a splitter that understands their grant model. The splitter +// must not be nil. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedKeySplitter(splitter KeySplitter) ChunkedWriterOption { + return func(c *chunkedWriterConfig) error { + if splitter == nil { + return errors.New("chunked: key splitter must not be nil") + } + c.splitter = splitter + return nil + } +} + +// withChunkedRand overrides the entropy source used to generate the +// DEK. The reader must not be nil. +func withChunkedRand(r io.Reader) ChunkedWriterOption { + return func(c *chunkedWriterConfig) error { + if r == nil { + return errors.New("chunked: rand must not be nil") + } + c.rand = r + return nil + } +} + +// WithChunkedAssertions attaches signed assertions to the produced +// TDF. Each assertion is bound to the payload's aggregate hash, so +// they are signed at Finalize once every segment is in. Assertions +// without their own SigningKey are signed with HS256 over the DEK. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedAssertions(assertions []AssertionConfig) ChunkedFinalizeOption { + return func(c *chunkedFinalizeConfig) error { + c.assertions = slices.Clone(assertions) + return nil + } +} + +// WithChunkedAttributes overrides the writer's initial attributes for +// this Finalize call. +// +// An empty or nil slice does not clear the writer's initial +// attributes -- it reads as "not specified" and the initial ones still +// apply. There is deliberately no way to finalize with no attributes +// once the writer was constructed with some: dropping attributes +// silently would loosen the policy on the data, which is the one +// mistake here that cannot be detected after the fact. Construct a +// writer without WithChunkedInitialAttributes instead. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedAttributes(values []*policy.Value) ChunkedFinalizeOption { + return func(c *chunkedFinalizeConfig) error { + c.attributes = slices.Clone(values) + return nil + } +} + +// WithChunkedDefaultKASForFinalize overrides the writer's initial +// default KAS for this Finalize call. +// +// A nil argument reads as "not specified", so the writer's initial +// default KAS still applies; there is no way to unset it for one call. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedDefaultKASForFinalize(kas *policy.SimpleKasKey) ChunkedFinalizeOption { + return func(c *chunkedFinalizeConfig) error { + c.defaultKAS = kas + return nil + } +} + +// WithChunkedEncryptedMetadata attaches AES-GCM-encrypted metadata to +// every KAO in the TDF. The metadata is keyed on the split share and +// only decryptable by a reader that has been granted access. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedEncryptedMetadata(metadata string) ChunkedFinalizeOption { + return func(c *chunkedFinalizeConfig) error { + c.encryptedMetadata = metadata + return nil + } +} + +// WithChunkedTargetMode targets a specific TDF spec version, given as +// a semver string such as "4.2.2". An empty mode selects the most +// recently available target version (4.3.0). +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedTargetMode(mode string) ChunkedWriterOption { + return func(c *chunkedWriterConfig) error { + if mode == "" { + c.useHex = false + c.excludeVersion = false + return nil + } + legacy, err := isLessThanSemver(mode, hexSemverThreshold) + if err != nil { + return fmt.Errorf("target mode %q: %w", mode, err) + } + c.useHex = legacy + c.excludeVersion = legacy + return nil + } +} + +// WithChunkedMimeType records the payload MIME type in the manifest. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedMimeType(mimeType string) ChunkedFinalizeOption { + return func(c *chunkedFinalizeConfig) error { + c.mimeType = mimeType + return nil + } +} + +// WithChunkedSegments sets the segments the finalized manifest +// describes. Passing no indices emits every written segment in +// ascending index order, which is what most callers want. +// +// The indices need not be contiguous. A caller that reserves a fixed +// block of indices per upload part -- part N owning +// [N*stride, (N+1)*stride) -- and fills only part of each block writes +// a sparse index set by construction; listing it here is fine. +// +// What the indices must be is a prefix of the written segments in +// ascending index order: they may drop from the end, but may not +// reorder or skip. The archive stores segments sorted by index, so a +// manifest that reorders them would not describe the bytes on disk, +// and one that skips a segment with bytes after it would misread every +// segment that follows. For the same reason the caller must +// concatenate each segment's TDFData in ascending index order. +// +// Dropping a segment from the manifest does not shrink the archive: +// the underlying writer never rolls back a completed segment's +// contribution to the payload's recorded size and CRC, so every +// segment that was actually written -- including ones this option +// excludes from the manifest -- must still be appended by the caller +// when assembling the final file. Skipping a dropped segment's bytes +// produces an archive whose central directory offsets overshoot. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedSegments(indices []int) ChunkedFinalizeOption { + return func(c *chunkedFinalizeConfig) error { + // Cloned because segmentOrderLocked validates the live slice before + // copying it, so a concurrent mutation could slip between the check + // and the use. + c.keepSegments = slices.Clone(indices) + return nil + } +} diff --git a/sdk/chunked_test.go b/sdk/chunked_test.go new file mode 100644 index 0000000000..8f5842afab --- /dev/null +++ b/sdk/chunked_test.go @@ -0,0 +1,2543 @@ +package sdk + +import ( + "archive/zip" + "bytes" + "context" + "crypto/ecdsa" + "crypto/rand" + "crypto/subtle" + "crypto/x509" + "encoding/hex" + "encoding/json" + "encoding/pem" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "slices" + "sync" + "testing" + "time" + + "connectrpc.com/connect" + "github.com/lestrrat-go/jwx/v2/jwt" + "github.com/opentdf/platform/lib/ocrypto" + kaspb "github.com/opentdf/platform/protocol/go/kas" + "github.com/opentdf/platform/protocol/go/kas/kasconnect" + "github.com/opentdf/platform/protocol/go/policy" + "github.com/opentdf/platform/sdk/internal/zipstream" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/encoding/protojson" +) + +// TestChunkedRoundTrip writes segments through NewChunkedWriter and +// reads the resulting TDF back through the mainline SDK.LoadTDF path +// (single-KAS, RSA-2048), verifying end-to-end interop. +func TestChunkedRoundTrip(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t) + + chunks := [][]byte{[]byte("hello, "), []byte("chunked "), []byte("world!")} + body := writeChunkedSegments(ctx, t, writer, chunks) + + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + require.NotNil(t, fin.Manifest) + + tdfBytes := bytes.Join([][]byte{body, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("hello, chunked world!"), plain) +} + +// TestChunkedKeepSegments verifies WithChunkedSegments trims trailing +// segments from the manifest and the mainline reader decrypts only the +// retained ones. +func TestChunkedKeepSegments(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t) + + body := writeChunkedSegments(ctx, t, writer, [][]byte{ + []byte("keep-0-"), []byte("keep-1-"), []byte("drop-2!"), + }) + fin, err := writer.Finalize(ctx, WithChunkedSegments([]int{0, 1})) + require.NoError(t, err) + require.Len(t, fin.Manifest.Segments, 2) + + tdfBytes := bytes.Join([][]byte{body, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("keep-0-keep-1-"), plain) +} + +// TestChunkedKeepSegmentsRequiresDroppedBytesAppended pins the +// invariant WithChunkedSegments documents: dropping a segment from +// the manifest does not shrink the archive, because CleanupSegment is +// never called for a trimmed index, so the archive's recorded size +// and CRC already include it. Omitting a dropped segment's bytes when +// assembling the file must not silently produce a readable TDF. +func TestChunkedKeepSegmentsRequiresDroppedBytesAppended(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t) + + seg0, err := writer.WriteSegment(ctx, 0, []byte("keep-0-")) + require.NoError(t, err) + seg0Bytes, err := io.ReadAll(seg0.TDFData) + require.NoError(t, err) + + seg1, err := writer.WriteSegment(ctx, 1, []byte("keep-1-")) + require.NoError(t, err) + seg1Bytes, err := io.ReadAll(seg1.TDFData) + require.NoError(t, err) + + // Written but dropped from the manifest below -- and its bytes are + // omitted from the assembled file too, the mistake the doc now + // warns against. + _, err = writer.WriteSegment(ctx, 2, []byte("drop-2!")) + require.NoError(t, err) + + fin, err := writer.Finalize(ctx, WithChunkedSegments([]int{0, 1})) + require.NoError(t, err) + require.Len(t, fin.Manifest.Segments, 2) + + shortBytes := bytes.Join([][]byte{seg0Bytes, seg1Bytes, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(shortBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + if err == nil { + _, err = io.ReadAll(reader) + } + require.Error(t, err, "an archive missing a written-but-dropped segment's bytes must not read back cleanly") +} + +// TestChunkedFinalizeSignsAssertions verifies assertions supplied to +// Finalize land in the manifest signed with the default HS256-over-DEK +// key, and that the mainline reader verifies them on the way back out. +func TestChunkedFinalizeSignsAssertions(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t) + + body := writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("asserted payload")}) + + fin, err := writer.Finalize(ctx, WithChunkedAssertions([]AssertionConfig{{ + ID: "a", + Type: BaseAssertion, + Scope: PayloadScope, + AppliesToState: Unencrypted, + Statement: Statement{Format: "json", Schema: "urn:test", Value: `{"k":"v"}`}, + }})) + require.NoError(t, err) + + require.Len(t, fin.Manifest.Assertions, 1) + got := fin.Manifest.Assertions[0] + assert.Equal(t, "a", got.ID) + assert.Equal(t, JWS.String(), got.Binding.Method) + assert.NotEmpty(t, got.Binding.Signature) + + // The reader recomputes the aggregate hash and re-verifies the + // binding, so a round trip is the real check that the assertion was + // bound to this payload and not merely well-formed. + tdfBytes := bytes.Join([][]byte{body, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("asserted payload"), plain) +} + +// TestChunkedOutOfOrderWrites exercises the core value proposition of +// ChunkedWriter: segments may be written in any order provided the +// caller concatenates TDFData in index order before Finalize.Data. +func TestChunkedOutOfOrderWrites(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t) + + chunks := [][]byte{[]byte("aaa-"), []byte("bbb-"), []byte("ccc-"), []byte("ddd!")} + + // Write in scrambled order: 2, 0, 3, 1. + segBytes := make([][]byte, len(chunks)) + for _, idx := range []int{2, 0, 3, 1} { + seg, err := writer.WriteSegment(ctx, idx, chunks[idx]) + require.NoError(t, err) + buf, err := io.ReadAll(seg.TDFData) + require.NoError(t, err) + segBytes[idx] = buf + } + + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + require.Equal(t, 4, fin.TotalSegments) + + // Concat in INDEX order (segment 0 carries the ZIP local header). + var body bytes.Buffer + for _, buf := range segBytes { + body.Write(buf) + } + body.Write(fin.Data) + + reader, err := s.LoadTDF(bytes.NewReader(body.Bytes()), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("aaa-bbb-ccc-ddd!"), plain) +} + +// TestChunkedClockThreadedToZipHeaders verifies withChunkedClock is +// threaded into the zipstream layer so every ZIP entry ModTime +// stamps from the injected clock rather than time.Now. This is the +// invariant that enables byte-deterministic ZIP headers for xtest +// fixtures (DEK / session-key randomness still varies the payload +// and KAS-wrap ciphertexts, which is not the scope of this test). +func TestChunkedClockThreadedToZipHeaders(t *testing.T) { + ctx := context.Background() + + // Pick a 2-second-aligned instant so DOS timestamp truncation is + // a no-op. + pinned := time.Date(2026, 6, 15, 12, 0, 0, 0, time.UTC) + w, _ := newChunkedWriterForTest(ctx, t, withChunkedClock(fixedClock{T: pinned})) + + body := writeChunkedSegments(ctx, t, w, [][]byte{[]byte("payload-abc")}) + fin, err := w.Finalize(ctx) + require.NoError(t, err) + + tdfBytes := bytes.Join([][]byte{body, fin.Data}, nil) + zr, err := zip.NewReader(bytes.NewReader(tdfBytes), int64(len(tdfBytes))) + require.NoError(t, err) + require.NotEmpty(t, zr.File) + + for _, f := range zr.File { + // archive/zip normalises DOS timestamps to the local zone; compare in UTC. + assert.Equal(t, pinned, f.Modified.UTC(), + "entry %q ModTime must match injected clock", f.Name) + } +} + +// TestChunkedRejectsInvalidSequencing pins the sentinel each of +// WriteSegment and Finalize returns for a call that is out of sequence +// with the writer's lifecycle. +func TestChunkedRejectsInvalidSequencing(t *testing.T) { + ctx := context.Background() + + for _, tc := range []struct { + name string + provoke func(t *testing.T, w ChunkedWriter) error + wantErr error + }{ + { + name: "duplicate segment", + provoke: func(t *testing.T, w ChunkedWriter) error { + t.Helper() + _, err := w.WriteSegment(ctx, 0, []byte("first")) + require.NoError(t, err) + _, err = w.WriteSegment(ctx, 0, []byte("second")) + return err + }, + wantErr: ErrChunkedSegmentAlreadyWritten, + }, + { + name: "negative index", + provoke: func(_ *testing.T, w ChunkedWriter) error { + _, err := w.WriteSegment(ctx, -1, []byte("x")) + return err + }, + wantErr: ErrChunkedInvalidSegmentIndex, + }, + { + name: "write after finalize", + provoke: func(t *testing.T, w ChunkedWriter) error { + t.Helper() + _, err := w.WriteSegment(ctx, 0, []byte("x")) + require.NoError(t, err) + _, err = w.Finalize(ctx) + require.NoError(t, err) + _, err = w.WriteSegment(ctx, 1, []byte("late")) + return err + }, + wantErr: ErrChunkedAlreadyFinalized, + }, + { + name: "double finalize", + provoke: func(t *testing.T, w ChunkedWriter) error { + t.Helper() + _, err := w.WriteSegment(ctx, 0, []byte("x")) + require.NoError(t, err) + _, err = w.Finalize(ctx) + require.NoError(t, err) + _, err = w.Finalize(ctx) + return err + }, + wantErr: ErrChunkedAlreadyFinalized, + }, + } { + t.Run(tc.name, func(t *testing.T) { + w, _ := newChunkedWriterForTest(ctx, t) + require.ErrorIs(t, tc.provoke(t, w), tc.wantErr) + }) + } +} + +// TestChunkedKeepSegmentsSparse verifies WithChunkedSegments accepts a +// sparse index set. This is the S3 multipart shape: each upload part +// reserves a fixed block of indices and fills only the front of it, so +// the written indices have large gaps but are still emitted in order. +func TestChunkedKeepSegmentsSparse(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + w, kasBundle := newChunkedWriterForTest(ctx, t) + + const stride = 5000 + chunks := map[int][]byte{ + 0: []byte("part1-a-"), + 1: []byte("part1-b-"), + stride: []byte("part2-a-"), + stride + 1: []byte("part2-b"), + } + indices := []int{0, 1, stride, stride + 1} + + // Write out of order to prove the index, not the call order, + // determines layout. + encrypted := make(map[int][]byte, len(indices)) + for _, idx := range []int{stride, 0, stride + 1, 1} { + seg, err := w.WriteSegment(ctx, idx, chunks[idx]) + require.NoError(t, err) + buf, err := io.ReadAll(seg.TDFData) + require.NoError(t, err) + encrypted[idx] = buf + } + + // Concatenate in ascending index order, as the contract requires. + var body bytes.Buffer + for _, idx := range indices { + body.Write(encrypted[idx]) + } + + fin, err := w.Finalize(ctx, WithChunkedSegments(indices)) + require.NoError(t, err) + require.Len(t, fin.Manifest.Segments, len(indices)) + + tdfBytes := bytes.Join([][]byte{body.Bytes(), fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("part1-a-part1-b-part2-a-part2-b"), plain) +} + +// TestChunkedKeepSegmentsRejects pins the validation Finalize applies to +// WithChunkedSegments: the keep list must name only written indices, in +// strictly ascending order, with no gaps that would shift a later +// segment's offset. +func TestChunkedKeepSegmentsRejects(t *testing.T) { + ctx := context.Background() + + for _, tc := range []struct { + name string + writes []int + keep []int + wantErrSub string + }{ + // Dropping segment 1 while keeping 2 would shift 2's offset and + // make the payload unreadable. + {"skips a written segment", []int{0, 1, 2}, []int{0, 2}, "ascending index order"}, + // Rejected even though every named index was written. + {"descending order", []int{0, 1}, []int{1, 0}, "ascending index order"}, + {"names an unwritten index", []int{0, 5}, []int{0, 1}, "not written"}, + {"longer than the written set", []int{0}, []int{0, 1}, "only 1 were written"}, + // A repeat cannot be ascending, so it fails the ordering rule + // rather than needing its own check. + {"names an index twice", []int{0, 1}, []int{0, 0}, "ascending index order"}, + // WriteSegment never accepts a negative index, so it can only + // ever be reported as unwritten. + {"negative index", []int{0, 1}, []int{-1, 0}, "not written"}, + } { + t.Run(tc.name, func(t *testing.T) { + w, _ := newChunkedWriterForTest(ctx, t) + for _, idx := range tc.writes { + _, err := w.WriteSegment(ctx, idx, []byte("x")) + require.NoError(t, err) + } + _, err := w.Finalize(ctx, WithChunkedSegments(tc.keep)) + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantErrSub) + }) + } +} + +// TestChunkedFinalizeRequiresSegmentZero pins the sentinel Finalize +// returns for a write set that omits segment 0, and that the rejection +// leaves the writer usable: the caller's only recovery is to write the +// missing segment and finalize again. +// +// The write set here -- indices 5 and 6, nothing lower -- is what a +// caller gets if it reserves a block of indices per upload part and +// part 0 never runs, or if it simply numbers parts from 1. Only +// segment 0 emits the payload's ZIP local file header (see +// zipstream.segmentWriter.WriteSegment), so the assembled bytes would +// not be a ZIP container at all; refusing is the only safe answer, +// since by Finalize the caller has already uploaded what it encrypted. +func TestChunkedFinalizeRequiresSegmentZero(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + w, kasBundle := newChunkedWriterForTest(ctx, t) + + encrypted := make(map[int][]byte, 3) + write := func(index int, chunk string) { + t.Helper() + seg, err := w.WriteSegment(ctx, index, []byte(chunk)) + require.NoError(t, err) + buf, err := io.ReadAll(seg.TDFData) + require.NoError(t, err) + encrypted[index] = buf + } + + write(5, "hello-") + write(6, "world!") + + _, err := w.Finalize(ctx) + require.ErrorIs(t, err, ErrChunkedMissingSegmentZero) + + // A rejected Finalize must not consume the writer, or the caller has + // no way back: the segments it already encrypted would be stranded. + write(0, "zero-") + fin, err := w.Finalize(ctx) + require.NoError(t, err) + require.Len(t, fin.Manifest.Segments, 3) + + var body bytes.Buffer + for _, idx := range []int{0, 5, 6} { + body.Write(encrypted[idx]) + } + body.Write(fin.Data) + + reader, err := s.LoadTDF(bytes.NewReader(body.Bytes()), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("zero-hello-world!"), plain) +} + +// TestChunkedFinalizeWithNoSegments verifies an untouched writer fails +// with the same sentinel rather than the archive layer's wrapped +// "segment missing". +func TestChunkedFinalizeWithNoSegments(t *testing.T) { + ctx := context.Background() + w, _ := newChunkedWriterForTest(ctx, t) + + _, err := w.Finalize(ctx) + require.ErrorIs(t, err, ErrChunkedMissingSegmentZero) +} + +// TestChunkedGetManifestWithoutSegmentZero guards the placement of the +// segment-0 check. GetManifest shares buildManifest with Finalize but +// is a pre-finalize snapshot, so it must keep working while segment 0 +// is still outstanding. +func TestChunkedGetManifestWithoutSegmentZero(t *testing.T) { + ctx := context.Background() + w, _ := newChunkedWriterForTest(ctx, t) + + _, err := w.WriteSegment(ctx, 5, []byte("hello-")) + require.NoError(t, err) + _, err = w.WriteSegment(ctx, 6, []byte("world!")) + require.NoError(t, err) + + snap, err := w.GetManifest(ctx) + require.NoError(t, err) + assert.Len(t, snap.Segments, 2) +} + +// TestChunkedGetManifestBeforeFinalize verifies GetManifest returns a +// snapshot of the currently-written segments prior to Finalize and +// the frozen manifest afterwards. +func TestChunkedGetManifestBeforeFinalize(t *testing.T) { + ctx := context.Background() + w, _ := newChunkedWriterForTest(ctx, t) + + _, err := w.WriteSegment(ctx, 0, []byte("first")) + require.NoError(t, err) + _, err = w.WriteSegment(ctx, 1, []byte("second")) + require.NoError(t, err) + + snap, err := w.GetManifest(ctx) + require.NoError(t, err) + require.NotNil(t, snap) + assert.Len(t, snap.Segments, 2) + + fin, err := w.Finalize(ctx) + require.NoError(t, err) + + frozen, err := w.GetManifest(ctx) + require.NoError(t, err) + assert.Equal(t, fin.Manifest.Method.Algorithm, frozen.Method.Algorithm) + assert.Len(t, frozen.Segments, 2) +} + +// writeChunkedSegments writes each element of segments as an ordered +// segment and returns the concatenated ciphertext produced by the +// writer. +func writeChunkedSegments(ctx context.Context, t *testing.T, w ChunkedWriter, segments [][]byte) []byte { + t.Helper() + var body bytes.Buffer + for i, chunk := range segments { + seg, err := w.WriteSegment(ctx, i, chunk) + require.NoError(t, err) + _, err = io.Copy(&body, seg.TDFData) + require.NoError(t, err) + } + return body.Bytes() +} + +// chunkedFakeKAS bundles an in-process RSA-2048 KAS + the httptest +// server it is registered on. Rewrap only handles the "wrapped" +// (RSA-OAEP) KeyType — matches what DefaultKeySplitter emits +// against an RSA-2048 KAS public key. +type chunkedFakeKAS struct { + kasconnect.UnimplementedAccessServiceHandler + privatePEM string + publicPEM string + kid string + url string + server *httptest.Server +} + +// newChunkedFakeKAS starts an httptest server hosting a fake KAS with +// a freshly-generated RSA-2048 keypair. +func newChunkedFakeKAS(t *testing.T) *chunkedFakeKAS { + t.Helper() + pair, err := ocrypto.NewRSAKeyPair(2048) + require.NoError(t, err) + pubPEM, err := pair.PublicKeyInPemFormat() + require.NoError(t, err) + privPEM, err := pair.PrivateKeyInPemFormat() + require.NoError(t, err) + + kas := &chunkedFakeKAS{ + privatePEM: privPEM, + publicPEM: pubPEM, + kid: "chunked-test-kid", + } + mux := http.NewServeMux() + path, handler := kasconnect.NewAccessServiceHandler(kas) + mux.Handle(path, handler) + kas.server = httptest.NewServer(mux) + kas.url = kas.server.URL + return kas +} + +// newChunkedWriterForTest starts a fake RSA-2048 KAS, registers its +// shutdown via t.Cleanup, and constructs a ChunkedWriter against it +// with opts layered on top of the default KAS option. Every case that +// expects NewChunkedWriter to succeed shares this setup; callers that +// need construction itself to fail build the fake KAS and call +// NewChunkedWriter directly instead. +func newChunkedWriterForTest(ctx context.Context, t *testing.T, opts ...ChunkedWriterOption) (ChunkedWriter, *chunkedFakeKAS) { + t.Helper() + kasBundle := newChunkedFakeKAS(t) + t.Cleanup(kasBundle.server.Close) + + all := append([]ChunkedWriterOption{WithChunkedDefaultKAS(kasBundle.simpleKey())}, opts...) + w, err := NewChunkedWriter(ctx, all...) + require.NoError(t, err) + return w, kasBundle +} + +// Rewrap unwraps every RSA-wrapped KAO under the KAS private key and +// re-wraps under the caller's session public key. +// +// It verifies the policy binding before doing so, as a real KAS does. That +// check is the whole reason the binding exists: it is what stops a holder of +// one TDF from swapping in a looser policy and replaying the KAOs against it. +// A fake that skipped it would accept a writer that bound the policy under the +// wrong key, or in the wrong encoding, and the round-trip tests would still +// pass -- the binding is never consulted on the decrypt path. +func (k *chunkedFakeKAS) Rewrap(_ context.Context, in *connect.Request[kaspb.RewrapRequest]) (*connect.Response[kaspb.RewrapResponse], error) { + tok, err := jwt.ParseInsecure([]byte(in.Msg.GetSignedRequestToken())) + if err != nil { + return nil, fmt.Errorf("parse jwt: %w", err) + } + rawBody, ok := tok.Get("requestBody") + if !ok { + return nil, errors.New("requestBody missing from token") + } + bodyStr, ok := rawBody.(string) + if !ok { + return nil, errors.New("requestBody not a string") + } + body := kaspb.UnsignedRewrapRequest{} + if err := protojson.Unmarshal([]byte(bodyStr), &body); err != nil { + return nil, fmt.Errorf("unmarshal request body: %w", err) + } + + dec, err := ocrypto.FromPrivatePEM(k.privatePEM) + if err != nil { + return nil, fmt.Errorf("kas priv: %w", err) + } + enc, err := ocrypto.FromPublicPEM(body.GetClientPublicKey()) + if err != nil { + return nil, fmt.Errorf("client pub: %w", err) + } + + resp := &kaspb.RewrapResponse{} + for _, req := range body.GetRequests() { + policyResult := &kaspb.PolicyRewrapResult{PolicyId: req.GetPolicy().GetId()} + for _, kaoReq := range req.GetKeyAccessObjects() { + kao := kaoReq.GetKeyAccessObject() + if kao.GetKeyType() != kWrapped { + return nil, fmt.Errorf("unsupported key type %q", kao.GetKeyType()) + } + share, err := dec.Decrypt(kao.GetWrappedKey()) + if err != nil { + return nil, fmt.Errorf("unwrap: %w", err) + } + // Checked against the share, not the DEK: each KAO binds the + // policy under its own split, which is what makes a binding + // non-transferable between KAOs. + if err := verifyChunkedPolicyBinding(kao.GetPolicyBinding(), share, req.GetPolicy().GetBody()); err != nil { + return nil, fmt.Errorf("kao %q: %w", kaoReq.GetKeyAccessObjectId(), err) + } + wrapped, err := enc.Encrypt(share) + if err != nil { + return nil, fmt.Errorf("rewrap: %w", err) + } + policyResult.Results = append(policyResult.Results, &kaspb.KeyAccessRewrapResult{ + Result: &kaspb.KeyAccessRewrapResult_KasWrappedKey{KasWrappedKey: wrapped}, + Status: "permit", + KeyAccessObjectId: kaoReq.GetKeyAccessObjectId(), + }) + } + resp.Responses = append(resp.Responses, policyResult) + } + return connect.NewResponse(resp), nil +} + +// verifyChunkedPolicyBinding recomputes the HMAC a KAS checks before it will +// release a key share, and compares it constant-time to the one the KAO +// carries. +func verifyChunkedPolicyBinding(binding *kaspb.PolicyBinding, share []byte, base64Policy string) error { + if binding == nil { + return errors.New("policy binding missing") + } + // A KAS dispatches on the algorithm rather than assuming HMAC; an unknown + // value is a refusal, not something to verify past. + if binding.GetAlgorithm() != hmacIntegrityAlgorithm { + return fmt.Errorf("unsupported policy binding algorithm %q", binding.GetAlgorithm()) + } + want := string(ocrypto.Base64Encode([]byte( + hex.EncodeToString(ocrypto.CalculateSHA256Hmac(share, []byte(base64Policy))), + ))) + if subtle.ConstantTimeCompare([]byte(want), []byte(binding.GetHash())) != 1 { + return errors.New("policy binding does not match the policy") + } + return nil +} + +// simpleKey returns the KAS descriptor the writer accepts. +func (k *chunkedFakeKAS) simpleKey() *policy.SimpleKasKey { + return &policy.SimpleKasKey{ + KasUri: k.url, + PublicKey: &policy.SimpleKasPublicKey{ + Algorithm: policy.Algorithm_ALGORITHM_RSA_2048, + Kid: k.kid, + Pem: k.publicPEM, + }, + } +} + +// newChunkedTestSDK builds a minimal SDK value for these tests. It is +// constructed from package-private fields to skip New()'s +// platform-lookup requirement, since LoadTDF only needs conn and +// tokenSource. +// +// Deliberately not wired to the fake KAS: the SDK never learns the KAS +// address. Each key access object carries its own URL, so the reader +// reaches the fake through the manifest. Pass the fake's URL to +// WithChunkedDefaultKAS when writing and WithKasAllowlist when +// reading. +func newChunkedTestSDK(t *testing.T) SDK { + t.Helper() + ats := getTokenSource(t) + return SDK{ + conn: &ConnectRPCConnection{Client: http.DefaultClient}, + tokenSource: ats, + } +} + +// TestChunkedKAOShape pins the key access object fields the chunked +// writer emits, so they cannot silently drift from the ones +// SDK.CreateTDF produces via the shared createKeyAccess helper. +func TestChunkedKAOShape(t *testing.T) { + ctx := context.Background() + writer, kasBundle := newChunkedWriterForTest(ctx, t) + + writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("payload")}) + fin, err := writer.Finalize(ctx, WithChunkedEncryptedMetadata("meta")) + require.NoError(t, err) + require.Len(t, fin.Manifest.KeyAccessObjs, 1) + + kao := fin.Manifest.KeyAccessObjs[0] + assert.Equal(t, kWrapped, kao.KeyType) + assert.Equal(t, kKasProtocol, kao.Protocol) + assert.Equal(t, keyAccessSchemaVersion, kao.SchemaVersion) + assert.Equal(t, kasBundle.url, kao.KasURL) + assert.Equal(t, kasBundle.kid, kao.KID) + assert.NotEmpty(t, kao.WrappedKey) + assert.NotEmpty(t, kao.EncryptedMetadata) + + binding, ok := kao.PolicyBinding.(PolicyBinding) + require.True(t, ok, "policy binding should be a PolicyBinding, got %T", kao.PolicyBinding) + assert.Equal(t, hmacIntegrityAlgorithm, binding.Alg) + assert.NotEmpty(t, binding.Hash) +} + +// TestChunkedECKeyAccess covers the EC wrapping path, which the +// round-trip tests miss because the fake KAS is RSA-only. It asserts +// the manifest key type is the one the real KAS dispatches on +// ("ec-wrapped", not "eccWrapped") and that the wrapped key actually +// decrypts under the KAS private key using the AES-GCM envelope the +// KAS rewrap path expects. +func TestChunkedECKeyAccess(t *testing.T) { + pair, err := ocrypto.NewECKeyPair(ocrypto.ECCModeSecp256r1) + require.NoError(t, err) + pubPEM, err := pair.PublicKeyInPemFormat() + require.NoError(t, err) + privPEM, err := pair.PrivateKeyInPemFormat() + require.NoError(t, err) + + const kasURL = "https://kas.example.com" + dek := make([]byte, kKeySize) + for i := range dek { + dek[i] = byte(i) + } + splits := &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{ + kasURL: { + Algorithm: ocrypto.EC256Key, + KID: "ec-kid", + PEM: pubPEM, + URL: kasURL, + }, + }, + Splits: []Split{{Data: dek, KASURLs: []string{kasURL}}}, + } + + kaos, err := buildChunkedKeyAccessObjects(splits, dek, []byte(`{"uuid":"test"}`), "") + require.NoError(t, err) + require.Len(t, kaos, 1) + + kao := kaos[0] + assert.Equal(t, kECWrapped, kao.KeyType, "KAS dispatches on this exact string") + require.NotEmpty(t, kao.EphemeralPublicKey) + + // Unwrap the way service/kas/access/rewrap.go does for "ec-wrapped". + keySize, err := ocrypto.GetECKeySize([]byte(kao.EphemeralPublicKey)) + require.NoError(t, err) + mode, err := ocrypto.ECSizeToMode(keySize) + require.NoError(t, err) + + block, _ := pem.Decode([]byte(kao.EphemeralPublicKey)) + require.NotNil(t, block) + pub, err := x509.ParsePKIXPublicKey(block.Bytes) + require.NoError(t, err) + ecPub, ok := pub.(*ecdsa.PublicKey) + require.True(t, ok) + compressed, err := ocrypto.CompressedECPublicKey(mode, *ecPub) + require.NoError(t, err) + + priv, err := ocrypto.ECPrivateKeyFromPem([]byte(privPEM)) + require.NoError(t, err) + dec, err := ocrypto.NewSaltedECDecryptor(priv, tdfSalt(), nil) + require.NoError(t, err) + + wrapped, err := ocrypto.Base64Decode([]byte(kao.WrappedKey)) + require.NoError(t, err) + unwrapped, err := dec.DecryptWithEphemeralKey(wrapped, compressed) + require.NoError(t, err, "KAS must be able to unwrap the EC-wrapped DEK") + assert.Equal(t, dek, unwrapped) +} + +// TestChunkedKeyAccessRejectsShareWithNoKAS pins the per-split +// invariant. Each split is one XOR share of the DEK, so a share that +// gets no key access object of its own is unrecoverable -- and the +// reader cannot tell: it builds its split set from the KAOs actually +// present in the manifest, so the missing share passes the +// completeness check, the DEK reconstructs wrong, and the failure +// surfaces as ErrRootSignatureFailure, which reads as tampering. +// Failing at creation is the only point where the cause is still +// visible. +func TestChunkedKeyAccessRejectsShareWithNoKAS(t *testing.T) { + pair, err := ocrypto.NewECKeyPair(ocrypto.ECCModeSecp256r1) + require.NoError(t, err) + pubPEM, err := pair.PublicKeyInPemFormat() + require.NoError(t, err) + + const kasURL = "https://kas.example.com" + dek := make([]byte, kKeySize) + splits := &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{ + kasURL: { + Algorithm: ocrypto.EC256Key, + KID: "ec-kid", + PEM: pubPEM, + URL: kasURL, + }, + }, + Splits: []Split{ + {ID: "wrapped", Data: dek, KASURLs: []string{kasURL}}, + {ID: "orphaned", Data: dek}, + }, + } + + _, err = buildChunkedKeyAccessObjects(splits, dek, []byte(`{"uuid":"test"}`), "") + require.Error(t, err, "a share with no KAS to unwrap it makes the DEK unrecoverable") + assert.Contains(t, err.Error(), "orphaned", "the error must name the split that cannot be recovered") +} + +// TestChunkedFinalizeManifestIsIndependent checks that the manifest +// Finalize hands back is not the one the writer keeps. GetManifest +// clones for this reason; returning the original from Finalize would +// let a caller's edit come back out of a later GetManifest, or race +// one. +func TestChunkedFinalizeManifestIsIndependent(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t) + + _, err := writer.WriteSegment(ctx, 0, []byte("payload")) + require.NoError(t, err) + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + require.NotEmpty(t, fin.Manifest.Segments) + + wantVersion := fin.Manifest.TDFVersion + wantHash := fin.Manifest.Segments[0].Hash + fin.Manifest.TDFVersion = "mutated" + fin.Manifest.Segments[0].Hash = "mutated" + + got, err := writer.GetManifest(ctx) + require.NoError(t, err) + assert.Equal(t, wantVersion, got.TDFVersion) + assert.Equal(t, wantHash, got.Segments[0].Hash) +} + +// TestChunkedLegacyTargetMode verifies that a pre-4.3.0 target mode +// produces the doubly-encoded (hex-then-base64) signatures that legacy +// readers require, and that the mainline reader -- which infers legacy +// mode solely from a missing schemaVersion -- still round-trips it. +func TestChunkedLegacyTargetMode(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t, WithChunkedTargetMode("4.2.2")) + + body := writeChunkedSegments(ctx, t, writer, [][]byte{ + []byte("legacy "), []byte("hex "), []byte("payload"), + }) + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + + // Absence of schemaVersion is the pre-4.3.0 marker readers key on. + assert.Empty(t, fin.Manifest.TDFVersion, "legacy manifest must omit schemaVersion") + + // A legacy signature is base64(hex(mac)), so it decodes to twice + // the MAC length: 64 for the 32-byte HS256 root, 32 for a 16-byte + // GMAC segment tag. The 4.3.0 form is base64(mac), half of each. + rootSig, err := ocrypto.Base64Decode([]byte(fin.Manifest.Signature)) + require.NoError(t, err) + assert.Len(t, rootSig, 64, "root signature must be hex-encoded before base64") + + for i, seg := range fin.Manifest.Segments { + segSig, err := ocrypto.Base64Decode([]byte(seg.Hash)) + require.NoError(t, err) + assert.Lenf(t, segSig, 2*kGMACPayloadLength, "segment %d hash must be hex-encoded before base64", i) + } + + tdfBytes := bytes.Join([][]byte{body, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("legacy hex payload"), plain) +} + +// TestChunkedCurrentTargetMode pins the 4.3.0-and-later form so a +// regression in either direction is caught. +func TestChunkedCurrentTargetMode(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t, WithChunkedTargetMode("4.3.0")) + + body := writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("current")}) + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + + assert.Equal(t, TDFSpecVersion, fin.Manifest.TDFVersion) + + rootSig, err := ocrypto.Base64Decode([]byte(fin.Manifest.Signature)) + require.NoError(t, err) + assert.Len(t, rootSig, 32, "root signature must be the raw HMAC, not hex") + + tdfBytes := bytes.Join([][]byte{body, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("current"), plain) +} + +// TestChunkedTargetModeInvalid rejects a non-semver target mode at +// construction rather than at Finalize. +func TestChunkedTargetModeInvalid(t *testing.T) { + ctx := context.Background() + kasBundle := newChunkedFakeKAS(t) + defer kasBundle.server.Close() + + _, err := NewChunkedWriter(ctx, + WithChunkedDefaultKAS(kasBundle.simpleKey()), + WithChunkedTargetMode("not-a-version"), + ) + require.Error(t, err) + assert.Contains(t, err.Error(), "not-a-version") +} + +// TestChunkedOptionsRejectNil checks that the injection-seam options +// refuse a nil value instead of storing it. A stored nil is +// indistinguishable from an unset field, so no default gets installed +// and the nil surfaces as a panic partway through writing -- for the +// key splitter, not until Finalize, after the caller has already +// encrypted and uploaded every segment. +func TestChunkedOptionsRejectNil(t *testing.T) { + ctx := context.Background() + kasBundle := newChunkedFakeKAS(t) + defer kasBundle.server.Close() + + for _, tc := range []struct { + name string + opt ChunkedWriterOption + }{ + {"archive writer factory", withChunkedArchiveWriterFactory(nil)}, + {"cipher factory", withChunkedCipherFactory(nil)}, + {"clock", withChunkedClock(nil)}, + {"key splitter", WithChunkedKeySplitter(nil)}, + {"rand", withChunkedRand(nil)}, + } { + t.Run(tc.name, func(t *testing.T) { + writer, err := NewChunkedWriter(ctx, + WithChunkedDefaultKAS(kasBundle.simpleKey()), + tc.opt, + ) + require.Error(t, err) + assert.Contains(t, err.Error(), "must not be nil") + assert.Nil(t, writer) + }) + } +} + +// TestChunkedIntegrityAlgorithmsAreFixed pins the only two algorithms +// the writer emits: an HS256 root over the aggregate hash, and GMAC +// segment hashes read out of the AEAD tag. +func TestChunkedIntegrityAlgorithmsAreFixed(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t) + + _, err := writer.WriteSegment(ctx, 0, []byte("first")) + require.NoError(t, err) + + // Segment 1, not 0: only segment 0's TDFData carries the payload's + // ZIP local file header, so 1 is exactly the bytes that were hashed. + second, err := writer.WriteSegment(ctx, 1, []byte("second")) + require.NoError(t, err) + sealed, err := io.ReadAll(second.TDFData) + require.NoError(t, err) + + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + + assert.Equal(t, hmacIntegrityAlgorithm, fin.Manifest.Algorithm) + assert.Equal(t, gmacIntegrityAlgorithm, fin.Manifest.SegmentHashAlgorithm) + + // The hash is the AEAD tag itself, not an HMAC over the same bytes. + segHash, err := ocrypto.Base64Decode([]byte(second.Hash)) + require.NoError(t, err) + require.Len(t, sealed, int(second.EncryptedSize)) + assert.Equal(t, sealed[len(sealed)-kGMACPayloadLength:], segHash) +} + +// errArchiveWriteFailed is the injected archive failure used to drive +// WriteSegment's error paths. +var ( + errArchiveWriteFailed = errors.New("archive write failed") + errArchiveCleanupFailed = errors.New("archive cleanup failed") +) + +// flakyArchiveWriter fails the first failures writes of one chosen +// segment index and delegates everything else to a real segment +// writer, so the archive itself stays consistent. +type flakyArchiveWriter struct { + zipstream.SegmentWriter + failIndex int + failures int +} + +func (f *flakyArchiveWriter) WriteSegment(ctx context.Context, index int, size uint64, crc32 uint32) ([]byte, error) { + if index == f.failIndex && f.failures > 0 { + f.failures-- + return nil, errArchiveWriteFailed + } + return f.SegmentWriter.WriteSegment(ctx, index, size, crc32) +} + +// flakyArchiveWriterFactory returns a withChunkedArchiveWriterFactory +// option whose archive fails the first failures writes to failIndex. +func flakyArchiveWriterFactory(failIndex, failures int) ChunkedWriterOption { + return withChunkedArchiveWriterFactory(func(clock clock) zipstream.SegmentWriter { + return &flakyArchiveWriter{ + SegmentWriter: defaultArchiveWriterFactory(clock), + failIndex: failIndex, + failures: failures, + } + }) +} + +// TestChunkedArchiveFailureKeepsManifestHonest checks that a segment +// whose bytes never reached the archive is not described in the +// manifest. A manifest covering bytes the caller never received leaves +// the reader mapping every later segment at the wrong payload offset, +// so the metadata is published only once the archive has accepted it. +// +// Skipping the index rather than retrying it is legal here — segment +// indices are ordering keys, not positions, so a sparse set finalizes +// normally (see segmentOrderLocked). Only index 0 is special, because +// it carries the ZIP local file header. +func TestChunkedArchiveFailureKeepsManifestHonest(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t, flakyArchiveWriterFactory(1, 1)) + + var body bytes.Buffer + write := func(index int, chunk string) error { + seg, err := writer.WriteSegment(ctx, index, []byte(chunk)) + if err != nil { + return err + } + _, err = io.Copy(&body, seg.TDFData) + return err + } + + require.NoError(t, write(0, "first ")) + + // Rejected by the archive, so the caller gets no bytes to append. + require.ErrorIs(t, write(1, "second "), errArchiveWriteFailed) + + require.NoError(t, write(2, "third")) + + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + body.Write(fin.Data) + + require.Len(t, fin.Manifest.Segments, 2, + "manifest must not describe the segment the archive rejected") + + reader, err := s.LoadTDF(bytes.NewReader(body.Bytes()), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, "first third", string(plain)) +} + +// TestChunkedSegmentRetryAfterArchiveFailure checks that a failed write +// releases its index, so a transient archive failure leaves the index +// writable rather than permanently blocking Finalize. +func TestChunkedSegmentRetryAfterArchiveFailure(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t, flakyArchiveWriterFactory(1, 1)) + + var body bytes.Buffer + write := func(index int, chunk string) error { + seg, err := writer.WriteSegment(ctx, index, []byte(chunk)) + if err != nil { + return err + } + _, err = io.Copy(&body, seg.TDFData) + require.NoError(t, err) + return nil + } + + require.NoError(t, write(0, "hello, ")) + + err := write(1, "chunked ") + require.ErrorIs(t, err, errArchiveWriteFailed) + + // The same index must be usable again. + require.NoError(t, write(1, "chunked "), "a failed segment must be retryable") + require.NoError(t, write(2, "world!")) + + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + + tdfBytes := bytes.Join([][]byte{body.Bytes(), fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("hello, chunked world!"), plain) +} + +// postWriteFailArchiveWriter delegates WriteSegment to a real writer +// -- so the archive's internal state is actually mutated -- and only +// then reports failure for one chosen index. Unlike +// flakyArchiveWriter, which fails before delegating, this exercises +// cleanup paths that only matter once the archive has partially +// accepted a write. cleanedUp records every CleanupSegment call so a +// test can assert release() actually rolls the archive back, not +// just the sdk-level reservation. +type postWriteFailArchiveWriter struct { + zipstream.SegmentWriter + failIndex int + failures int + cleanedUp []int + + // writer, when set, lets CleanupSegment observe the sdk-level + // reservation for the index it is rolling back. Each cleanup call + // appends what it saw to reservedDuringCleanup. + writer *chunkedWriter + reservedDuringCleanup []bool +} + +func (f *postWriteFailArchiveWriter) WriteSegment(ctx context.Context, index int, size uint64, crc32 uint32) ([]byte, error) { + header, err := f.SegmentWriter.WriteSegment(ctx, index, size, crc32) + if err != nil { + return header, err + } + if index == f.failIndex && f.failures > 0 { + f.failures-- + return nil, errArchiveWriteFailed + } + return header, nil +} + +func (f *postWriteFailArchiveWriter) CleanupSegment(index int) error { + f.cleanedUp = append(f.cleanedUp, index) + if f.writer != nil { + f.writer.mu.RLock() + _, held := f.writer.segments[index] + f.writer.mu.RUnlock() + f.reservedDuringCleanup = append(f.reservedDuringCleanup, held) + } + return f.SegmentWriter.CleanupSegment(index) +} + +// TestChunkedWriteSegmentCleansUpArchiveOnFailure checks that a +// segment the archive already accepted internally, but that +// WriteSegment then reports as failed, is rolled back via +// CleanupSegment rather than merely dropped from the sdk-level +// reservation map. Without this, a retry sees the archive's own +// bookkeeping for the index still present and fails a second time +// with an unrelated duplicate-segment error instead of succeeding. +// +// The rollback has to undo the archive's size and CRC accounting too, +// not just its record of the index, so the test carries on to a real +// TDF: the failed attempt's bytes were never handed to the caller, and +// if the archive still counted them every offset after segment 1 would +// be wrong. Only reading the payload back proves it does not. +func TestChunkedWriteSegmentCleansUpArchiveOnFailure(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + archive := &postWriteFailArchiveWriter{failIndex: 1, failures: 1} + writer, kasBundle := newChunkedWriterForTest(ctx, t, withChunkedArchiveWriterFactory(func(c clock) zipstream.SegmentWriter { + archive.SegmentWriter = defaultArchiveWriterFactory(c) + return archive + })) + + var body bytes.Buffer + write := func(index int, chunk string) error { + seg, err := writer.WriteSegment(ctx, index, []byte(chunk)) + if err != nil { + return err + } + _, err = io.Copy(&body, seg.TDFData) + require.NoError(t, err) + return nil + } + + require.NoError(t, write(0, "hello, ")) + + require.ErrorIs(t, write(1, "doomed"), errArchiveWriteFailed) + assert.Equal(t, []int{1}, archive.cleanedUp) + + // Deliberately a different length from the failed attempt: if the + // archive kept the doomed write's size, the mismatch shows up as a + // corrupt payload rather than being masked by identical accounting. + require.NoError(t, write(1, "chunked "), + "the archive's own record of the failed attempt must be rolled back, not just the sdk-level reservation") + require.NoError(t, write(2, "world!")) + + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + require.Len(t, fin.Manifest.Segments, 3) + body.Write(fin.Data) + + reader, err := s.LoadTDF(bytes.NewReader(body.Bytes()), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("hello, chunked world!"), plain) +} + +// TestChunkedCleanupRunsBeforeReleasingIndex pins the ordering of the +// two halves of the rollback. CleanupSegment addresses an index, not a +// particular attempt at it, so the reservation has to outlive it: if +// the index were released first, a racing write could claim it and get +// its bytes accepted by the archive, only for the losing attempt's +// cleanup to delete that record and roll back its size accounting. +// The winner would then publish segment metadata for bytes the archive +// no longer knows about. +// +// Observing the reservation from inside CleanupSegment pins the +// ordering directly, without needing a goroutine to lose the race +// often enough to be reliable. +func TestChunkedCleanupRunsBeforeReleasingIndex(t *testing.T) { + ctx := context.Background() + archive := &postWriteFailArchiveWriter{failIndex: 1, failures: 1} + writer, _ := newChunkedWriterForTest(ctx, t, withChunkedArchiveWriterFactory(func(c clock) zipstream.SegmentWriter { + archive.SegmentWriter = defaultArchiveWriterFactory(c) + return archive + })) + inner, ok := writer.(*chunkedWriter) + require.True(t, ok) + archive.writer = inner + + _, err := writer.WriteSegment(ctx, 1, []byte("doomed")) + require.ErrorIs(t, err, errArchiveWriteFailed) + assert.Equal(t, []bool{true}, archive.reservedDuringCleanup, + "the index must stay reserved until the archive rollback has finished") +} + +// TestChunkedConcurrentWrites exercises the contract WriteSegment +// documents but nothing tested: distinct indices may be written +// concurrently. Every other out-of-order test drives a single +// goroutine, so -race never saw the locking around w.mu, and neither +// the reservation nor the rollback path was observed under contention. +func TestChunkedConcurrentWrites(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t) + + const segments = 16 + chunks := make([][]byte, segments) + var want bytes.Buffer + for i := range chunks { + chunks[i] = []byte(fmt.Sprintf("segment-%02d;", i)) + want.Write(chunks[i]) + } + + // Index-keyed slices, so the goroutines share no mutable state of + // this test's own making and any race -race reports belongs to the + // writer. + segBytes := make([][]byte, segments) + errs := make([]error, segments) + + var wg sync.WaitGroup + start := make(chan struct{}) + for i := range segments { + wg.Add(1) + go func() { + defer wg.Done() + <-start // widen the window in which the writes overlap + seg, err := writer.WriteSegment(ctx, i, chunks[i]) + if err != nil { + errs[i] = err + return + } + segBytes[i], errs[i] = io.ReadAll(seg.TDFData) + }() + } + close(start) + wg.Wait() + + for i, err := range errs { + require.NoError(t, err, "segment %d", i) + } + + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + require.Equal(t, segments, fin.TotalSegments) + + // Concatenation is in index order regardless of write order. + var body bytes.Buffer + for _, buf := range segBytes { + body.Write(buf) + } + body.Write(fin.Data) + + reader, err := s.LoadTDF(bytes.NewReader(body.Bytes()), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, want.String(), string(plain)) +} + +// TestChunkedConcurrentDuplicateIndex checks the other half of the +// contract: when several goroutines race on one index, exactly one +// wins and the rest are rejected. The reservation is what makes this +// deterministic, so it is worth pinning under -race. +func TestChunkedConcurrentDuplicateIndex(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t) + + const racers = 8 + errs := make([]error, racers) + var wg sync.WaitGroup + start := make(chan struct{}) + for i := range racers { + wg.Add(1) + go func() { + defer wg.Done() + <-start + _, errs[i] = writer.WriteSegment(ctx, 0, []byte("contested")) + }() + } + close(start) + wg.Wait() + + var won int + for i, err := range errs { + if err == nil { + won++ + continue + } + require.ErrorIs(t, err, ErrChunkedSegmentAlreadyWritten, "racer %d", i) + } + assert.Equal(t, 1, won, "exactly one writer may claim an index") +} + +// errArchiveFinalizeFailed and errArchiveCloseFailed are the injected +// trailer-time failures used to drive Finalize's fencing paths. +var ( + errArchiveFinalizeFailed = errors.New("archive finalize failed") + errArchiveCloseFailed = errors.New("archive close failed") +) + +// trailerFailArchiveWriter delegates to a real segment writer and then +// reports failure from Finalize or Close. Delegating first is the +// point: it reproduces the state the sdk cannot recover from, where the +// archive has already mutated itself -- appended the payload entry to +// its central directory, written the manifest -- and only then failed. +type trailerFailArchiveWriter struct { + zipstream.SegmentWriter + failFinalize bool + failClose bool +} + +func (f *trailerFailArchiveWriter) Finalize(ctx context.Context, manifest []byte) ([]byte, error) { + out, err := f.SegmentWriter.Finalize(ctx, manifest) + if err != nil { + return out, err + } + if f.failFinalize { + return nil, errArchiveFinalizeFailed + } + return out, nil +} + +func (f *trailerFailArchiveWriter) Close() error { + if err := f.SegmentWriter.Close(); err != nil { + return err + } + if f.failClose { + return errArchiveCloseFailed + } + return nil +} + +// TestChunkedTrailerFailureFencesWriter pins that a failure from the +// archive's Finalize or from the Close after it leaves the writer +// permanently unusable, with every later call returning the same +// sentinel and the underlying cause still wrapped inside it. +// +// Neither is retryable. zipstream's Finalize appends the payload entry +// to its central directory partway through and can still fail +// afterwards without marking itself finalized, so a second attempt +// would append that entry twice; and an archive whose Finalize +// succeeded is terminally finalized regardless of what Close returned. +// The writer must therefore say so rather than look retryable -- the +// alternative is a caller looping on an operation that can only +// produce a corrupt archive or the same error forever. +func TestChunkedTrailerFailureFencesWriter(t *testing.T) { + for _, tc := range []struct { + name string + finalize bool + close bool + want error + cause error + }{ + {"finalize", true, false, ErrChunkedFinalizeFailed, errArchiveFinalizeFailed}, + {"close after finalize", false, true, ErrChunkedCloseFailed, errArchiveCloseFailed}, + } { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t, withChunkedArchiveWriterFactory(func(c clock) zipstream.SegmentWriter { + return &trailerFailArchiveWriter{ + SegmentWriter: defaultArchiveWriterFactory(c), + failFinalize: tc.finalize, + failClose: tc.close, + } + })) + + _, err := writer.WriteSegment(ctx, 0, []byte("payload")) + require.NoError(t, err) + + _, err = writer.Finalize(ctx) + require.ErrorIs(t, err, tc.want) + require.ErrorIs(t, err, tc.cause, "the sentinel must not swallow the archive's own error") + + // Every door is closed, and closed with the same error: a + // caller that kept a handle to the writer learns why rather + // than getting ErrChunkedAlreadyFinalized, which would be a + // lie -- nothing was finalized. + _, err = writer.Finalize(ctx) + require.ErrorIs(t, err, tc.want) + require.NotErrorIs(t, err, ErrChunkedAlreadyFinalized) + + _, err = writer.WriteSegment(ctx, 1, []byte("more")) + require.ErrorIs(t, err, tc.want) + + _, err = writer.GetManifest(ctx) + require.ErrorIs(t, err, tc.want) + }) + } +} + +// nonMutatingFinalizeArchiveWriter reports the other kind of trailer +// failure: one the archive caught before it changed anything, flagged as +// such. zipstream does raise these -- every refusal above its SetOrder +// call carries Mutated false -- but chunkedWriter's own preconditions +// pre-empt each one (closed writer, missing segment, missing segment 0, +// cancelled ctx), so reaching them through this package needs an injected +// writer. What that makes this test is a contract test, not a reachability +// one: given Mutated false, chunkedWriter must leave the writer retryable. +type nonMutatingFinalizeArchiveWriter struct { + zipstream.SegmentWriter + failures int +} + +func (f *nonMutatingFinalizeArchiveWriter) Finalize(ctx context.Context, manifest []byte) ([]byte, error) { + if f.failures > 0 { + f.failures-- + return nil, &zipstream.Error{ + Op: "finalize", + Type: "segment", + Err: errArchiveFinalizeFailed, + Mutated: false, + } + } + return f.SegmentWriter.Finalize(ctx, manifest) +} + +// TestChunkedNonMutatingTrailerFailureKeepsWriterUsable is the +// counterpart to TestChunkedTrailerFailureFencesWriter: an archive that +// says it failed before touching itself must not cost the caller the +// payload. Fencing here would discard an already-uploaded archive over +// a failure that leaves it byte-for-byte retryable, and the fence being +// too wide went unnoticed precisely because only the post-mutation case +// was exercised. +func TestChunkedNonMutatingTrailerFailureKeepsWriterUsable(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t, withChunkedArchiveWriterFactory(func(c clock) zipstream.SegmentWriter { + return &nonMutatingFinalizeArchiveWriter{ + SegmentWriter: defaultArchiveWriterFactory(c), + failures: 1, + } + })) + + seg0, err := writer.WriteSegment(ctx, 0, []byte("survives ")) + require.NoError(t, err) + body, err := io.ReadAll(seg0.TDFData) + require.NoError(t, err) + + _, err = writer.Finalize(ctx) + require.ErrorIs(t, err, errArchiveFinalizeFailed, "the archive's own error must still reach the caller") + require.NotErrorIs(t, err, ErrChunkedFinalizeFailed, "an unmutated archive must not fence the writer") + + // Usable in every direction the fence would have closed. + seg1, err := writer.WriteSegment(ctx, 1, []byte("the refusal")) + require.NoError(t, err) + tail, err := io.ReadAll(seg1.TDFData) + require.NoError(t, err) + + _, err = writer.GetManifest(ctx) + require.NoError(t, err) + + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + require.Len(t, fin.Manifest.Segments, 2) + + tdfBytes := bytes.Join([][]byte{body, tail, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("survives the refusal"), plain) +} + +// TestChunkedFinalizeContextCancellationIsRetryable pins the one +// non-mutating Finalize refusal reachable in production: a request-scoped +// deadline firing before the archive is touched. +// +// The refusal comes from chunkedWriter's own ctx.Err() check, not from +// zipstream's -- chunkedWriter checks immediately before handing off, which +// is what makes "nothing was mutated" structural rather than something +// inferred from an error the archive raised. What matters is the +// consequence: fencing here would destroy an already-encrypted, +// already-uploaded payload and blame the archive for a deadline. +func TestChunkedFinalizeContextCancellationIsRetryable(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t) + + seg0, err := writer.WriteSegment(ctx, 0, []byte("deadline ")) + require.NoError(t, err) + body, err := io.ReadAll(seg0.TDFData) + require.NoError(t, err) + + expired, cancel := context.WithDeadline(ctx, time.Now().Add(-time.Second)) + defer cancel() + + _, err = writer.Finalize(expired) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.NotErrorIs(t, err, ErrChunkedFinalizeFailed, "nothing was written, so nothing to fence") + + seg1, err := writer.WriteSegment(ctx, 1, []byte("survived")) + require.NoError(t, err) + tail, err := io.ReadAll(seg1.TDFData) + require.NoError(t, err) + + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + require.Len(t, fin.Manifest.Segments, 2) + + tdfBytes := bytes.Join([][]byte{body, tail, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("deadline survived"), plain) +} + +// TestChunkedWriteSegmentContextCancellationFreesIndex checks the same +// property one level down: a cancelled write must not wedge its index. +// A reservation left behind would make the index unwritable forever and +// now also blocks Finalize outright, so the recovery path matters more +// than it did. +func TestChunkedWriteSegmentContextCancellationFreesIndex(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t) + + expired, cancel := context.WithCancel(ctx) + cancel() + _, err := writer.WriteSegment(expired, 0, []byte("cancelled")) + require.ErrorIs(t, err, context.Canceled) + + _, err = writer.WriteSegment(ctx, 0, []byte("retried")) + require.NoError(t, err, "the cancelled write must have released its reservation") + + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + assert.Len(t, fin.Manifest.Segments, 1) +} + +// errCipherFailed is the injected cipher failure. +var errCipherFailed = errors.New("cipher failed") + +// failingCipher fails or panics on every call, depending on how it is +// built. Segment encryption is the one step before the archive is +// touched, so it exercises the reservation-rollback path with +// archiveWriteAttempted still false. +type failingCipher struct { + panics bool +} + +func (c failingCipher) EncryptInPlace(_ []byte) ([]byte, []byte, error) { + if c.panics { + panic("cipher exploded") + } + return nil, nil, errCipherFailed +} + +// TestChunkedCipherFailureReleasesIndex checks the reservation is +// released when encryption fails before the archive is involved, so the +// index stays writable. A wedged index is unrecoverable for the caller: +// every retry returns ErrChunkedSegmentAlreadyWritten for a segment +// that was never written. +func TestChunkedCipherFailureReleasesIndex(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t, withChunkedCipherFactory(func([]byte) (segmentCipher, error) { + return failingCipher{}, nil + })) + + _, err := writer.WriteSegment(ctx, 0, []byte("doomed")) + require.ErrorIs(t, err, errCipherFailed) + + inner, ok := writer.(*chunkedWriter) + require.True(t, ok) + inner.mu.RLock() + _, held := inner.segments[0] + inner.mu.RUnlock() + assert.False(t, held, "a failed encrypt must leave the index free to retry") +} + +// TestChunkedCipherPanicReleasesIndex covers the same rollback for a +// panic rather than an error. The release runs from a defer precisely +// so an injected cipher that panics -- or any future panic between the +// reservation and the commit -- cannot strand the index; a caller that +// recovers and retries has to find it free. +func TestChunkedCipherPanicReleasesIndex(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t, withChunkedCipherFactory(func([]byte) (segmentCipher, error) { + return failingCipher{panics: true}, nil + })) + + func() { + defer func() { + assert.NotNil(t, recover(), "the cipher was supposed to panic") + }() + _, _ = writer.WriteSegment(ctx, 0, []byte("doomed")) + }() + + inner, ok := writer.(*chunkedWriter) + require.True(t, ok) + inner.mu.RLock() + _, held := inner.segments[0] + inner.mu.RUnlock() + assert.False(t, held, "a panic between reservation and commit must still release the index") +} + +// xorSplitter splits the DEK into one share per KAS: n-1 random shares +// plus a final share chosen so the XOR of all of them is the DEK. This +// is the multi-KAS AND shape DefaultKeySplitter does not produce. With +// one split the reader's XOR loop runs over a single operand and the +// split-id grouping is never exercised. +type xorSplitter struct { + kases []*policy.SimpleKasKey +} + +func (s *xorSplitter) Split(_ context.Context, _ []*policy.Value, dek []byte, _ *policy.SimpleKasKey) (*SplitResult, error) { + out := &SplitResult{KASPublicKeys: make(map[string]KASPublicKey, len(s.kases))} + accum := make([]byte, len(dek)) + for i, kas := range s.kases { + share := make([]byte, len(dek)) + if i == len(s.kases)-1 { + // Last share closes the XOR back onto the DEK. + for j := range share { + share[j] = dek[j] ^ accum[j] + } + } else { + if _, err := io.ReadFull(rand.Reader, share); err != nil { + return nil, err + } + for j := range share { + accum[j] ^= share[j] + } + } + url := kas.GetKasUri() + alg, err := PolicyAlgorithmToKeyType(kas.GetPublicKey().GetAlgorithm()) + if err != nil { + return nil, err + } + out.KASPublicKeys[url] = KASPublicKey{ + Algorithm: alg, + KID: kas.GetPublicKey().GetKid(), + PEM: kas.GetPublicKey().GetPem(), + URL: url, + } + out.Splits = append(out.Splits, Split{ + Data: share, + ID: fmt.Sprintf("split-%d", i), + KASURLs: []string{url}, + }) + } + return out, nil +} + +// TestChunkedMultiSplitRoundTrip drives two splits against two +// independent KASes end to end, so the reader has to collect more than +// one share, group the key access objects by sid, and XOR the results. +// A writer-side change that emitted the right number of KAOs with the +// wrong sids fails nowhere until decryption. +func TestChunkedMultiSplitRoundTrip(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + + kasA := newChunkedFakeKAS(t) + t.Cleanup(kasA.server.Close) + kasB := newChunkedFakeKAS(t) + t.Cleanup(kasB.server.Close) + + writer, err := NewChunkedWriter(ctx, + WithChunkedKeySplitter(&xorSplitter{kases: []*policy.SimpleKasKey{kasA.simpleKey(), kasB.simpleKey()}}), + ) + require.NoError(t, err) + + body := writeChunkedSegments(ctx, t, writer, [][]byte{ + []byte("two "), []byte("kas "), []byte("split"), + }) + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + + require.Len(t, fin.Manifest.KeyAccessObjs, 2) + sids := make(map[string]string, 2) + for _, kao := range fin.Manifest.KeyAccessObjs { + sids[kao.SplitID] = kao.KasURL + } + assert.Equal(t, map[string]string{ + "split-0": kasA.url, + "split-1": kasB.url, + }, sids, "each share must be wrapped to its own KAS under its own sid") + + tdfBytes := bytes.Join([][]byte{body, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasA.url, kasB.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("two kas split"), plain) +} + +// TestChunkedAttributes covers the three attribute/KAS options: the +// writer-level default, the Finalize-level override, and the +// Finalize-level KAS. They reach the manifest by two different +// routes -- attributes become the policy body, the KAS decides which +// endpoint each key access object names -- so a regression in the +// override precedence is invisible without checking both. +func TestChunkedAttributes(t *testing.T) { + ctx := context.Background() + + initial := []*policy.Value{{Fqn: "https://example.com/attr/initial/value/one"}} + override := []*policy.Value{ + {Fqn: "https://example.com/attr/override/value/a"}, + {Fqn: "https://example.com/attr/override/value/b"}, + } + + // policyFQNs decodes the base64 policy the manifest carries and + // returns the attribute FQNs inside it. + policyFQNs := func(t *testing.T, m *Manifest) []string { + t.Helper() + raw, err := ocrypto.Base64Decode([]byte(m.Policy)) + require.NoError(t, err) + var p PolicyObject + require.NoError(t, json.Unmarshal(raw, &p)) + out := make([]string, 0, len(p.Body.DataAttributes)) + for _, a := range p.Body.DataAttributes { + out = append(out, a.Attribute) + } + return out + } + + t.Run("initial attributes reach the policy", func(t *testing.T) { + writer, _ := newChunkedWriterForTest(ctx, t, WithChunkedInitialAttributes(initial)) + writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("payload")}) + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + assert.Equal(t, []string{initial[0].GetFqn()}, policyFQNs(t, fin.Manifest)) + }) + + t.Run("finalize attributes replace them", func(t *testing.T) { + writer, _ := newChunkedWriterForTest(ctx, t, WithChunkedInitialAttributes(initial)) + writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("payload")}) + fin, err := writer.Finalize(ctx, WithChunkedAttributes(override)) + require.NoError(t, err) + // Replace, not merge: the initial set must be gone entirely. + assert.Equal(t, []string{override[0].GetFqn(), override[1].GetFqn()}, policyFQNs(t, fin.Manifest)) + }) + + // An empty override reads as "not specified", so the initial set + // survives. Documented on WithChunkedAttributes, and pinned here + // because the alternative -- letting an empty slice win -- would + // silently loosen the policy on the data, and nothing downstream + // could tell that from a writer that never had attributes. + for _, tc := range []struct { + name string + values []*policy.Value + }{ + {"nil finalize attributes keep the initial set", nil}, + {"empty finalize attributes keep the initial set", []*policy.Value{}}, + } { + t.Run(tc.name, func(t *testing.T) { + writer, _ := newChunkedWriterForTest(ctx, t, WithChunkedInitialAttributes(initial)) + writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("payload")}) + fin, err := writer.Finalize(ctx, WithChunkedAttributes(tc.values)) + require.NoError(t, err) + assert.Equal(t, []string{initial[0].GetFqn()}, policyFQNs(t, fin.Manifest)) + }) + } + + t.Run("finalize default KAS overrides the writer's", func(t *testing.T) { + writer, initialKAS := newChunkedWriterForTest(ctx, t) + lateKAS := newChunkedFakeKAS(t) + t.Cleanup(lateKAS.server.Close) + + writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("payload")}) + fin, err := writer.Finalize(ctx, WithChunkedDefaultKASForFinalize(lateKAS.simpleKey())) + require.NoError(t, err) + + require.Len(t, fin.Manifest.KeyAccessObjs, 1) + assert.Equal(t, lateKAS.url, fin.Manifest.KeyAccessObjs[0].KasURL) + assert.NotEqual(t, initialKAS.url, fin.Manifest.KeyAccessObjs[0].KasURL) + }) +} + +// blockingCipher lets a test hold one WriteSegment inside encryption +// while it inspects writer state from another goroutine. started fires +// once the call is inside; release unblocks it. +type blockingCipher struct { + inner segmentCipher + target int + calls int + mu sync.Mutex + started chan struct{} + release chan struct{} +} + +func (c *blockingCipher) EncryptInPlace(data []byte) ([]byte, []byte, error) { + c.mu.Lock() + n := c.calls + c.calls++ + c.mu.Unlock() + if n == c.target { + close(c.started) + <-c.release + } + return c.inner.EncryptInPlace(data) +} + +// TestChunkedGetManifestDuringInFlightWrite pins that a reserved but +// not-yet-written index is invisible to the manifest. The reservation +// exists from the moment WriteSegment claims the index, well before the +// archive has accepted anything, so a manifest built in that window +// would otherwise either describe a segment with no hash and no size or +// fail outright -- and GetManifest running alongside in-flight writes is +// the whole point of the snapshot. +func TestChunkedGetManifestDuringInFlightWrite(t *testing.T) { + ctx := context.Background() + cipher := &blockingCipher{ + // Segments 0 and 2 go through; the third call -- index 1 -- is + // held. Leaving a written index above the blocked one is what + // lets the keepSegments check below reach its per-index branch + // rather than stopping at the count. + target: 2, + started: 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 + })) + + _, err := writer.WriteSegment(ctx, 0, []byte("first")) + require.NoError(t, err) + _, err = writer.WriteSegment(ctx, 2, []byte("third")) + require.NoError(t, err) + + done := make(chan error, 1) + go func() { + _, err := writer.WriteSegment(ctx, 1, []byte("second")) + done <- err + }() + <-cipher.started + + snap, err := writer.GetManifest(ctx) + require.NoError(t, err, "a reservation in flight must not fail the snapshot") + assert.Len(t, snap.Segments, 2, "the snapshot must describe only segments the archive has accepted") + + // The same index must also be refused as a keepSegments member while + // it is merely reserved, for the same reason: naming it would put a + // segment in the manifest that has no bytes behind it. + _, err = writer.GetManifest(ctx, WithChunkedSegments([]int{0, 1})) + require.Error(t, err) + assert.Contains(t, err.Error(), "references segment 1 which was not written") + + close(cipher.release) + require.NoError(t, <-done) + + snap, err = writer.GetManifest(ctx) + require.NoError(t, err) + assert.Len(t, snap.Segments, 3, "the segment appears once the archive has it") +} + +// TestChunkedFinalizeRejectsInFlightWrite pins the answer to the race +// GetManifest is allowed to ignore. A manifest is a snapshot and may be +// one segment short; an archive may not. zipstream has already counted +// the in-flight segment's bytes into the payload entry by the time it +// is visible here, so a trailer built now records offsets that overshoot +// unless the caller appends bytes the manifest never mentions -- while +// the racing WriteSegment returns success. Refusing costs a retry; +// proceeding costs the payload. +func TestChunkedFinalizeRejectsInFlightWrite(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + cipher := &blockingCipher{ + // Call 0 is segment 0; call 1 -- segment 1 -- is held open. + target: 1, + started: make(chan struct{}), + release: make(chan struct{}), + } + writer, kasBundle := 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 + })) + + seg0, err := writer.WriteSegment(ctx, 0, []byte("held ")) + require.NoError(t, err) + body, err := io.ReadAll(seg0.TDFData) + require.NoError(t, err) + + done := make(chan *ChunkedSegmentResult, 1) + go func() { + seg, err := writer.WriteSegment(ctx, 1, []byte("open")) + assert.NoError(t, err) + done <- seg + }() + <-cipher.started + + _, err = writer.Finalize(ctx) + require.ErrorIs(t, err, ErrChunkedWriteInFlight) + assert.Contains(t, err.Error(), "[1]", "the error must name the index the caller has to wait on") + + // GetManifest must not inherit the rejection. It produces nothing + // durable, so a snapshot one segment short is a correct snapshot of + // this instant -- and running alongside in-flight writes is its + // entire purpose. + snap, err := writer.GetManifest(ctx) + require.NoError(t, err) + assert.Len(t, snap.Segments, 1) + + close(cipher.release) + seg1 := <-done + require.NotNil(t, seg1) + tail, err := io.ReadAll(seg1.TDFData) + require.NoError(t, err) + + // Not fenced: the refusal happened before anything was touched, so + // joining the goroutine and retrying is all the recovery needed. + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + require.Len(t, fin.Manifest.Segments, 2) + + tdfBytes := bytes.Join([][]byte{body, tail, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("held open"), plain) +} + +// TestChunkedRealisticSegmentSizes runs a 2 MiB segment size rather +// than the handful of bytes every other test uses. Segment +// sizing is where the ZIP64 thresholds, the CRC combine over full-size +// segments, and the offset arithmetic actually get exercised; a +// three-byte payload passes all of them trivially. +func TestChunkedRealisticSegmentSizes(t *testing.T) { + if testing.Short() { + t.Skip("allocates ~6 MiB and encrypts it") + } + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t) + + const segSize = 2 * 1024 * 1024 + chunks := make([][]byte, 3) + var want bytes.Buffer + for i := range chunks { + chunk := make([]byte, segSize) + // Varied, reproducible content: a constant fill would hide a + // segment written at the wrong offset. + for j := range chunk { + chunk[j] = byte(i*7 + j) + } + chunks[i] = chunk + want.Write(chunk) + } + + body := writeChunkedSegments(ctx, t, writer, chunks) + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + assert.Equal(t, int64(3*segSize), fin.TotalSize) + assert.Equal(t, int64(segSize), fin.Manifest.DefaultSegmentSize) + + tdfBytes := bytes.Join([][]byte{body, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, want.Bytes(), plain) +} + +// partialSplitter names two KAS URLs on one split but resolves a public +// key for only the first, the shape a splitter produces when a KAS +// lookup fails and the failure is swallowed upstream. +type partialSplitter struct { + known *policy.SimpleKasKey + // missingURL is listed on the split but absent from KASPublicKeys. + missingURL string +} + +func (s partialSplitter) Split(_ context.Context, _ []*policy.Value, dek []byte, _ *policy.SimpleKasKey) (*SplitResult, error) { + url := s.known.GetKasUri() + share := make([]byte, len(dek)) + copy(share, dek) + return &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{ + url: { + Algorithm: ocrypto.RSA2048Key, + KID: s.known.GetPublicKey().GetKid(), + PEM: s.known.GetPublicKey().GetPem(), + URL: url, + }, + }, + Splits: []Split{{ + Data: share, + KASURLs: []string{url, s.missingURL}, + }}, + }, nil +} + +// TestChunkedFinalizeRejectsUnresolvedKAS checks that a split naming a +// KAS with no resolved public key fails Finalize. Skipping it would +// emit a TDF whose KAO set silently omits that KAS -- and if every URL +// on a split were missing, the share would be unrecoverable. +func TestChunkedFinalizeRejectsUnresolvedKAS(t *testing.T) { + ctx := context.Background() + kasBundle := newChunkedFakeKAS(t) + defer kasBundle.server.Close() + + w, err := NewChunkedWriter(ctx, + WithChunkedDefaultKAS(kasBundle.simpleKey()), + WithChunkedKeySplitter(partialSplitter{ + known: kasBundle.simpleKey(), + missingURL: "https://unresolved.example.com", + }), + ) + require.NoError(t, err) + + _, err = w.WriteSegment(ctx, 0, []byte("payload")) + require.NoError(t, err) + + _, err = w.Finalize(ctx) + require.ErrorIs(t, err, errKasPubKeyMissing) + assert.Contains(t, err.Error(), "https://unresolved.example.com") +} + +// brokenSplitter returns a share that passes every structural check +// and still does not rebuild the DEK -- the failure mode SplitResult's +// own doc describes, where the reader collects every share it is told +// about, XORs them, and reports a root signature failure that reads as +// tampering. +type brokenSplitter struct { + kas *policy.SimpleKasKey + // short truncates the share instead of corrupting it. A single + // short share agrees with itself on length, so Validate accepts it + // and the DEK's tail is left unmasked. + short bool +} + +func (s brokenSplitter) Split(_ context.Context, _ []*policy.Value, dek []byte, _ *policy.SimpleKasKey) (*SplitResult, error) { + share := slices.Clone(dek) + if s.short { + share = share[:len(share)/2] + } else { + share[0] ^= 0xFF + } + url := s.kas.GetKasUri() + return &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{ + url: { + Algorithm: ocrypto.RSA2048Key, + KID: s.kas.GetPublicKey().GetKid(), + PEM: s.kas.GetPublicKey().GetPem(), + URL: url, + }, + }, + Splits: []Split{{Data: share, KASURLs: []string{url}}}, + }, nil +} + +// TestChunkedFinalizeRejectsNonReconstructingSplitter checks that a +// splitter whose shares do not rebuild the DEK is caught at Finalize +// rather than at the reader's root signature check, hours later and on +// someone else's machine. +func TestChunkedFinalizeRejectsNonReconstructingSplitter(t *testing.T) { + for _, tc := range []struct { + name string + short bool + wantSub string + }{ + {"share of the wrong value", false, "do not XOR back to the DEK"}, + {"share shorter than the DEK", true, "the length of the DEK"}, + } { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + kasBundle := newChunkedFakeKAS(t) + t.Cleanup(kasBundle.server.Close) + + w, err := NewChunkedWriter(ctx, + WithChunkedDefaultKAS(kasBundle.simpleKey()), + WithChunkedKeySplitter(brokenSplitter{kas: kasBundle.simpleKey(), short: tc.short}), + ) + require.NoError(t, err) + + _, err = w.WriteSegment(ctx, 0, []byte("payload")) + require.NoError(t, err) + + _, err = w.Finalize(ctx) + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantSub) + + // The splits are built before the archive is touched, so a + // splitter fault is a correctable mistake rather than a lost + // payload. Moving the key access build below the archive call + // would make it terminal, silently. + require.NotErrorIs(t, err, ErrChunkedFinalizeFailed) + _, err = w.WriteSegment(ctx, 1, []byte("more")) + require.NoError(t, err, "a splitter fault must not fence the writer") + }) + } +} + +// dekZeroingSplitter wipes the slice it is handed once it has taken a +// copy, the way a splitter treating the DEK as scratch space would. Its +// output is correct; only the caller's buffer is damaged. +type dekZeroingSplitter struct { + kas *policy.SimpleKasKey +} + +func (s dekZeroingSplitter) Split(_ context.Context, _ []*policy.Value, dek []byte, _ *policy.SimpleKasKey) (*SplitResult, error) { + share := slices.Clone(dek) + clear(dek) + url := s.kas.GetKasUri() + return &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{ + url: { + Algorithm: ocrypto.RSA2048Key, + KID: s.kas.GetPublicKey().GetKid(), + PEM: s.kas.GetPublicKey().GetPem(), + URL: url, + }, + }, + Splits: []Split{{Data: share, KASURLs: []string{url}}}, + }, nil +} + +// TestChunkedSplitterCannotDamageTheDEK checks that the writer hands +// the splitter a copy. The splitter runs between the segment signatures +// (already computed against the DEK) and the root signature (computed +// after it), so a splitter that mutates its argument desynchronizes the +// two and produces a TDF that fails verification exactly as a tampered +// one does -- with the wrapped key still correct, so the KAS grants +// access and the failure surfaces only at the last step. +func TestChunkedSplitterCannotDamageTheDEK(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + kasBundle := newChunkedFakeKAS(t) + t.Cleanup(kasBundle.server.Close) + + writer, err := NewChunkedWriter(ctx, + WithChunkedDefaultKAS(kasBundle.simpleKey()), + WithChunkedKeySplitter(dekZeroingSplitter{kas: kasBundle.simpleKey()}), + ) + require.NoError(t, err) + + body := writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("intact "), []byte("key")}) + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + + tdfBytes := bytes.Join([][]byte{body, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, []byte("intact key"), plain) +} + +// TestChunkedInitialAttributesAreCloned checks that reusing the slice +// passed at construction cannot change the policy the writer later +// writes. The writer holds it for its whole lifetime, and a policy +// quietly loosened after the fact is the one mistake here that nothing +// downstream can detect. +func TestChunkedInitialAttributesAreCloned(t *testing.T) { + ctx := context.Background() + const want = "https://example.com/attr/clone/value/original" + + attrs := []*policy.Value{{Fqn: want}} + writer, _ := newChunkedWriterForTest(ctx, t, WithChunkedInitialAttributes(attrs)) + attrs[0] = &policy.Value{Fqn: "https://example.com/attr/clone/value/swapped"} + + writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("payload")}) + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + + raw, err := ocrypto.Base64Decode([]byte(fin.Manifest.Policy)) + require.NoError(t, err) + var p PolicyObject + require.NoError(t, json.Unmarshal(raw, &p)) + require.Len(t, p.Body.DataAttributes, 1) + assert.Equal(t, want, p.Body.DataAttributes[0].Attribute) +} + +// TestChunkedEncryptedSizeIsNotTheByteCount pins the trap in +// ChunkedFinalizeResult.EncryptedSize: it describes the manifest, not the +// file. A caller who sizes an upload from it truncates the payload, and the +// corruption surfaces only at decrypt. +// +// Two distinct shortfalls stack here, and the test measures them separately so +// a regression names which one moved. Segment 0's TDFData carries the +// payload's ZIP local file header, which EncryptedSize does not count; and +// WithChunkedSegments drops segment 2 from the manifest while the archive +// still counts its bytes and the caller must still append them. +func TestChunkedEncryptedSizeIsNotTheByteCount(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t) + + chunks := [][]byte{[]byte("keep-0-"), []byte("keep-1-"), []byte("drop-2!")} + var appended, keptManifestBytes, droppedBytes, headerBytes int64 + for i, chunk := range chunks { + seg, err := writer.WriteSegment(ctx, i, chunk) + require.NoError(t, err) + n, err := io.Copy(io.Discard, seg.TDFData) + require.NoError(t, err) + appended += n + if i == 0 { + // Whatever TDFData carries beyond the segment's own ciphertext is + // the local file header. Derived, not hardcoded: the header's + // length depends on the archive writer's naming, and pinning a + // literal here would only assert the default factory. + headerBytes = n - seg.EncryptedSize + } + if i < 2 { + keptManifestBytes += seg.EncryptedSize + } else { + droppedBytes += n + } + } + require.Positive(t, headerBytes, "segment 0's TDFData must carry the payload header") + require.Positive(t, droppedBytes) + + fin, err := writer.Finalize(ctx, WithChunkedSegments([]int{0, 1})) + require.NoError(t, err) + require.Len(t, fin.Manifest.Segments, 2) + + assert.Equal(t, keptManifestBytes, fin.EncryptedSize, + "EncryptedSize sums exactly the segments the manifest describes") + assert.Equal(t, fin.EncryptedSize+headerBytes+droppedBytes, appended, + "the bytes to append exceed EncryptedSize by the payload header plus every dropped segment") + assert.Less(t, fin.EncryptedSize, appended, + "sizing an upload from EncryptedSize truncates it") +} + +// TestChunkedMimeType covers the default and the override together: +// the default is what a caller who never sets one ships, so it is as +// much a part of the contract as the option. +func TestChunkedMimeType(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t) + writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("payload")}) + + snap, err := writer.GetManifest(ctx) + require.NoError(t, err) + assert.Equal(t, "application/octet-stream", snap.MimeType) + + fin, err := writer.Finalize(ctx, WithChunkedMimeType("application/pdf")) + require.NoError(t, err) + assert.Equal(t, "application/pdf", fin.Manifest.MimeType) +} + +// TestChunkedZeroLengthSegments exercises the invariant segmentSlot's +// two-field design exists for: a written segment is distinguished from +// an absent one by a flag, not by its size, so an empty segment must be +// a segment and not a gap. An interior empty segment also checks that +// the offsets of everything after it are unaffected. +func TestChunkedZeroLengthSegments(t *testing.T) { + for _, tc := range []struct { + name string + chunks [][]byte + want string + }{ + {"empty interior segment", [][]byte{[]byte("a"), {}, []byte("b")}, "ab"}, + {"empty trailing segment", [][]byte{[]byte("ab"), {}}, "ab"}, + {"empty payload", [][]byte{{}}, ""}, + } { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + writer, kasBundle := newChunkedWriterForTest(ctx, t) + + body := writeChunkedSegments(ctx, t, writer, tc.chunks) + fin, err := writer.Finalize(ctx) + require.NoError(t, err) + require.Len(t, fin.Manifest.Segments, len(tc.chunks), + "an empty segment is still a segment") + + tdfBytes := bytes.Join([][]byte{body, fin.Data}, nil) + reader, err := s.LoadTDF(bytes.NewReader(tdfBytes), + WithKasAllowlist([]string{kasBundle.url}), + ) + require.NoError(t, err) + + plain, err := io.ReadAll(reader) + require.NoError(t, err) + assert.Equal(t, tc.want, string(plain)) + }) + } +} + +// TestChunkedManifestCloneIsDeep pins that the slices hanging off a +// returned manifest are copies too. Segments is already covered by the +// round trips; KeyAccessObjs and Assertions are not, so dropping either +// clone would leave a caller's edit -- or a data race against a +// concurrent GetManifest -- in the writer's own record of what the +// archive bytes say. +func TestChunkedManifestCloneIsDeep(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t) + writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("payload")}) + + fin, err := writer.Finalize(ctx, WithChunkedAssertions([]AssertionConfig{{ + ID: "a", + Type: BaseAssertion, + Scope: PayloadScope, + AppliesToState: Unencrypted, + Statement: Statement{Format: "json", Schema: "urn:test", Value: `{"k":"v"}`}, + }})) + require.NoError(t, err) + require.Len(t, fin.Manifest.KeyAccessObjs, 1) + require.Len(t, fin.Manifest.Assertions, 1) + require.Len(t, fin.Manifest.Segments, 1) + + wantKAS := fin.Manifest.KeyAccessObjs[0].KasURL + wantAssertion := fin.Manifest.Assertions[0].ID + wantHash := fin.Manifest.Segments[0].Hash + + fin.Manifest.KeyAccessObjs[0].KasURL = "https://tampered.example.com" + fin.Manifest.Assertions[0].ID = "tampered" + fin.Manifest.Segments[0].Hash = "tampered" + + again, err := writer.GetManifest(ctx) + require.NoError(t, err) + assert.Equal(t, wantKAS, again.KeyAccessObjs[0].KasURL) + assert.Equal(t, wantAssertion, again.Assertions[0].ID) + assert.Equal(t, wantHash, again.Segments[0].Hash) +} + +// TestSummarizeSegmentIndices covers the truncation branch, which only +// fires on a writer with more segments in flight than any test would +// realistically hold open. Unbounded, the ErrChunkedWriteInFlight +// message on a ten-thousand-segment writer is measured in kilobytes. +func TestSummarizeSegmentIndices(t *testing.T) { + for _, tc := range []struct { + name string + n int + want string + }{ + {"empty", 0, "[]"}, + {"at the cap", maxReportedSegmentIndices, "[0 1 2 3 4 5 6 7]"}, + {"past the cap", maxReportedSegmentIndices + 3, "[0 1 2 3 4 5 6 7] ... (3 more)"}, + } { + t.Run(tc.name, func(t *testing.T) { + indices := make([]int, tc.n) + for i := range indices { + indices[i] = i + } + assert.Equal(t, tc.want, summarizeSegmentIndices(indices)) + }) + } +} + +// TestChunkedEmptyAttributeFQNRejected checks that Finalize refuses an +// attribute value whose FQN is empty rather than writing it into the policy. +// +// GetFqn is nil-receiver safe, so both a nil element and a value that never +// had its FQN populated reach buildChunkedPolicy as "". Emitting that produces +// a policy naming no attribute: it fails closed, but only at decrypt, after +// the caller has encrypted and uploaded every byte. Nothing between here and +// there looks at it. +func TestChunkedEmptyAttributeFQNRejected(t *testing.T) { + for _, tc := range []struct { + name string + attrs []*policy.Value + }{ + {"unpopulated FQN", []*policy.Value{{Fqn: ""}}}, + {"nil value", []*policy.Value{nil}}, + { + "one good, one empty", + []*policy.Value{{Fqn: "https://example.com/attr/a/value/b"}, {Fqn: ""}}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t, WithChunkedInitialAttributes(tc.attrs)) + writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("payload")}) + + _, err := writer.Finalize(ctx) + require.ErrorContains(t, err, "empty FQN") + + // Not fenced: nothing was written, so a caller who can supply the + // real attributes may still finalize this payload. + fin, err := writer.Finalize(ctx, + WithChunkedAttributes([]*policy.Value{{Fqn: "https://example.com/attr/a/value/b"}})) + require.NoError(t, err) + require.NotNil(t, fin) + }) + } +} + +// cleanupFailArchiveWriter fails the rollback of a failed write, leaving the +// archive in a state this package cannot characterize. +type cleanupFailArchiveWriter struct { + zipstream.SegmentWriter + failIndex int +} + +func (f *cleanupFailArchiveWriter) WriteSegment(ctx context.Context, index int, size uint64, crc32 uint32) ([]byte, error) { + if index == f.failIndex { + return nil, errArchiveWriteFailed + } + return f.SegmentWriter.WriteSegment(ctx, index, size, crc32) +} + +func (f *cleanupFailArchiveWriter) CleanupSegment(index int) error { + if index == f.failIndex { + return errArchiveCleanupFailed + } + return f.SegmentWriter.CleanupSegment(index) +} + +// TestChunkedCleanupFailureFencesWriter checks that a rollback the archive +// could not perform stops the writer instead of being swallowed. +// +// CleanupSegment is what makes a failed WriteSegment retryable. If it fails, +// the archive may still carry the attempt's contribution to the payload's size +// and CRC, and nothing here can tell how much -- so every later segment would +// be placed at an offset computed from a length the manifest never describes. +// That assembles into a file readers accept and misread, which is exactly the +// failure mode worth trading a dead writer for. +func TestChunkedCleanupFailureFencesWriter(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t, + withChunkedArchiveWriterFactory(func(clock clock) zipstream.SegmentWriter { + return &cleanupFailArchiveWriter{ + SegmentWriter: defaultArchiveWriterFactory(clock), + failIndex: 1, + } + })) + + seg0, err := writer.WriteSegment(ctx, 0, []byte("first ")) + require.NoError(t, err) + _, err = io.Copy(io.Discard, seg0.TDFData) + require.NoError(t, err) + + // The caller sees the write error, not the rollback error: the write is + // the cause, and reporting the rollback instead would hide it. + _, err = writer.WriteSegment(ctx, 1, []byte("second ")) + require.ErrorIs(t, err, errArchiveWriteFailed) + require.NotErrorIs(t, err, errArchiveCleanupFailed) + + // Every later call reports the fence, wrapping the rollback's cause. + for _, call := range []struct { + name string + run func() error + }{ + {"WriteSegment", func() error { _, err := writer.WriteSegment(ctx, 2, []byte("third")); return err }}, + {"Finalize", func() error { _, err := writer.Finalize(ctx); return err }}, + {"GetManifest", func() error { _, err := writer.GetManifest(ctx); return err }}, + {"retry the failed index", func() error { + _, err := writer.WriteSegment(ctx, 1, []byte("second ")) + return err + }}, + } { + t.Run(call.name, func(t *testing.T) { + err := call.run() + require.ErrorIs(t, err, ErrChunkedCleanupFailed) + assert.ErrorIs(t, err, errArchiveCleanupFailed, "the fence must carry its cause") + }) + } +} + +// TestChunkedEmptyMimeTypeFallsBackToDefault checks that +// WithChunkedMimeType("") is treated as "not specified" rather than taken +// literally. The empty string is what a caller forwarding an unset field +// passes, and a manifest with no MIME type at all is not what they meant. +func TestChunkedEmptyMimeTypeFallsBackToDefault(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t) + writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("payload")}) + + fin, err := writer.Finalize(ctx, WithChunkedMimeType("")) + require.NoError(t, err) + assert.Equal(t, defaultMimeType, fin.Manifest.MimeType) +} diff --git a/sdk/chunked_writer.go b/sdk/chunked_writer.go new file mode 100644 index 0000000000..2225c868c9 --- /dev/null +++ b/sdk/chunked_writer.go @@ -0,0 +1,1193 @@ +package sdk + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/json" + "errors" + "fmt" + "hash/crc32" + "io" + "slices" + "sync" + "time" + + "github.com/google/uuid" + "github.com/opentdf/platform/lib/ocrypto" + "github.com/opentdf/platform/protocol/go/policy" + "github.com/opentdf/platform/sdk/internal/zipstream" +) + +// The injection seams below — the clock, the segment cipher, the archive +// writer and the entropy source — are unexported on purpose. Each exists so +// in-package tests can pin non-deterministic behavior; none is reachable from +// outside. For archiveWriterFactory it is the signature that closes the seam, +// not the return type's implementability: a struct declared anywhere with the +// right method set satisfies zipstream.SegmentWriter, but no code outside this +// package can spell `func(clock) zipstream.SegmentWriter` — clock is +// unexported and zipstream lives under internal/ — so a factory value cannot +// be constructed to pass in. + +// clock supplies the current time to the chunked writer and, through it, to +// the zipstream layer that stamps ZIP header timestamps. Injected so tests can +// pin timestamps and produce byte-for-byte deterministic TDF output. +type clock interface { + // Now returns the current wall-clock time. + Now() time.Time +} + +// systemClock returns time.Now(). Production default. +type systemClock struct{} + +// Now returns the current wall-clock time. +func (systemClock) Now() time.Time { return time.Now() } + +// fixedClock returns the same time on every call, for deterministic ZIP +// output in tests. +type fixedClock struct { + // T is the wall-clock time to return from Now. + T time.Time +} + +// Now returns the pinned time. +func (c fixedClock) Now() time.Time { return c.T } + +// segmentCipher encrypts a single payload segment. Implementations must be +// safe for concurrent use by segment writers. +// +// The output must be AEAD in the shape the TDF reader expects: a fresh nonce +// per call, and a ciphertext ending in a 16-byte authentication tag. +// WriteSegment concatenates nonce+ciphertext and hands the result to +// segmentIntegrity, which under SegmentGMAC reads the tag straight off the +// tail; a cipher that omits the tag or returns a repeated nonce produces a +// manifest that verifies against nothing. +type segmentCipher interface { + // EncryptInPlace returns (ciphertext, nonce, error). Despite the name -- + // inherited from ocrypto.AesGcm -- nothing is encrypted in place: the + // implementation must allocate its output and must neither retain nor + // modify data, which WriteSegment passes through from its caller. + EncryptInPlace(data []byte) ([]byte, []byte, error) +} + +// segmentCipherFactory builds a segmentCipher from the writer-generated DEK. +// Tests inject deterministic ciphers for reproducible fixtures. +type segmentCipherFactory func(dek []byte) (segmentCipher, error) + +// defaultSegmentCipherFactory wraps ocrypto.NewAESGcm (AES-256-GCM). +func defaultSegmentCipherFactory(dek []byte) (segmentCipher, error) { + return ocrypto.NewAESGcm(dek) +} + +// archiveWriterFactory builds a zipstream.SegmentWriter for a new TDF. It +// receives the writer's clock so ZIP header timestamps stay injectable +// end-to-end. +type archiveWriterFactory func(c clock) zipstream.SegmentWriter + +// defaultArchiveWriterFactory returns a ZIP64-enabled segment writer with its +// clock plumbed to the caller-supplied clock. +// +// The 1 is expectedSegments, which reads like a capacity hint but is not: +// nothing raises it, so a writer that goes on to accept five hundred segments +// still reports ExpectedCount 1. It is harmless only because that count is +// consulted solely when no explicit segment order has been set, and +// zipstream's Finalize always derives an order from the segments actually +// present before it checks completeness. +func defaultArchiveWriterFactory(c clock) zipstream.SegmentWriter { + return zipstream.NewSegmentTDFWriter(1, + zipstream.WithZip64(), + zipstream.WithClock(c.Now), + ) +} + +// Sentinel errors returned by [ChunkedWriter]. +// +// Experimental: not part of the stable SDK API; may change or be removed. +var ( + // ErrChunkedAlreadyFinalized is returned when a ChunkedWriter + // method is called after Finalize has already succeeded. + ErrChunkedAlreadyFinalized = errors.New("chunked: writer already finalized") + + // ErrChunkedCleanupFailed is returned when a failed WriteSegment + // could not roll its partial write back out of the archive, and by + // every later call on that writer. + // + // CleanupSegment is what makes a failed WriteSegment retryable: it + // returns the archive to a state indistinguishable from never + // having attempted the index. If it fails, the archive may still + // carry that attempt's record and its contribution to the payload's + // size and CRC, and nothing here can tell how much. Continuing + // would write later segments at offsets computed from a payload + // length that counts bytes the manifest never describes, producing + // an archive that assembles cleanly and fails at decrypt. Fencing + // the writer surfaces it at the next call instead. + // + // Unreachable with the default archive writer, whose CleanupSegment + // cannot fail; it is the injected-factory case this guards. + ErrChunkedCleanupFailed = errors.New("chunked: segment cleanup failed after a failed write; writer is unusable") + + // ErrChunkedCloseFailed is returned when the archive's Close fails + // after its Finalize already succeeded. The archive is terminally + // finalized internally at that point regardless, so the writer is + // unusable: every subsequent call returns this same error rather + // than retrying against an archive that can only fail again. + // + // Note what is being discarded. Finalize had already returned a + // complete, valid trailer; it is dropped because the archive + // reported a fault afterwards, not because the bytes were bad. The + // payload segments already uploaded are still sound -- the TDF is + // unfinishable through this writer, but nothing about the data + // itself is in question. + // + // Unreachable with the default archive writer, whose Close cannot + // fail; it is the injected-factory case this guards. + ErrChunkedCloseFailed = errors.New("chunked: archive close failed after finalize; writer is unusable") + + // ErrChunkedFinalizeFailed is returned when the archive reports + // that it had already mutated itself before failing -- or reports + // nothing either way -- and by every later call on that writer. + // Such a failure is not retryable: zipstream appends the payload + // entry to the central directory partway through finalizing, then + // can still fail on the manifest entry, on a ZIP64 requirement, or + // on generating the directory bytes, and it does not roll that back + // or mark itself finalized. A second attempt would append the + // payload entry again, yielding a central directory with two + // entries for one payload: an archive some readers accept and + // silently misread. Failing fast is the only safe answer. + // + // An archive failure raised before any mutation does not produce + // this error and does not fence the writer; see Finalize. + ErrChunkedFinalizeFailed = errors.New("chunked: archive finalize failed; writer is unusable") + + // ErrChunkedInvalidSegmentIndex is returned when WriteSegment + // receives a negative index. + ErrChunkedInvalidSegmentIndex = errors.New("chunked: invalid segment index") + + // ErrChunkedMissingSegmentZero is returned when Finalize is called + // on a writer that never wrote segment 0. Only segment 0 emits the + // payload's ZIP local file header, and every offset in the manifest + // and central directory is measured from it. + ErrChunkedMissingSegmentZero = errors.New("chunked: segment 0 was never written; it carries the payload's ZIP local file header") + + // ErrChunkedSegmentAlreadyWritten is returned when WriteSegment + // receives an index that was already written. + ErrChunkedSegmentAlreadyWritten = errors.New("chunked: segment already written") + + // ErrChunkedWriteInFlight is returned by Finalize when a + // WriteSegment has reserved an index but the archive has not yet + // accepted its bytes. Finalizing inside that window would omit the + // segment from the manifest while the racing WriteSegment still + // returned success, so the caller would have no signal at all. + // + // Nothing was written and the writer is not fenced: join the + // outstanding goroutines and call Finalize again. + ErrChunkedWriteInFlight = errors.New("chunked: a segment write is still in flight; every WriteSegment must return before Finalize") +) + +// ChunkedWriter creates a TDF from segments that may arrive in any +// order. Callers write each segment independently — typically +// off-thread or in parallel — then call Finalize to close the +// archive. Contrast with SDK.CreateTDF, which requires the full +// plaintext up front. +// +// Experimental: not part of the stable SDK API; may change or be removed. +type ChunkedWriter interface { + // Finalize completes TDF creation. Every option applies only to + // this Finalize call; writer-level defaults set at NewChunked* + // remain otherwise. Returns the closing bytes (the payload's data + // descriptor, the embedded manifest entry, and the central + // directory + end-of-central-directory record) that must be + // appended after every segment's TDFData. Returns + // ErrChunkedMissingSegmentZero if segment 0 was never written, + // ErrChunkedAlreadyFinalized on a second call, and + // ErrChunkedWriteInFlight if a WriteSegment has not yet returned. + // + // Only a mutation of the archive fences the writer. A refusal made + // before the archive is touched -- an in-flight write, a cancelled + // ctx, a rejected option, a splitter fault -- leaves the writer + // usable: fix the cause and call Finalize again. A failure from the + // archive itself after it has begun mutating, or a Close failure + // after Finalize succeeded, leaves the writer unusable + // (ErrChunkedFinalizeFailed / ErrChunkedCloseFailed), and every + // subsequent call on it returns that same error rather than + // retrying. + // + // A permitted retry is safe, not idempotent: each call mints a + // fresh policy UUID and re-splits the DEK, so the manifest it + // produces differs from the one the failed call would have. + // + // Finalize must happen-after every WriteSegment call returns. + // Sequencing that is the caller's job -- wait on the errgroup, + // close the channel, join the pool. Finalize does not block waiting + // for in-flight writes (it cannot know how many are still coming); + // it refuses with ErrChunkedWriteInFlight instead, so a missed join + // is a deterministic error rather than a silently short TDF. + Finalize(ctx context.Context, opts ...ChunkedFinalizeOption) (*ChunkedFinalizeResult, error) + + // GetManifest returns the manifest for the TDF. After Finalize it + // is exactly the manifest that was written. + // + // Before Finalize it is a preview of the manifest's *shape* -- + // segment count, sizes, hashes, which KASes appear -- and nothing + // more. It is not byte-stable and is not what Finalize will + // produce: each call re-runs the whole manifest build, minting a + // fresh policy UUID, re-splitting the DEK (so ephemeral wrapping + // keys, and therefore every key access object, differ), and + // re-signing any assertions. Two back-to-back calls with identical + // options disagree, and so will Finalize. Use it to inspect + // progress, never to predict or cache the final manifest. + // + // It reflects only segments the archive has already accepted; + // writes still in flight are skipped rather than waited on or + // reported. Unlike Finalize it produces nothing durable, so a + // snapshot one segment short is a correct snapshot of that instant + // -- which is why, unlike Finalize, it does not refuse them. + // + // It is safe to call concurrently with WriteSegment, but it is not + // free: GetManifest and Finalize both hold the writer lock across + // KeySplitter.Split, so a splitter that does I/O stalls concurrent + // WriteSegment calls at their commit step for its duration. + GetManifest(ctx context.Context, opts ...ChunkedFinalizeOption) (*Manifest, error) + + // WriteSegment encrypts data as segment index and returns the ZIP + // bytes for that segment: segment 0 is preceded by the payload's + // ZIP local file header, every other segment is nonce + ciphertext + // only. Callers upload or buffer those bytes; Finalize does not + // re-emit them. Indices need not arrive in order and need not be + // contiguous, but index 0 is mandatory: it carries that local file + // header, so Finalize refuses a write set without it. + // + // data is neither retained nor modified; the caller may reuse the + // buffer as soon as the call returns. + WriteSegment(ctx context.Context, index int, data []byte) (*ChunkedSegmentResult, error) +} + +// ChunkedSegmentResult carries the ZIP bytes for one segment plus its +// integrity metadata. +// +// Experimental: not part of the stable SDK API; may change or be removed. +type ChunkedSegmentResult struct { + // EncryptedSize is the ciphertext byte length including nonce and + // GCM tag. It is the size the manifest records for this segment, + // which for segment 0 is *not* the number of bytes TDFData yields: + // that reader is prefixed with the payload's ZIP local file header, + // which the manifest does not count. Size an upload from TDFData + // itself (or from len(header)+EncryptedSize), never from + // EncryptedSize alone. + EncryptedSize int64 + + // Hash is the base64-encoded segment integrity hash. + Hash string + + // Index is the zero-based segment index. + Index int + + // PlaintextSize is the byte length of the pre-encryption input. + PlaintextSize int64 + + // TDFData is a reader over the segment's ZIP-embedded ciphertext: + // for segment 0 this is the payload's local file header + nonce + + // AES-GCM output; every other segment omits the local header and + // is nonce + AES-GCM output only. Callers assemble the TDF by + // concatenating each segment's TDFData in emission order followed + // by ChunkedFinalizeResult.Data. + // + // It may be read only once -- it is a stream over the buffers this + // call produced, not a rewindable view of them. A second read + // yields zero bytes and no error, so an upload that retries by + // re-reading it silently writes nothing. Buffer the bytes yourself + // if you need them twice. + TDFData io.Reader +} + +// ChunkedFinalizeResult carries the finalized TDF's closing bytes and +// metadata about what was written. +// +// Experimental: not part of the stable SDK API; may change or be removed. +type ChunkedFinalizeResult struct { + // Data is the ZIP closing bytes, in order: the payload's data + // descriptor, the embedded manifest entry (its own local file + // header + JSON data), and the central directory + EOCD record. + // Append after every segment's TDFData to form the complete TDF + // file -- including any segment WithChunkedSegments excluded from + // the manifest; the archive's recorded size and CRC already + // account for it regardless. + Data []byte + + // EncryptedSize is the total ciphertext byte length across the + // segments the manifest describes -- post-trim if + // WithChunkedSegments was used. + // + // This is a property of the manifest, not of the file. It is not + // the number of bytes to append: it counts nonce+ciphertext only, + // where segment 0's TDFData additionally carries the payload's ZIP + // local file header, and it omits any segment WithChunkedSegments + // dropped, whose bytes the caller must still append. Size an upload + // by measuring the TDFData readers as they are consumed -- the only + // quantity that counts every byte exactly once. Sizing it from this + // field truncates the payload, and the corruption surfaces only at + // decrypt. + EncryptedSize int64 + + // Manifest is the finalized manifest that was serialized into the + // archive. + Manifest *Manifest + + // TotalSegments is the number of segments in the finalized + // manifest (post-trim if WithChunkedSegments was used). + TotalSegments int + + // TotalSize is the total plaintext byte length across the segments + // the manifest describes, post-trim on the same basis as + // EncryptedSize. + TotalSize int64 +} + +// chunkedWriterConfig captures the settings supplied at +// NewChunkedWriter time. Fields are unexported; use options. +type chunkedWriterConfig struct { + // archiveFactory builds the ZIP archive writer that lays out the + // TDF. Defaults to defaultArchiveWriterFactory. + archiveFactory archiveWriterFactory + + // cipherFactory builds the segment cipher from the DEK. Defaults + // to defaultSegmentCipherFactory (AES-256-GCM). + cipherFactory segmentCipherFactory + + // clock supplies the current time to the writer and the + // underlying zipstream. Defaults to systemClock. Tests inject + // fixedClock for deterministic ZIP output. + clock clock + + // excludeVersion omits the schemaVersion field from the manifest. + // Set together with useHex by WithChunkedTargetMode; readers use + // the field's absence as the pre-4.3.0 marker, so the two must + // agree. + excludeVersion bool + + // initialAttributes are the attribute values used at Finalize + // when the Finalize call does not supply its own. + initialAttributes []*policy.Value + + // initialDefaultKAS is the default KAS used at Finalize when the + // Finalize call does not supply its own. + initialDefaultKAS *policy.SimpleKasKey + + // rand is the entropy source used to generate the DEK. Defaults + // to crypto/rand.Reader. + rand io.Reader + + // splitter maps attribute values to DEK splits at Finalize time. + // Defaults to DefaultKeySplitter (single-KAS only). + splitter KeySplitter + + // useHex hex-encodes segment, root, and assertion signatures + // before base64, producing the doubly-encoded form that readers + // older than 4.3.0 require. Set by WithChunkedTargetMode. + useHex bool +} + +// chunkedFinalizeConfig captures Finalize-time overrides. +type chunkedFinalizeConfig struct { + // assertions to sign and attach to the produced TDF. Each + // AssertionConfig may carry a SigningKey; one that does not is + // signed with HS256 over the writer's DEK. + assertions []AssertionConfig + + // attributes overrides the writer's initialAttributes for this + // Finalize call. + attributes []*policy.Value + + // defaultKAS overrides the writer's initialDefaultKAS for this + // Finalize call. + defaultKAS *policy.SimpleKasKey + + // encryptedMetadata is opaque metadata AES-GCM-encrypted on each + // KAO with the split share. + encryptedMetadata string + + // excludeVersion omits the schemaVersion field from the manifest + // for compatibility with older readers. Defaults to the writer's + // setting; see WithChunkedTargetMode. + excludeVersion bool + + // keepSegments names the segments the finalized manifest + // describes: must be a prefix of the written segments in ascending + // index order, dropping only from the end. Empty means every + // written segment, ascending. See WithChunkedSegments. + keepSegments []int + + // mimeType records the payload MIME type in the manifest. + // Defaults to [defaultMimeType]. + mimeType string +} + +// ChunkedWriterOption configures a ChunkedWriter at construction +// time. +// +// Experimental: not part of the stable SDK API; may change or be removed. +type ChunkedWriterOption func(*chunkedWriterConfig) error + +// ChunkedFinalizeOption configures a single Finalize call. +// +// Experimental: not part of the stable SDK API; may change or be removed. +type ChunkedFinalizeOption func(*chunkedFinalizeConfig) error + +// segmentSlot is one entry in the writer's segment table. The slot +// appears the moment WriteSegment reserves an index, which is what +// makes a second write to that index fail, but written stays false +// until the archive has durably accepted the bytes. Only a written +// slot is visible to the manifest: a reservation describes bytes that +// may yet never exist. +// +// The two are separate fields rather than a sentinel in seg.Size +// because zero-length segments are legal, so no size value is free to +// mean "not yet". +type segmentSlot struct { + // seg is the manifest metadata for this segment. Meaningful only + // once written is true. + seg Segment + + // written is true once the archive has accepted the segment's + // bytes and seg has been filled in. + written bool +} + +// chunkedWriter is the concrete ChunkedWriter. +type chunkedWriter struct { + // archiveWriter handles the underlying ZIP archive creation. + archiveWriter zipstream.SegmentWriter + + // block is the segment cipher built from the DEK. + block segmentCipher + + // dek is the Data Encryption Key. 32 bytes (AES-256). + dek []byte + + // excludeVersion omits schemaVersion from the manifest unless a + // Finalize option overrides it. + excludeVersion bool + + // unusable is non-nil once the archive has been left in a state no + // later call can recover from: its Finalize failed partway through + // its own mutations, its Close failed after Finalize had succeeded, + // or a failed WriteSegment could not roll its partial write back + // out. In each case w.finalized stays false, yet the archive can no + // longer be trusted, so every method returns this error instead of + // continuing. It holds the wrapped original, so callers see both + // the sentinel and the underlying cause. + // + // Only an unrecoverable mutation sets it. A Finalize that fails + // before the archive is touched -- an in-flight write, a cancelled + // ctx, a bad option, a splitter fault, or an archive error reported + // as non-mutating -- leaves this nil and the writer retryable, as + // does a failed WriteSegment whose cleanup succeeded. + unusable error + + // finalized is true once Finalize returns successfully. + finalized bool + + // initialAttributes captured at construction; used by Finalize + // when the caller does not override. + initialAttributes []*policy.Value + + // initialDefaultKAS captured at construction; used by Finalize + // when the caller does not override. + initialDefaultKAS *policy.SimpleKasKey + + // manifest holds the finalized manifest for post-Finalize + // GetManifest calls. + manifest *Manifest + + // mu guards writer state that spans WriteSegment and Finalize. + mu sync.RWMutex + + // segments records per-index slots, reserved and then written. + segments map[int]*segmentSlot + + // splitter converts attributes + DEK into key splits at + // Finalize time. + splitter KeySplitter + + // useHex selects the pre-4.3.0 doubly-encoded signature form. + // Read by WriteSegment, so it is fixed at construction rather + // than at Finalize. + useHex bool +} + +// NewChunkedWriter constructs a per-segment TDF writer. WriteSegment +// may be called from several goroutines at once so long as each +// targets a distinct segment index; two concurrent calls for the same +// index are not allowed, and one of them will fail with +// ErrChunkedSegmentAlreadyWritten rather than corrupt the archive. +// +// Finalize must happen-after every WriteSegment returns. The writer +// cannot sequence that for you, but it will not paper over a missed +// join either: Finalize refuses with ErrChunkedWriteInFlight while any +// index is reserved but not yet accepted by the archive, rather than +// emitting a TDF that is short by however many writes lost the race. +// It does not block — it has no way to know how many more writes are +// coming — so join your goroutines and call it again. +// +// No SDK value is needed. The key splitter (WithChunkedKeySplitter) +// and the attribute and KAS defaults (WithChunkedInitialAttributes, +// WithChunkedDefaultKAS) are supplied through options; the clock, +// cipher, archive-writer and entropy seams are unexported test seams +// and are not reachable from outside this package. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func NewChunkedWriter(_ context.Context, opts ...ChunkedWriterOption) (ChunkedWriter, error) { + cfg := chunkedWriterConfig{ + archiveFactory: defaultArchiveWriterFactory, + cipherFactory: defaultSegmentCipherFactory, + clock: systemClock{}, + rand: rand.Reader, + splitter: DefaultKeySplitter(), + } + for _, opt := range opts { + if err := opt(&cfg); err != nil { + return nil, err + } + } + + dek := make([]byte, kKeySize) + if _, err := io.ReadFull(cfg.rand, dek); err != nil { + return nil, fmt.Errorf("generate DEK: %w", err) + } + block, err := cfg.cipherFactory(dek) + if err != nil { + return nil, fmt.Errorf("build segment cipher: %w", err) + } + return &chunkedWriter{ + archiveWriter: cfg.archiveFactory(cfg.clock), + block: block, + dek: dek, + excludeVersion: cfg.excludeVersion, + initialAttributes: cfg.initialAttributes, + initialDefaultKAS: cfg.initialDefaultKAS, + segments: make(map[int]*segmentSlot), + splitter: cfg.splitter, + useHex: cfg.useHex, + }, nil +} + +// Finalize serializes the manifest, closes the archive, and returns +// the trailing bytes. +func (w *chunkedWriter) Finalize(ctx context.Context, opts ...ChunkedFinalizeOption) (*ChunkedFinalizeResult, error) { + w.mu.Lock() + defer w.mu.Unlock() + if w.unusable != nil { + return nil, w.unusable + } + if w.finalized { + return nil, ErrChunkedAlreadyFinalized + } + + // Refuse before diagnosing anything else. A write still in flight makes + // every later check unsound: the segment-0 check below would report a + // missing segment 0 when segment 0 is merely unfinished, sending the caller + // hunting an indexing bug. Nothing has been touched at this point, so the + // writer stays usable -- join the goroutines and call again. + if reserved := w.reservedIndicesLocked(); len(reserved) > 0 { + return nil, fmt.Errorf("%w: %d in flight, indices %s", + ErrChunkedWriteInFlight, len(reserved), summarizeSegmentIndices(reserved)) + } + + // Segment 0 is the only one that emits the payload's ZIP local file + // header (see zipstream.segmentWriter.WriteSegment), and the archive + // measures every offset it records from that header being at the + // front of the assembled stream. It cannot be synthesized here: by + // Finalize the caller has already encrypted and uploaded the bytes it + // would have to precede. + if slot, ok := w.segments[0]; !ok || !slot.written { + return nil, ErrChunkedMissingSegmentZero + } + + cfg, err := w.applyFinalizeOptions(opts) + if err != nil { + return nil, err + } + + manifest, totals, err := w.buildManifest(ctx, cfg) + if err != nil { + return nil, err + } + manifestBytes, err := json.Marshal(manifest) + if err != nil { + return nil, fmt.Errorf("marshal manifest: %w", err) + } + + // ctx is checked here rather than inferred from the archive's error. A + // cancellation the archive notices after it has started mutating is + // indistinguishable at the error site from one it notices before, and only + // the second is safe to retry. Checking immediately beforehand makes the + // distinction structural: nothing has been touched, so nothing to fence. + if err := ctx.Err(); err != nil { + return nil, fmt.Errorf("chunked: finalize: %w", err) + } + + finalBytes, err := w.archiveWriter.Finalize(ctx, manifestBytes) + if err != nil { + // Fence unless the archive affirmatively reports that it failed before + // changing itself. Fail safe: "we mutated" and "we did not say" get the + // same answer, because the cost of guessing wrong is a central + // directory with two entries for one payload -- an archive some readers + // accept and silently misread. See ErrChunkedFinalizeFailed. + var zerr *zipstream.Error + if errors.As(err, &zerr) && !zerr.Mutated { + // Not fenced: w.unusable and w.finalized are untouched, so the + // caller may write more segments and finalize again. + return nil, fmt.Errorf("chunked: finalize archive: %w", err) + } + w.unusable = fmt.Errorf("%w: %w", ErrChunkedFinalizeFailed, err) + return nil, w.unusable + } + if err := w.archiveWriter.Close(); err != nil { + // The archive is terminally finalized internally regardless of + // this error, so a retry can only ever hit the same failure. + w.unusable = fmt.Errorf("%w: %w", ErrChunkedCloseFailed, err) + return nil, w.unusable + } + + w.finalized = true + w.manifest = manifest + return &ChunkedFinalizeResult{ + Data: finalBytes, + EncryptedSize: totals.encrypted, + // Cloned for the same reason GetManifest clones: w.manifest is + // the writer's own record of what the archive bytes say, and + // handing out that pointer would let a caller edit it and see + // the edit come back from a later GetManifest, or race one. + Manifest: cloneChunkedManifest(manifest), + TotalSegments: len(manifest.Segments), + TotalSize: totals.plaintext, + }, nil +} + +// GetManifest returns the manifest snapshot. +func (w *chunkedWriter) GetManifest(ctx context.Context, opts ...ChunkedFinalizeOption) (*Manifest, error) { + w.mu.RLock() + defer w.mu.RUnlock() + if w.unusable != nil { + return nil, w.unusable + } + if w.finalized { + if w.manifest == nil { + // Unreachable unless Finalize is changed to set finalized without + // recording what it wrote. Refuse rather than fall through to the + // rebuild below: after finalize the caller is asking what shipped, + // and a rebuild does not answer that. It mints a fresh policy UUID + // and fresh key splits, so it would hand back a plausible manifest + // that does not describe the bytes -- and a caller that stored it + // alongside them could not decrypt. + return nil, errors.New("chunked: writer is finalized but recorded no manifest") + } + return cloneChunkedManifest(w.manifest), nil + } + cfg, err := w.applyFinalizeOptions(opts) + if err != nil { + return nil, err + } + // No in-flight check here, deliberately: see GetManifest's interface doc. + manifest, _, err := w.buildManifest(ctx, cfg) + if err != nil { + return nil, err + } + return manifest, nil +} + +// WriteSegment encrypts data as segment index and returns the ZIP +// bytes for that segment. +func (w *chunkedWriter) WriteSegment(ctx context.Context, index int, data []byte) (*ChunkedSegmentResult, error) { + w.mu.Lock() + if w.unusable != nil { + err := w.unusable + w.mu.Unlock() + return nil, err + } + if w.finalized { + w.mu.Unlock() + return nil, ErrChunkedAlreadyFinalized + } + if index < 0 { + w.mu.Unlock() + return nil, ErrChunkedInvalidSegmentIndex + } + if _, ok := w.segments[index]; ok { + w.mu.Unlock() + return nil, ErrChunkedSegmentAlreadyWritten + } + // Reserve the index so a concurrent write to the same one is + // rejected, but leave the slot unwritten: the segment does not + // count as written until its bytes are in the archive. + slot := &segmentSlot{} + w.segments[index] = slot + w.mu.Unlock() + + // release drops the reservation so the caller can retry this index + // after a failure. It matches on identity and on the placeholder + // still being unwritten, so it can never discard a segment some + // other call has since completed. + // cleanupErr, when non-nil, additionally fences the writer: the + // reservation is gone but the archive may still carry the failed + // attempt, so no later call can be trusted. + release := func(cleanupErr error) { + w.mu.Lock() + if cur, ok := w.segments[index]; ok && cur == slot && !cur.written { + delete(w.segments, index) + } + if cleanupErr != nil && w.unusable == nil { + w.unusable = fmt.Errorf("%w: segment %d: %w", ErrChunkedCleanupFailed, index, cleanupErr) + } + w.mu.Unlock() + } + + // committed marks the point past which the archive has durably + // accepted the write. Until then, every exit path -- including a + // panic unwinding through an injected cipher or archive-writer seam + // -- must release the reservation, or the index is wedged forever: + // retries see ErrChunkedSegmentAlreadyWritten and default-mode + // Finalize can never find every index accounted for. + committed := false + archiveWriteAttempted := false + defer func() { + if committed { + return + } + var cleanupErr error + if archiveWriteAttempted { + // The archive may have partially recorded the write before + // failing; CleanupSegment undoes that so a retry starts + // from a state indistinguishable from never having been + // attempted (see zipstream.SegmentWriter's contract). The + // concrete writer's CleanupSegment cannot itself fail; a + // custom archiveWriterFactory that does fail here leaves + // the archive in a state this code cannot characterize, so + // release() fences the writer. It does not change what this + // call returns -- the original write error is the cause, + // and reporting the rollback instead would hide it. + // + // This must precede release(): the reservation is the only + // thing keeping another goroutine out of this index, and + // CleanupSegment addresses the index rather than a + // particular attempt at it. Releasing first lets a racing + // write claim the index and get its bytes accepted by the + // archive, whereupon this call deletes that attempt's + // record and rolls back its size accounting -- leaving the + // winner to publish segment metadata for bytes the archive + // no longer knows about. + // + // The two locks are taken sequentially here -- CleanupSegment + // takes and drops the archive's, then release() takes and drops + // w.mu -- never nested. Finalize nests them the other way, w.mu + // across archiveWriter.Finalize, so hoisting w.mu.Lock() above + // this call to "simplify" the defer closes the cycle and + // deadlocks the writer. + cleanupErr = w.archiveWriter.CleanupSegment(index) + } + release(cleanupErr) + }() + + ciphertext, nonce, err := w.block.EncryptInPlace(data) + if err != nil { + return nil, fmt.Errorf("encrypt segment %d: %w", index, err) + } + sig, err := segmentIntegrity(ciphertext, w.dek, SegmentGMAC, w.useHex) + if err != nil { + return nil, fmt.Errorf("segment %d signature: %w", index, err) + } + hash := string(ocrypto.Base64Encode([]byte(sig))) + encryptedSize := int64(len(nonce) + len(ciphertext)) + + crc := crc32.Update(crc32.ChecksumIEEE(nonce), crc32.IEEETable, ciphertext) + // Deliberately outside w.mu. Finalize's ErrChunkedWriteInFlight exists + // because this call is unsynchronized against it; holding w.mu here would + // close that window by serializing every segment write, which is the one + // thing this writer exists not to do. + archiveWriteAttempted = true + header, err := w.archiveWriter.WriteSegment(ctx, index, uint64(encryptedSize), crc) + if err != nil { + return nil, fmt.Errorf("write segment %d to archive: %w", index, err) + } + + // Commit only once the archive has accepted the segment. Publishing + // the metadata earlier would let Finalize emit a manifest that + // describes bytes the archive never received. + w.mu.Lock() + slot.seg = Segment{ + EncryptedSize: encryptedSize, + Hash: hash, + Size: int64(len(data)), + } + slot.written = true + // Set under the same lock that publishes the slot, not after it. Once + // written is true, release() declines to drop the reservation, so a panic + // unwinding between the two would run CleanupSegment -- deleting the + // archive's record and rolling back its size accounting -- while leaving a + // slot the manifest still describes. Publication and commit are one step. + committed = true + w.mu.Unlock() + + var reader io.Reader + if len(header) == 0 { + reader = io.MultiReader(bytes.NewReader(nonce), bytes.NewReader(ciphertext)) + } else { + reader = io.MultiReader(bytes.NewReader(header), bytes.NewReader(nonce), bytes.NewReader(ciphertext)) + } + // Reported from the locals rather than from seg, which is shared with + // concurrent readers of w.segments once the lock is released. + return &ChunkedSegmentResult{ + EncryptedSize: encryptedSize, + Hash: hash, + Index: index, + PlaintextSize: int64(len(data)), + TDFData: reader, + }, nil +} + +// reservedIndicesLocked returns the indices WriteSegment has claimed but the +// archive has not accepted, in ascending order. Caller holds mu. +// +// This is the in-flight set, computed rather than counted. A separate counter +// would have to be incremented with the reservation and decremented at both +// exits -- commit and release -- and would drift silently if either exit were +// ever missed; the drift direction that matters (undercount) reopens exactly +// the race this guards. The map is the reservation, so reading it cannot +// disagree with it. +// +// Cost is one pass over w.segments, on a path that already sorts it +// (segmentOrderLocked) and walks it again (buildManifest). +func (w *chunkedWriter) reservedIndicesLocked() []int { + var reserved []int + for idx, slot := range w.segments { + if !slot.written { + reserved = append(reserved, idx) + } + } + slices.Sort(reserved) + return reserved +} + +// maxReportedSegmentIndices caps how many indices an error message lists, so a +// writer with thousands of segments in flight does not produce an error string +// measured in kilobytes. +const maxReportedSegmentIndices = 8 + +// summarizeSegmentIndices renders indices for an error message, truncating +// past maxReportedSegmentIndices. +func summarizeSegmentIndices(indices []int) string { + if len(indices) <= maxReportedSegmentIndices { + return fmt.Sprint(indices) + } + return fmt.Sprintf("%v ... (%d more)", + indices[:maxReportedSegmentIndices], len(indices)-maxReportedSegmentIndices) +} + +// applyFinalizeOptions builds a chunkedFinalizeConfig with defaults +// then applies each option in order. +func (w *chunkedWriter) applyFinalizeOptions(opts []ChunkedFinalizeOption) (*chunkedFinalizeConfig, error) { + cfg := &chunkedFinalizeConfig{ + attributes: nil, + encryptedMetadata: "", + excludeVersion: w.excludeVersion, + } + for _, opt := range opts { + if err := opt(cfg); err != nil { + return nil, err + } + } + // Every default below is resolved after the option loop, so the zero value + // uniformly reads as "not specified" and falls back. Defaulting mimeType + // before the loop instead would make WithChunkedMimeType("") the one option + // whose empty argument is taken literally, writing a manifest with no MIME + // type rather than the default one. + if len(cfg.attributes) == 0 && len(w.initialAttributes) > 0 { + cfg.attributes = w.initialAttributes + } + if cfg.defaultKAS == nil && w.initialDefaultKAS != nil { + cfg.defaultKAS = w.initialDefaultKAS + } + if cfg.mimeType == "" { + cfg.mimeType = defaultMimeType + } + return cfg, nil +} + +// chunkedTotals are the byte counts across the segments the manifest +// describes. They are reported for information only; neither is the number of +// bytes the caller must append, which is what the TDFData readers yield and +// nothing here tracks. See ChunkedFinalizeResult.EncryptedSize. +type chunkedTotals struct { + // encrypted is the ciphertext length across the segments in the manifest. + encrypted int64 + + // plaintext is the plaintext length across the segments in the manifest. + plaintext int64 +} + +// 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, chunkedTotals, error) { + var totals chunkedTotals + order, err := w.segmentOrderLocked(cfg.keepSegments) + if err != nil { + return nil, totals, err + } + + // Hand the splitter copies of both slices. A splitter that zeroes or + // rewrites the DEK it is given -- scrubbing what it thinks is its own + // working buffer, say -- would desynchronize the DEK from the segment + // signatures already computed against it and from the root signature + // computed below, producing a TDF that fails verification exactly as a + // tampered one would. + // + // The attributes need the same defense for a different reason: on the + // fallback path cfg.attributes *is* w.initialAttributes, the writer's own + // retained slice, so a splitter that rewrites an element changes the policy + // buildChunkedPolicy writes just below -- and every later Finalize and + // GetManifest on this writer. + splits, err := w.splitter.Split(ctx, slices.Clone(cfg.attributes), slices.Clone(w.dek), cfg.defaultKAS) + if err != nil { + return nil, totals, err + } + policyBytes, err := buildChunkedPolicy(cfg.attributes) + if err != nil { + return nil, totals, err + } + kaos, err := buildChunkedKeyAccessObjects(splits, w.dek, policyBytes, cfg.encryptedMetadata) + if err != nil { + return nil, totals, err + } + + encInfo := EncryptionInformation{ + KeyAccessObjs: kaos, + KeyAccessType: kSplitKeyType, + Policy: string(ocrypto.Base64Encode(policyBytes)), + Method: Method{ + Algorithm: kGCMCipherAlgorithm, + IsStreamable: true, + }, + IntegrityInformation: IntegrityInformation{ + SegmentHashAlgorithm: SegmentGMAC.String(), + Segments: make([]Segment, len(order)), + }, + } + + var aggregate bytes.Buffer + for i, idx := range order { + // segmentOrderLocked only ever names written slots, whether it + // derived the order itself or validated a caller-supplied one. + slot, ok := w.segments[idx] + if !ok || !slot.written { + return nil, totals, fmt.Errorf("segment %d not written; cannot finalize", idx) + } + if slot.seg.Hash == "" { + return nil, totals, fmt.Errorf("segment %d has empty hash", idx) + } + encInfo.Segments[i] = slot.seg + totals.plaintext += slot.seg.Size + totals.encrypted += slot.seg.EncryptedSize + decoded, err := ocrypto.Base64Decode([]byte(slot.seg.Hash)) + if err != nil { + return nil, totals, fmt.Errorf("decode segment %d hash: %w", idx, err) + } + aggregate.Write(decoded) + } + + if len(order) > 0 { + if first, ok := w.segments[order[0]]; ok { + encInfo.DefaultEncryptedSegSize = first.seg.EncryptedSize + encInfo.DefaultSegmentSize = first.seg.Size + } + } + + rootSig, err := rootIntegrity(aggregate.Bytes(), w.dek, RootHS256, w.useHex) + if err != nil { + return nil, totals, err + } + encInfo.RootSignature = RootSignature{ + Algorithm: RootHS256.String(), + Signature: string(ocrypto.Base64Encode([]byte(rootSig))), + } + + // Assertions bind to the same aggregate hash the root signature + // covers, so they can only be signed once every segment is in. + assertions, err := signAssertions(aggregate.Bytes(), cfg.assertions, w.dek, w.useHex) + if err != nil { + return nil, totals, err + } + + manifest := &Manifest{ + Assertions: assertions, + EncryptionInformation: encInfo, + Payload: Payload{ + IsEncrypted: true, + MimeType: cfg.mimeType, + Protocol: tdfAsZip, + Type: tdfZipReference, + URL: zipstream.TDFPayloadFileName, + }, + } + if !cfg.excludeVersion { + manifest.TDFVersion = TDFSpecVersion + } + return manifest, totals, nil +} + +// segmentOrderLocked returns the emission order given the current +// writer state and an optional keepSegments subset. Caller holds mu. +// +// With no subset, every written segment is emitted in ascending index +// order. Indices that are merely reserved -- WriteSegment has claimed +// them but the archive has not accepted their bytes yet -- are not +// written and are skipped. That is what lets GetManifest run +// concurrently with in-flight writes, which is its whole purpose: a +// reservation describes bytes that may never exist, so including one +// would make every such call fail on a segment that is not late, just +// unfinished. +// +// A supplied subset must be a prefix of that same ascending sequence. +// Note this constrains position, not value: the written indices +// themselves may be sparse (a caller reserving a block of indices per +// upload part and filling only part of each block writes e.g. +// 0,1,5000,5001), and any such set is accepted so long as the subset +// names its members in order and drops only from the end. Whether +// index 0 is among them is Finalize's business, not this function's: +// GetManifest shares this path and legitimately runs before segment 0 +// has been written. +// +// Both halves of that rule are forced by the archive layout, which +// 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. +func (w *chunkedWriter) segmentOrderLocked(keep []int) ([]int, error) { + written := make([]int, 0, len(w.segments)) + for idx, slot := range w.segments { + if !slot.written { + continue + } + written = append(written, idx) + } + slices.Sort(written) + if len(keep) == 0 { + return written, nil + } + if len(keep) > len(written) { + return nil, fmt.Errorf("WithChunkedSegments names %d segments but only %d were written", len(keep), len(written)) + } + for i, idx := range keep { + if idx == written[i] { + continue + } + // A reserved-but-unwritten index reads as "not written" here, + // the same as one never claimed at all: naming it in the + // manifest would describe bytes the archive has not accepted. + if slot, ok := w.segments[idx]; !ok || !slot.written { + return nil, fmt.Errorf("WithChunkedSegments references segment %d which was not written", idx) + } + return nil, fmt.Errorf( + "WithChunkedSegments must name written segments in ascending index order and may drop only from the end; got %d at position %d where %d was expected", + idx, i, written[i], + ) + } + out := make([]int, len(keep)) + copy(out, keep) + return out, nil +} + +// buildChunkedKeyAccessObjects wraps each split share to each KAS +// listed by the splitter. +func buildChunkedKeyAccessObjects(splits *SplitResult, dek, policyBytes []byte, metadata string) ([]KeyAccess, error) { + // This is the one place caller-supplied split data is turned into + // manifest content, so it is where the splitter's contract is + // enforced -- for the default splitter and for anything injected + // through WithChunkedKeySplitter alike. Validate guarantees every + // invariant the loop below relies on: at least one split, every + // split naming at least one KAS, every named KAS resolving to a + // usable wrapping key, and split ids that the reader can group on. + if err := splits.Validate(); err != nil { + return nil, err + } + // Validate first, then this: Validate names the specific structural fault + // (no splits, empty share, duplicate id, unresolved KAS), all of which + // VerifyReconstruction can only report as a length or value mismatch. + if err := splits.VerifyReconstruction(dek); err != nil { + return nil, err + } + base64Policy := ocrypto.Base64Encode(policyBytes) + + out := make([]KeyAccess, 0, len(splits.Splits)) + for _, split := range splits.Splits { + // Policy binding and metadata are keyed on the split share, not + // on the KAS, so compute them once per split rather than once + // per KAS URL in an OR-group. + policyBinding := createPolicyBinding(split.Data, base64Policy) + var encMeta string + if metadata != "" { + m, err := encryptMetadata(split.Data, metadata) + if err != nil { + return nil, fmt.Errorf("encrypt metadata for split %s: %w", split.ID, err) + } + encMeta = m + } + for _, url := range split.KASURLs { + // Validate resolved every URL, so the lookup cannot miss. + pk := splits.KASPublicKeys[url] + kao, err := createKeyAccess(pk.toKASInfo(), split.Data, policyBinding, encMeta, split.ID) + if err != nil { + return nil, fmt.Errorf("wrap key for %s: %w", url, err) + } + out = append(out, kao) + } + } + return out, nil +} + +// buildChunkedPolicy composes the TDF Policy document from attribute +// values. +func buildChunkedPolicy(values []*policy.Value) ([]byte, error) { + p := PolicyObject{UUID: uuid.NewString()} + p.Body.DataAttributes = make([]attributeObject, 0, len(values)) + p.Body.Dissem = make([]string, 0) + for i, v := range values { + // GetFqn is nil-receiver safe, so a nil element or a value that never + // had its FQN populated yields "" rather than panicking. Emitting that + // writes {"attribute": ""} into the bound policy: the PDP denies it, so + // the TDF fails closed rather than open, but it fails at decrypt with + // nothing naming the cause -- and it does so after the caller has + // encrypted and uploaded the whole payload. + fqn := v.GetFqn() + if fqn == "" { + return nil, fmt.Errorf("chunked: attribute value at index %d has an empty FQN; the policy it produces names no attribute and no KAS could evaluate it", i) + } + p.Body.DataAttributes = append(p.Body.DataAttributes, attributeObject{ + Attribute: fqn, + }) + } + return json.Marshal(p) +} + +// cloneChunkedManifest copies the manifest struct and clones its three +// slices, which is what makes the result safe to hand to a caller: it +// can append to or overwrite elements of Segments, KeyAccessObjs and +// Assertions without the writer's copy changing. The clones are +// shallow, so a caller reaching into a pointer or map inside an +// element -- an Assertion's binding, a KeyAccess's policy binding -- +// still shares that state. Nothing in this package hands out a +// manifest a caller has any reason to mutate that deeply. +func cloneChunkedManifest(in *Manifest) *Manifest { + if in == nil { + return nil + } + out := *in + if in.KeyAccessObjs != nil { + out.KeyAccessObjs = slices.Clone(in.KeyAccessObjs) + } + if in.Segments != nil { + out.Segments = slices.Clone(in.Segments) + } + if in.Assertions != nil { + out.Assertions = slices.Clone(in.Assertions) + } + return &out +} diff --git a/sdk/internal/zipstream/segment_writer.go b/sdk/internal/zipstream/segment_writer.go index 34b3da57ac..420473a16d 100644 --- a/sdk/internal/zipstream/segment_writer.go +++ b/sdk/internal/zipstream/segment_writer.go @@ -12,6 +12,43 @@ import ( "sync" ) +// The three constructors below exist so the mutation boundary inside +// segmentWriter.Finalize is greppable rather than inferred. Finalize's first +// irreversible change is SetOrder, which is retained and skips re-derivation on +// a later call; its first archive-corrupting change is adding the payload entry +// to the central directory. The reviewable rule they buy is: every return at or +// below the SetOrder call uses finalizeErrMutated. +// +// The sentinels cannot substitute for the flag -- ErrSegmentMissing and +// ErrZip64Required are each returned from both sides of the boundary, so +// matching on them classifies nothing. + +// writerTypeSegment is the Error.Type every error from this file carries. +const writerTypeSegment = "segment" + +// finalizeErr reports a Finalize failure that left the writer untouched. Every +// return site lexically above the SetOrder call uses this one. +func finalizeErr(err error) *Error { + return &Error{Op: "finalize", Type: writerTypeSegment, Err: err} +} + +// finalizeErrMutated reports a Finalize failure that occurred after the writer +// had already changed its own state -- the derived segment order is retained, +// or the payload entry is already in the central directory. Retrying compounds +// the change; a caller that fences on this is doing the right thing. +func finalizeErrMutated(err error) *Error { + return &Error{Op: "finalize", Type: writerTypeSegment, Err: err, Mutated: true} +} + +// writeSegmentErr reports a WriteSegment failure. There is deliberately no +// mutated variant: AddSegment validates before it records, and the payload size +// accounting runs only after it has returned, so no path out of WriteSegment +// leaves the writer half-changed. A future mutation added above a return site +// must add one. +func writeSegmentErr(err error) *Error { + return &Error{Op: "write-segment", Type: writerTypeSegment, Err: err} +} + // segmentWriter implements the SegmentWriter interface for out-of-order segment writing type segmentWriter struct { *baseWriter @@ -57,27 +94,27 @@ func (sw *segmentWriter) WriteSegment(ctx context.Context, index int, size uint6 // Check if writer is closed or finalized if err := sw.checkClosed(); err != nil { - return nil, &Error{Op: "write-segment", Type: "segment", Err: err} + return nil, writeSegmentErr(err) } if sw.finalized { - return nil, &Error{Op: "write-segment", Type: "segment", Err: ErrWriterClosed} + return nil, writeSegmentErr(ErrWriterClosed) } // Validate segment index (allow dynamic expansion for streaming use cases) if index < 0 { - return nil, &Error{Op: "write-segment", Type: "segment", Err: ErrInvalidSegment} + return nil, writeSegmentErr(ErrInvalidSegment) } // Check for duplicate segment if _, exists := sw.metadata.Segments[index]; exists { - return nil, &Error{Op: "write-segment", Type: "segment", Err: ErrDuplicateSegment} + return nil, writeSegmentErr(ErrDuplicateSegment) } // Check context cancellation select { case <-ctx.Done(): - return nil, &Error{Op: "write-segment", Type: "segment", Err: ctx.Err()} + return nil, writeSegmentErr(ctx.Err()) default: } @@ -90,14 +127,14 @@ func (sw *segmentWriter) WriteSegment(ctx context.Context, index int, size uint6 if index == 0 { // Segment 0: Write local file header + encrypted data if err := sw.writeLocalFileHeader(buffer); err != nil { - return nil, &Error{Op: "write-segment", Type: "segment", Err: err} + return nil, writeSegmentErr(err) } } // Record segment metadata only (no payload retention). Payload bytes are returned // to the caller and may be uploaded; we keep only CRC and size for finalize. if err := sw.metadata.AddSegment(index, size, crc32); err != nil { - return nil, &Error{Op: "write-segment", Type: "segment", Err: err} + return nil, writeSegmentErr(err) } // Update payload entry metadata @@ -113,26 +150,29 @@ func (sw *segmentWriter) Finalize(ctx context.Context, manifest []byte) ([]byte, sw.mu.Lock() defer sw.mu.Unlock() + // Every return from here down to the SetOrder call is finalizeErr: these + // are all pure reads, so the writer is untouched and the caller may retry. + // Check if writer is closed or already finalized if err := sw.checkClosed(); err != nil { - return nil, &Error{Op: "finalize", Type: "segment", Err: err} + return nil, finalizeErr(err) } if sw.finalized { - return nil, &Error{Op: "finalize", Type: "segment", Err: ErrWriterClosed} + return nil, finalizeErr(ErrWriterClosed) } // Check context cancellation select { case <-ctx.Done(): - return nil, &Error{Op: "finalize", Type: "segment", Err: ctx.Err()} + return nil, finalizeErr(ctx.Err()) default: } // Nothing arrived at all: report the general incomplete-input error // rather than the segment-0-specific one below. if len(sw.metadata.Segments) == 0 { - return nil, &Error{Op: "finalize", Type: "segment", Err: ErrSegmentMissing} + return nil, finalizeErr(ErrSegmentMissing) } // Only segment 0 emits the payload's local file header, and every offset @@ -149,9 +189,14 @@ func (sw *segmentWriter) Finalize(ctx context.Context, manifest []byte) ([]byte, // absence either, since that derived order is self-consistent by // construction. if _, ok := sw.metadata.Segments[0]; !ok { - return nil, &Error{Op: "finalize", Type: "segment", Err: ErrNoSegmentZero} + return nil, finalizeErr(ErrNoSegmentZero) } + // SetOrder is the mutation boundary: from here down, every return is + // finalizeErrMutated. Order is retained and the check above skips + // re-derivation on a later call, so a retry inherits an order that predates + // any segment added since. + // If no explicit order was provided, derive order from present indices (sorted). if len(sw.metadata.Order) == 0 { order := make([]int, 0, len(sw.metadata.Segments)) @@ -160,8 +205,12 @@ func (sw *segmentWriter) Finalize(ctx context.Context, manifest []byte) ([]byte, } sort.Ints(order) if err := sw.metadata.SetOrder(order); err != nil { - // This should be an unreachable state, but handle it defensively. - return nil, &Error{Op: "finalize", Type: "segment", Err: fmt.Errorf("internal error setting segment order: %w", err)} + // Mutated, even though SetOrder validates the whole order before + // assigning and so leaves Order nil on failure. Marking it true + // keeps the positional rule exception-free at no cost: the site is + // unreachable (the order is derived sorted and deduplicated) and + // true is the fail-safe direction. Do not "correct" this. + return nil, finalizeErrMutated(fmt.Errorf("internal error setting segment order: %w", err)) } } @@ -170,7 +219,7 @@ func (sw *segmentWriter) Finalize(ctx context.Context, manifest []byte) ([]byte, // complete by construction, and the empty set already returned. Kept for // a future caller that supplies an explicit order. if !sw.metadata.IsComplete() { - return nil, &Error{Op: "finalize", Type: "segment", Err: ErrSegmentMissing} + return nil, finalizeErrMutated(ErrSegmentMissing) } // Compute final CRC32 by combining per-segment CRCs now that all are present @@ -203,13 +252,17 @@ func (sw *segmentWriter) Finalize(ctx context.Context, manifest []byte) ([]byte, // 1. Write data descriptor for payload (fail if Zip64Never but required) if sw.config.Zip64 == Zip64Never && needZip64ForPayload { - return nil, &Error{Op: "finalize", Type: "segment", Err: ErrZip64Required} + return nil, finalizeErrMutated(ErrZip64Required) } if err := sw.writeDataDescriptor(buffer, needZip64ForPayload); err != nil { - return nil, &Error{Op: "finalize", Type: "segment", Err: err} + return nil, finalizeErrMutated(err) } - // 2. Update payload entry CRC32 and add to central directory + // 2. Update payload entry CRC32 and add to central directory. + // + // Second boundary: below this, a retry calls AddFile again and produces a + // central directory with two entries for one payload -- an archive some + // readers accept and silently misread. sw.payloadEntry.CRC32 = sw.metadata.TotalCRC32 sw.centralDir.AddFile(*sw.payloadEntry) @@ -225,7 +278,7 @@ func (sw *segmentWriter) Finalize(ctx context.Context, manifest []byte) ([]byte, } if err := sw.writeManifestFile(buffer, manifest, manifestEntry); err != nil { - return nil, &Error{Op: "finalize", Type: "segment", Err: err} + return nil, finalizeErrMutated(err) } // 4. Add manifest entry to central directory @@ -239,15 +292,15 @@ func (sw *segmentWriter) Finalize(ctx context.Context, manifest []byte) ([]byte, sw.centralDir.Offset > maxNonZip64 || len(sw.centralDir.Entries) >= zip64MagicVal16 if sw.config.Zip64 == Zip64Never && needZip64ForCD { - return nil, &Error{Op: "finalize", Type: "segment", Err: ErrZip64Required} + return nil, finalizeErrMutated(ErrZip64Required) } cdBytes, err := sw.centralDir.GenerateBytes(needZip64ForCD) if err != nil { - return nil, &Error{Op: "finalize", Type: "segment", Err: err} + return nil, finalizeErrMutated(err) } if _, err := buffer.Write(cdBytes); err != nil { - return nil, &Error{Op: "finalize", Type: "segment", Err: err} + return nil, finalizeErrMutated(err) } sw.finalized = true diff --git a/sdk/internal/zipstream/segment_writer_test.go b/sdk/internal/zipstream/segment_writer_test.go index a4319e66f5..d08546d86e 100644 --- a/sdk/internal/zipstream/segment_writer_test.go +++ b/sdk/internal/zipstream/segment_writer_test.go @@ -798,3 +798,216 @@ func TestWriterClampsOutOfRangeClock(t *testing.T) { "%s should be clamped to the base year, not wrapped into the future", f.Name) } } + +// TestSegmentWriterErrorMutated pins Error.Mutated across every refusal a +// caller can provoke through the public surface. +// +// The flag is the entire basis on which a wrapping writer decides between +// retrying and fencing, and it is not recoverable from anything else in the +// error: ErrSegmentMissing and ErrZip64Required are each returned from both +// sides of the mutation boundary, so matching on the sentinel classifies +// nothing. Getting it backwards is silent in both directions -- false when it +// should be true lets a retry append a second payload entry to the central +// directory, yielding an archive some readers accept and misread; true when it +// should be false discards an already-encrypted, already-uploaded payload over +// a refusal that changed nothing. +func TestSegmentWriterErrorMutated(t *testing.T) { + // A segment whose recorded size alone crosses the ZIP64 threshold. Only + // the size is recorded, not the bytes, so the threshold can be crossed + // without allocating anything. + const smallSegment = 4 + + for _, tc := range []struct { + name string + // provoke returns the error to classify, having driven the writer + // into the state that produces it. + provoke func(t *testing.T, ctx context.Context) error + wantErr error + wantMutated bool + }{ + { + name: "closed writer", + provoke: func(t *testing.T, ctx context.Context) error { + w := NewSegmentTDFWriter(1) + require.NoError(t, w.Close()) + _, err := w.Finalize(ctx, []byte(`{}`)) + return err + }, + wantErr: ErrWriterClosed, + wantMutated: false, + }, + { + name: "already finalized", + provoke: func(t *testing.T, ctx context.Context) error { + w := NewSegmentTDFWriter(1) + _, err := w.WriteSegment(ctx, 0, smallSegment, 0) + require.NoError(t, err) + _, err = w.Finalize(ctx, []byte(`{}`)) + require.NoError(t, err) + _, err = w.Finalize(ctx, []byte(`{}`)) + return err + }, + wantErr: ErrWriterClosed, + wantMutated: false, + }, + { + name: "cancelled context", + provoke: func(t *testing.T, ctx context.Context) error { + w := NewSegmentTDFWriter(1) + _, err := w.WriteSegment(ctx, 0, smallSegment, 0) + require.NoError(t, err) + cancelled, cancel := context.WithCancel(ctx) + cancel() + _, err = w.Finalize(cancelled, []byte(`{}`)) + return err + }, + wantErr: context.Canceled, + wantMutated: false, + }, + { + name: "no segments at all", + provoke: func(_ *testing.T, ctx context.Context) error { + _, err := NewSegmentTDFWriter(1).Finalize(ctx, []byte(`{}`)) + return err + }, + wantErr: ErrSegmentMissing, + wantMutated: false, + }, + { + // The companion to the mutated ErrSegmentMissing below: same + // sentinel, opposite verdict, which is why the flag exists. + name: "missing segment zero", + provoke: func(t *testing.T, ctx context.Context) error { + w := NewSegmentTDFWriter(2) + _, err := w.WriteSegment(ctx, 1, smallSegment, 0) + require.NoError(t, err) + _, err = w.Finalize(ctx, []byte(`{}`)) + return err + }, + wantErr: ErrNoSegmentZero, + wantMutated: false, + }, + { + // Past SetOrder: the derived order is retained, so a retry that + // supplied the missing bytes would finalize against an order + // computed before they arrived. + name: "payload needs zip64 under Zip64Never", + provoke: func(t *testing.T, ctx context.Context) error { + w := NewSegmentTDFWriter(1, + WithZip64Mode(Zip64Never), + WithMaxNonZip64Value(smallSegment-1)) + _, err := w.WriteSegment(ctx, 0, smallSegment, 0) + require.NoError(t, err) + _, err = w.Finalize(ctx, []byte(`{}`)) + return err + }, + wantErr: ErrZip64Required, + wantMutated: true, + }, + { + // Past the second boundary: the payload entry is already in the + // central directory, so a retry adds a duplicate. The payload + // itself fits under the threshold here; only the directory's + // offset, which the local file header pushes past it, does not. + name: "central directory offset needs zip64 under Zip64Never", + provoke: func(t *testing.T, ctx context.Context) error { + w := NewSegmentTDFWriter(1, + WithZip64Mode(Zip64Never), + WithMaxNonZip64Value(smallSegment+1)) + _, err := w.WriteSegment(ctx, 0, smallSegment, 0) + require.NoError(t, err) + _, err = w.Finalize(ctx, []byte(`{}`)) + return err + }, + wantErr: ErrZip64Required, + wantMutated: true, + }, + { + name: "write segment on a closed writer", + provoke: func(t *testing.T, ctx context.Context) error { + w := NewSegmentTDFWriter(1) + require.NoError(t, w.Close()) + _, err := w.WriteSegment(ctx, 0, smallSegment, 0) + return err + }, + wantErr: ErrWriterClosed, + wantMutated: false, + }, + { + name: "duplicate segment", + provoke: func(t *testing.T, ctx context.Context) error { + w := NewSegmentTDFWriter(1) + _, err := w.WriteSegment(ctx, 0, smallSegment, 0) + require.NoError(t, err) + _, err = w.WriteSegment(ctx, 0, smallSegment, 0) + return err + }, + wantErr: ErrDuplicateSegment, + wantMutated: false, + }, + { + name: "negative segment index", + provoke: func(_ *testing.T, ctx context.Context) error { + _, err := NewSegmentTDFWriter(1).WriteSegment(ctx, -1, smallSegment, 0) + return err + }, + wantErr: ErrInvalidSegment, + wantMutated: false, + }, + } { + t.Run(tc.name, func(t *testing.T) { + err := tc.provoke(t, t.Context()) + require.ErrorIs(t, err, tc.wantErr) + + // Asserted as a typed *Error, not via any interface: a caller can + // only read the flag if the concrete type survives wrapping. + var zerr *Error + require.ErrorAs(t, err, &zerr) + assert.Equal(t, writerTypeSegment, zerr.Type) + assert.Equal(t, tc.wantMutated, zerr.Mutated, + "Mutated decides retry vs. fence; %q is on the %s side of the boundary", + tc.name, map[bool]string{true: "mutated", false: "untouched"}[tc.wantMutated]) + }) + } +} + +// TestSegmentWriterNonMutatingFinalizeIsRetryable checks the promise Mutated +// false actually makes. Asserting the flag alone would pass just as well if +// the writer were quietly wrecked, so this drives a refused Finalize through +// to a successful one and reads the archive back. +func TestSegmentWriterNonMutatingFinalizeIsRetryable(t *testing.T) { + ctx := t.Context() + w := NewSegmentTDFWriter(2) + + payload := []byte("second segment arrives late") + tail, err := w.WriteSegment(ctx, 1, uint64(len(payload)), crc32.ChecksumIEEE(payload)) + require.NoError(t, err) + + _, err = w.Finalize(ctx, []byte(`{"manifest":true}`)) + require.ErrorIs(t, err, ErrNoSegmentZero) + var zerr *Error + require.ErrorAs(t, err, &zerr) + require.False(t, zerr.Mutated) + + // Segment 0 arrives afterwards. Had the refused Finalize retained a + // derived order, this segment would be absent from it and the archive + // below would be short one entry or fail outright. + head := []byte("first ") + header, err := w.WriteSegment(ctx, 0, uint64(len(head)), crc32.ChecksumIEEE(head)) + require.NoError(t, err) + + final, err := w.Finalize(ctx, []byte(`{"manifest":true}`)) + require.NoError(t, err) + + all := bytes.Join([][]byte{header, head, tail, payload, final}, nil) + zipReader, err := zip.NewReader(bytes.NewReader(all), int64(len(all))) + require.NoError(t, err, "the retried finalize must produce a readable archive") + require.Len(t, zipReader.File, 2) + + f, err := zipReader.Open(TDFPayloadFileName) + require.NoError(t, err) + t.Cleanup(func() { assert.NoError(t, f.Close()) }) + got, err := io.ReadAll(f) + require.NoError(t, err) + assert.Equal(t, append(append([]byte{}, head...), payload...), got) +} diff --git a/sdk/internal/zipstream/writer.go b/sdk/internal/zipstream/writer.go index 3763fedc91..e4c73f4437 100644 --- a/sdk/internal/zipstream/writer.go +++ b/sdk/internal/zipstream/writer.go @@ -48,6 +48,20 @@ type Error struct { Op string // Operation that failed Type string // Writer type: "sequential", "streaming", "segment" Err error // Underlying error + + // Mutated reports whether the writer had already changed its own state + // when the operation failed. False means the writer is byte-for-byte as it + // was before the call and the caller may retry; true means the partial + // change is not rolled back and a retry compounds it rather than + // recovering from it. + // + // Only Finalize sets it. WriteSegment validates everything before it + // mutates anything, so every error it returns is Mutated false. + // + // An implementation that wraps a zipstream writer must not pass the + // delegate's error through unchanged if it performed mutations of its own: + // the flag describes the whole operation, not one layer of it. + Mutated bool } func (e *Error) Error() string { diff --git a/sdk/key_splitter.go b/sdk/key_splitter.go new file mode 100644 index 0000000000..e25002c20e --- /dev/null +++ b/sdk/key_splitter.go @@ -0,0 +1,308 @@ +package sdk + +import ( + "context" + "crypto/subtle" + "errors" + "fmt" + + "github.com/opentdf/platform/lib/ocrypto" + "github.com/opentdf/platform/protocol/go/policy" +) + +// KeySplitter converts attribute values plus a DEK into one or more +// key splits, each addressed to one or more KAS servers. Injected on +// the chunked Writer so tests can substitute an identity splitter +// without touching real attribute grants. +// +// Experimental: not part of the stable SDK API; may change or be removed. +type KeySplitter interface { + // Split evaluates the ABAC policy expressed by attrs, produces N + // splits of dek per the resulting boolean expression, and returns + // each split alongside the KAS public keys it must be wrapped to. + Split(ctx context.Context, attrs []*policy.Value, dek []byte, defaultKAS *policy.SimpleKasKey) (*SplitResult, error) +} + +// Split is one XOR share of the DEK bound to one or more KAS +// servers. +// +// Experimental: not part of the stable SDK API; may change or be removed. +type Split struct { + // Data is this share of the DEK: the value that, XOR'd with every + // other share in the result, reproduces the DEK. With a single + // split -- what DefaultKeySplitter always produces -- there is + // nothing to XOR against, so it is the DEK verbatim. + Data []byte + + // ID uniquely identifies this split within a SplitResult. It is + // what the reader groups shares on, so it is required once there + // is more than one split and must be distinct across them. With a + // single split there is nothing to distinguish, and it is + // conventionally left empty. + ID string + + // KASURLs lists every KAS that can unwrap this split. Multiple + // URLs mean any one KAS is sufficient (OR semantics). + KASURLs []string +} + +// SplitResult is what KeySplitter.Split returns: the shares plus the +// KAS wrapping keys needed to encrypt each share into a KeyAccess +// object. +// +// Experimental: not part of the stable SDK API; may change or be removed. +type SplitResult struct { + // KASPublicKeys maps KAS URL to the wrapping key to use for that + // URL. Populated for every URL referenced by any split. + KASPublicKeys map[string]KASPublicKey + + // Splits are the DEK shares in emission order. + Splits []Split +} + +// Validate reports whether the result can be turned into a key access +// array that a reader can actually reconstruct the DEK from. The +// chunked writer calls it before wrapping anything, and third-party +// KeySplitter implementations can call it as a self-check. +// +// Every rule here exists because breaking it produces a TDF that is +// accepted at creation and fails later in a way that does not name the +// cause. The reader derives the set of splits it must collect from the +// key access objects present in the manifest, so a share that never +// reached the manifest is invisible to its completeness check: it XORs +// what it has, reconstructs the wrong DEK, and reports a root +// signature failure, which reads as tampering. By then the plaintext +// is gone. +// +// Validate sees only the shares, so it is half the contract: it can +// confirm they agree with each other, but not that they agree with the +// key. VerifyReconstruction is the other half. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func (r *SplitResult) Validate() error { + if r == nil || len(r.Splits) == 0 { + return errors.New("chunked: splitter returned no splits") + } + + ids := make(map[string]struct{}, len(r.Splits)) + shareLen := len(r.Splits[0].Data) + for i, split := range r.Splits { + // A share is one XOR operand. A short or empty one silently + // leaves part of the DEK unmasked, since the reader XORs only + // as many bytes as the share carries. + if len(split.Data) == 0 { + return fmt.Errorf("chunked: split %d (id %q) has no share data", i, split.ID) + } + if len(split.Data) != shareLen { + return fmt.Errorf("chunked: split %d (id %q) has a %d-byte share; split 0 has %d and all shares must agree", + i, split.ID, len(split.Data), shareLen) + } + + // A single split carries no sid, because there is nothing to + // distinguish it from. With more than one, sid is what the + // reader groups on: an empty or duplicated id collapses two + // shares into one group, and the second is never XOR'd in. + if len(r.Splits) > 1 && split.ID == "" { + return fmt.Errorf("chunked: split %d has an empty id; ids are required when there is more than one split", i) + } + if _, dup := ids[split.ID]; dup { + return fmt.Errorf("chunked: split id %q is used by more than one split", split.ID) + } + ids[split.ID] = struct{}{} + + if len(split.KASURLs) == 0 { + return fmt.Errorf("chunked: split %d (id %q) names no KAS; its share could never be unwrapped", i, split.ID) + } + for _, url := range split.KASURLs { + pk, ok := r.KASPublicKeys[url] + if !ok { + // The same sentinel the mainline writer uses for an + // unresolved KAS, so callers can match one error whichever + // path produced it. + return fmt.Errorf("chunked: splitID:[%s], kas:[%s]: no entry in KASPublicKeys: %w", split.ID, url, errKasPubKeyMissing) + } + if err := pk.validate(url); err != nil { + return err + } + } + } + return nil +} + +// VerifyReconstruction reports whether the shares actually rebuild +// dek. It is the half of the splitter's contract Validate cannot +// check: Validate sees only the shares, so it can confirm they agree +// with each other but not that they agree with the key. +// +// Mirrors the self-check the XOR splitter performs on its own output. +// A splitter that returns shares of the wrong length, or of the right +// length and the wrong value, produces a TDF that wraps, uploads and +// rewraps without complaint, then fails at the reader's root signature +// check -- reported as tampering, after the plaintext is gone. +// +// With a single split there is nothing to XOR against, so the share +// must equal dek verbatim; the loop below expresses that as the +// degenerate case rather than special-casing it. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func (r *SplitResult) VerifyReconstruction(dek []byte) error { + if r == nil || len(r.Splits) == 0 { + return errors.New("chunked: splitter returned no splits") + } + if len(dek) == 0 { + return errors.New("chunked: cannot verify split reconstruction against an empty DEK") + } + + reconstructed := make([]byte, len(dek)) + for i, split := range r.Splits { + // Length is checked against the DEK, not against Splits[0] as Validate + // does: shares that agree with each other but are shorter than the key + // leave its tail unmasked, and the reader XORs only as many bytes as + // each share carries. + if len(split.Data) != len(dek) { + return fmt.Errorf("chunked: split %d (id %q) has a %d-byte share; every share must be %d bytes, the length of the DEK", + i, split.ID, len(split.Data), len(dek)) + } + subtle.XORBytes(reconstructed, reconstructed, split.Data) + } + + // Constant-time, and deliberately silent about where the mismatch is. The + // XOR splitter's equivalent reports the first differing byte index, which + // is a prefix-match oracle over key material in exchange for nothing: the + // caller cannot act on the offset. + if subtle.ConstantTimeCompare(reconstructed, dek) != 1 { + return errors.New("chunked: the splits do not XOR back to the DEK; a reader that collected every share would reconstruct the wrong key and report a root signature failure") + } + return nil +} + +// KASPublicKey is the wrapping key resolved for one KAS URL. +// +// Experimental: not part of the stable SDK API; may change or be removed. +type KASPublicKey struct { + // Algorithm identifies the wrapping scheme. It must be one of the + // values ocrypto.ParseKeyType accepts, e.g. ocrypto.RSA2048Key or + // ocrypto.EC256Key -- use PolicyAlgorithmToKeyType to derive it + // from a policy.Algorithm. An unrecognized value falls through + // createKeyAccess's RSA default, which sniffs the PEM rather than + // honoring the declared scheme, and produces a KAO that cannot be + // decrypted; SplitResult.Validate rejects it first. + Algorithm ocrypto.KeyType + + // KID identifies which key at that KAS to use. + KID string + + // PEM is the wrapping key in PEM form. + PEM string + + // URL of the KAS. + URL string +} + +// validate checks one resolved wrapping key. mapKey is the +// KASPublicKeys key it was found under, which is reported in errors +// and must match URL: the KAO is built from URL, so the two +// disagreeing means the split names one KAS and the manifest points a +// reader at another. +func (k KASPublicKey) validate(mapKey string) error { + if k.PEM == "" { + return fmt.Errorf("chunked: kas:[%s]: %w", mapKey, errKasPubKeyMissing) + } + if k.URL == "" { + return fmt.Errorf("chunked: kas %q has an empty URL field; every key access object built from it would name no endpoint", mapKey) + } + if k.URL != mapKey { + return fmt.Errorf("chunked: kas keyed under %q carries URL %q; key access objects are built from the URL field", mapKey, k.URL) + } + if _, err := ocrypto.ParseKeyType(string(k.Algorithm)); err != nil { + return fmt.Errorf("chunked: kas %q: %w", mapKey, err) + } + return nil +} + +// toKASInfo adapts the splitter's wrapping-key descriptor to the +// KASInfo shape consumed by createKeyAccess. Default is not carried +// over; it plays no part in building a key access object. +func (k KASPublicKey) toKASInfo() KASInfo { + return KASInfo{ + URL: k.URL, + PublicKey: k.PEM, + KID: k.KID, + Algorithm: string(k.Algorithm), + } +} + +// ErrSplitterRequiresDefaultKAS is returned by the default key +// splitter when no default KAS was supplied. The default splitter is +// single-KAS only; multi-attribute splits require injecting a full +// splitter via WithChunkedKeySplitter. +var ErrSplitterRequiresDefaultKAS = errors.New("chunked: default splitter requires a default KAS; supply WithChunkedDefaultKAS or WithChunkedKeySplitter") + +// ErrSplitterUnsupportedAlgorithm is returned by the default key +// splitter when the default KAS advertises a key algorithm this SDK +// has no wrapping scheme for. +var ErrSplitterUnsupportedAlgorithm = errors.New("chunked: unsupported KAS key algorithm") + +// DefaultKeySplitter returns a single-KAS single-split splitter. +// Attributes are ignored; the entire DEK is bound to the caller's +// default KAS. Callers with attribute-based key splits requirements +// should inject their own splitter via WithChunkedKeySplitter. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func DefaultKeySplitter() KeySplitter { return &singleKASSplitter{} } + +// singleKASSplitter binds the full DEK to a single KAS. Attributes +// are ignored; splitting into multi-KAS OR-of-AND shares is beyond +// this default's scope. +type singleKASSplitter struct{} + +// Split returns one split covering the full DEK, addressed to +// defaultKAS. Errors when defaultKAS is nil, has no public key or +// URI, or names an algorithm this SDK cannot wrap for. +func (s *singleKASSplitter) Split(_ context.Context, _ []*policy.Value, dek []byte, defaultKAS *policy.SimpleKasKey) (*SplitResult, error) { + if defaultKAS == nil || defaultKAS.GetPublicKey() == nil || defaultKAS.GetPublicKey().GetPem() == "" { + return nil, ErrSplitterRequiresDefaultKAS + } + url := defaultKAS.GetKasUri() + if url == "" { + // An empty URI would still produce a PEM-valid split, but the + // resulting KeyAccess.KasURL leaves a reader with no endpoint to + // send a rewrap request to. + return nil, fmt.Errorf("%w: kas uri is empty", ErrSplitterRequiresDefaultKAS) + } + + // Reject an unmappable algorithm here rather than letting the empty + // string formatAlg returns on failure reach createKeyAccess. There + // it selects the RSA branch by default, and ocrypto.FromPublicPEM + // sniffs the PEM instead of honoring that choice: an EC or ML-KEM + // key parses successfully and wraps, but the KAO is left claiming + // keyType "wrapped" with no ephemeral public key. That produces a + // TDF nothing can decrypt, which is far worse to debug than a + // failure at creation time. + // + // SplitResult.Validate enforces the same rule for every splitter, + // including injected ones. Keeping the check here too lets the + // error name ErrSplitterUnsupportedAlgorithm and the offending + // policy.Algorithm, which the generic check cannot see. + alg, err := PolicyAlgorithmToKeyType(defaultKAS.GetPublicKey().GetAlgorithm()) + if err != nil { + return nil, fmt.Errorf("%w: kas %s: %w", ErrSplitterUnsupportedAlgorithm, url, err) + } + + share := make([]byte, len(dek)) + copy(share, dek) + return &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{ + url: { + Algorithm: alg, + KID: defaultKAS.GetPublicKey().GetKid(), + PEM: defaultKAS.GetPublicKey().GetPem(), + URL: url, + }, + }, + Splits: []Split{{ + Data: share, + KASURLs: []string{url}, + }}, + }, nil +} diff --git a/sdk/key_splitter_test.go b/sdk/key_splitter_test.go new file mode 100644 index 0000000000..c05ea576c1 --- /dev/null +++ b/sdk/key_splitter_test.go @@ -0,0 +1,435 @@ +package sdk + +import ( + "context" + "fmt" + "slices" + "testing" + + "github.com/opentdf/platform/lib/ocrypto" + "github.com/opentdf/platform/protocol/go/policy" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// A PEM is required to get past the earlier nil/empty checks; these +// tests never wrap anything, so its contents only need to parse as a +// PEM block, not match the advertised algorithm. +const splitterTestPEM = `-----BEGIN PUBLIC KEY----- +MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEjhRuJUUiBTLBmYuIJ6vGz1L8k+d3 +0j9RGVOM3G8mUJDPuOwLZLwJqDGmvHkyTa8k3lWK8v5nOSGN3nOJ8t2gEg== +-----END PUBLIC KEY-----` + +func TestSingleKASSplitterRequiresDefaultKAS(t *testing.T) { + for _, tc := range []struct { + name string + kas *policy.SimpleKasKey + }{ + {"nil KAS", nil}, + {"nil public key", &policy.SimpleKasKey{KasUri: "https://kas.example.com"}}, + {"empty PEM", &policy.SimpleKasKey{ + KasUri: "https://kas.example.com", + PublicKey: &policy.SimpleKasPublicKey{Algorithm: policy.Algorithm_ALGORITHM_RSA_2048, Kid: "k1"}, + }}, + {"empty KAS URI", &policy.SimpleKasKey{ + PublicKey: &policy.SimpleKasPublicKey{Algorithm: policy.Algorithm_ALGORITHM_RSA_2048, Kid: "k1", Pem: splitterTestPEM}, + }}, + } { + t.Run(tc.name, func(t *testing.T) { + splitter := DefaultKeySplitter() + res, err := splitter.Split(context.Background(), nil, []byte("0123456789abcdef"), tc.kas) + + require.ErrorIs(t, err, ErrSplitterRequiresDefaultKAS) + assert.Nil(t, res) + }) + } +} + +func TestSingleKASSplitterRejectsUnmappableAlgorithm(t *testing.T) { + // An algorithm the SDK has no wrapping scheme for used to yield the + // empty string, which createKeyAccess reads as a request for RSA. + // The resulting KAO claims keyType "wrapped" with no ephemeral + // public key, so the TDF is built successfully and then cannot be + // decrypted by anything. Fail at creation time instead. + for _, tc := range []struct { + name string + alg policy.Algorithm + }{ + {"unspecified", policy.Algorithm_ALGORITHM_UNSPECIFIED}, + {"out of range", policy.Algorithm(9999)}, + } { + t.Run(tc.name, func(t *testing.T) { + splitter := DefaultKeySplitter() + res, err := splitter.Split(context.Background(), nil, []byte("0123456789abcdef"), + &policy.SimpleKasKey{ + KasUri: "https://kas.example.com", + PublicKey: &policy.SimpleKasPublicKey{ + Algorithm: tc.alg, + Kid: "k1", + Pem: splitterTestPEM, + }, + }) + + require.ErrorIs(t, err, ErrSplitterUnsupportedAlgorithm) + assert.Nil(t, res) + // The KAS URL is in the message so an operator can tell which + // of several KASes is misconfigured. + assert.Contains(t, err.Error(), "https://kas.example.com") + }) + } +} + +func TestSingleKASSplitterAcceptsKnownAlgorithms(t *testing.T) { + for _, tc := range []struct { + name string + alg policy.Algorithm + want ocrypto.KeyType + }{ + {"rsa 2048", policy.Algorithm_ALGORITHM_RSA_2048, ocrypto.RSA2048Key}, + {"ec p256", policy.Algorithm_ALGORITHM_EC_P256, ocrypto.EC256Key}, + } { + t.Run(tc.name, func(t *testing.T) { + splitter := DefaultKeySplitter() + dek := []byte("0123456789abcdef") + res, err := splitter.Split(context.Background(), nil, dek, + &policy.SimpleKasKey{ + KasUri: "https://kas.example.com", + PublicKey: &policy.SimpleKasPublicKey{ + Algorithm: tc.alg, + Kid: "k1", + Pem: splitterTestPEM, + }, + }) + + require.NoError(t, err) + require.Len(t, res.Splits, 1) + assert.Equal(t, dek, res.Splits[0].Data) + assert.Equal(t, tc.want, res.KASPublicKeys["https://kas.example.com"].Algorithm) + }) + } +} + +// TestSplitResultValidate covers the contract every KeySplitter must +// meet, including the ones it is injected for. Each case below is a +// result that would be accepted at creation and fail only at decryption +// time, with an error naming nothing useful -- the reader derives the +// set of shares it needs from the key access objects present in the +// manifest, so a share that never reached the manifest is invisible to +// its completeness check and surfaces as a root signature failure. +func TestSplitResultValidate(t *testing.T) { + const ( + urlA = "https://kas-a.example.com" + urlB = "https://kas-b.example.com" + ) + keyA := KASPublicKey{Algorithm: ocrypto.RSA2048Key, KID: "a", PEM: splitterTestPEM, URL: urlA} + keyB := KASPublicKey{Algorithm: ocrypto.RSA2048Key, KID: "b", PEM: splitterTestPEM, URL: urlB} + share := func(b byte) []byte { return []byte{b, b, b, b} } + + for _, tc := range []struct { + name string + in *SplitResult + wantSub string + }{ + { + name: "single split", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: keyA}, + Splits: []Split{{Data: share(1), KASURLs: []string{urlA}}}, + }, + }, + { + name: "two splits, one KAS each", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: keyA, urlB: keyB}, + Splits: []Split{ + {ID: "s0", Data: share(1), KASURLs: []string{urlA}}, + {ID: "s1", Data: share(2), KASURLs: []string{urlB}}, + }, + }, + }, + { + name: "one split, two KASes (OR group)", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: keyA, urlB: keyB}, + Splits: []Split{{Data: share(1), KASURLs: []string{urlA, urlB}}}, + }, + }, + { + name: "nil result", + in: nil, + wantSub: "no splits", + }, + { + name: "no splits", + in: &SplitResult{KASPublicKeys: map[string]KASPublicKey{urlA: keyA}}, + wantSub: "no splits", + }, + { + name: "empty share", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: keyA}, + Splits: []Split{{Data: nil, KASURLs: []string{urlA}}}, + }, + wantSub: "no share data", + }, + { + // The reader XORs only as many bytes as each share carries, + // so a short one leaves the tail of the DEK unmasked. + name: "shares of different lengths", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: keyA, urlB: keyB}, + Splits: []Split{ + {ID: "s0", Data: share(1), KASURLs: []string{urlA}}, + {ID: "s1", Data: []byte{2, 2}, KASURLs: []string{urlB}}, + }, + }, + wantSub: "all shares must agree", + }, + { + // sid is what the reader groups on; an empty one on a + // multi-split result collapses the two into one group. + name: "empty id with more than one split", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: keyA, urlB: keyB}, + Splits: []Split{ + {Data: share(1), KASURLs: []string{urlA}}, + {ID: "s1", Data: share(2), KASURLs: []string{urlB}}, + }, + }, + wantSub: "empty id", + }, + { + name: "duplicate ids", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: keyA, urlB: keyB}, + Splits: []Split{ + {ID: "same", Data: share(1), KASURLs: []string{urlA}}, + {ID: "same", Data: share(2), KASURLs: []string{urlB}}, + }, + }, + wantSub: "more than one split", + }, + { + name: "split names no KAS", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: keyA}, + Splits: []Split{ + {ID: "s0", Data: share(1), KASURLs: []string{urlA}}, + {ID: "s1", Data: share(2)}, + }, + }, + wantSub: "names no KAS", + }, + { + name: "split names an unresolved KAS", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: keyA}, + Splits: []Split{{Data: share(1), KASURLs: []string{urlA, urlB}}}, + }, + wantSub: urlB, + }, + { + name: "resolved KAS has no PEM", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: {Algorithm: ocrypto.RSA2048Key, URL: urlA}}, + Splits: []Split{{Data: share(1), KASURLs: []string{urlA}}}, + }, + wantSub: "kas public key is missing", + }, + { + // KeyAccess.KasURL is built from the URL field, so a blank + // one points the reader at no endpoint at all. + name: "resolved KAS has no URL field", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: {Algorithm: ocrypto.RSA2048Key, PEM: splitterTestPEM}}, + Splits: []Split{{Data: share(1), KASURLs: []string{urlA}}}, + }, + wantSub: "empty URL field", + }, + { + name: "URL field disagrees with the map key", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: {Algorithm: ocrypto.RSA2048Key, PEM: splitterTestPEM, URL: urlB}}, + Splits: []Split{{Data: share(1), KASURLs: []string{urlA}}}, + }, + wantSub: "key access objects are built from the URL field", + }, + { + // An unparseable algorithm falls through createKeyAccess's + // RSA default, which sniffs the PEM instead of honoring the + // declared scheme and emits a KAO nothing can decrypt. + name: "unparseable algorithm", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: {Algorithm: "rsa:9999", PEM: splitterTestPEM, URL: urlA}}, + Splits: []Split{{Data: share(1), KASURLs: []string{urlA}}}, + }, + wantSub: "rsa:9999", + }, + { + // The zero value of the field, so this is what an injected + // splitter that simply never set Algorithm produces. + name: "empty algorithm", + in: &SplitResult{ + KASPublicKeys: map[string]KASPublicKey{urlA: {PEM: splitterTestPEM, URL: urlA}}, + Splits: []Split{{Data: share(1), KASURLs: []string{urlA}}}, + }, + wantSub: "unrecognized key type", + }, + } { + t.Run(tc.name, func(t *testing.T) { + err := tc.in.Validate() + if tc.wantSub == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantSub) + }) + } +} + +// TestSplitResultVerifyReconstruction covers the half of the contract +// Validate cannot reach. Every case below passes Validate -- the shares +// are non-empty, agree on length with each other, and carry distinct +// ids -- and still yields a key no reader can rebuild, which surfaces +// only as a root signature failure at decryption. +func TestSplitResultVerifyReconstruction(t *testing.T) { + dek := []byte("0123456789abcdef") + + // xorShares splits dek into n shares that close the XOR exactly, + // the shape a correct multi-KAS splitter produces. + xorShares := func(n int) []Split { + out := make([]Split, n) + accum := make([]byte, len(dek)) + for i := range n - 1 { + share := make([]byte, len(dek)) + for j := range share { + share[j] = byte(i*31 + j*7 + 1) + accum[j] ^= share[j] + } + out[i] = Split{ID: fmt.Sprintf("s%d", i), Data: share} + } + last := make([]byte, len(dek)) + for j := range last { + last[j] = dek[j] ^ accum[j] + } + out[n-1] = Split{ID: fmt.Sprintf("s%d", n-1), Data: last} + return out + } + + // mutate flips one byte of the first share, so the shares still + // agree on length and only the value is wrong. + mutate := func(splits []Split) []Split { + splits[0].Data[0] ^= 0xFF + return splits + } + + for _, tc := range []struct { + name string + in *SplitResult + dek []byte + wantSub string + }{ + { + name: "single split equal to the DEK", + in: &SplitResult{Splits: []Split{{Data: slices.Clone(dek)}}}, + dek: dek, + }, + { + name: "real multi-way XOR", + in: &SplitResult{Splits: xorShares(3)}, + dek: dek, + }, + { + name: "single split with the wrong value", + in: &SplitResult{Splits: mutate([]Split{{Data: slices.Clone(dek)}})}, + dek: dek, + wantSub: "do not XOR back to the DEK", + }, + { + name: "multi-way XOR that does not close", + in: &SplitResult{Splits: mutate(xorShares(3))}, + dek: dek, + wantSub: "do not XOR back to the DEK", + }, + { + // Mutually-agreeing short shares pass Validate, which only + // compares shares against each other, and leave the DEK's + // tail unmasked. + name: "short shares", + in: &SplitResult{Splits: []Split{ + {ID: "s0", Data: slices.Clone(dek[:8])}, + {ID: "s1", Data: make([]byte, 8)}, + }}, + dek: dek, + wantSub: "the length of the DEK", + }, + { + name: "nil receiver", + in: nil, + dek: dek, + wantSub: "no splits", + }, + { + name: "no splits", + in: &SplitResult{}, + dek: dek, + wantSub: "no splits", + }, + { + // Refused rather than passed: an empty DEK would make every + // share trivially "reconstruct" nothing. + name: "empty DEK", + in: &SplitResult{Splits: []Split{{Data: slices.Clone(dek)}}}, + dek: nil, + wantSub: "empty DEK", + }, + } { + t.Run(tc.name, func(t *testing.T) { + if tc.in != nil && len(tc.in.Splits) > 0 && len(tc.dek) > 0 { + // Every case must be one Validate would let through, or + // it proves nothing about the gap this method closes. + populateKASKeys(tc.in) + require.NoError(t, tc.in.Validate(), "case must pass Validate to be interesting") + } + + err := tc.in.VerifyReconstruction(tc.dek) + if tc.wantSub == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantSub) + }) + } +} + +// populateKASKeys gives every split a resolvable KAS so Validate's +// structural checks pass and only the reconstruction check is left to +// fail. +func populateKASKeys(r *SplitResult) { + const url = "https://kas-a.example.com" + r.KASPublicKeys = map[string]KASPublicKey{ + url: {Algorithm: ocrypto.RSA2048Key, KID: "a", PEM: splitterTestPEM, URL: url}, + } + for i := range r.Splits { + r.Splits[i].KASURLs = []string{url} + } +} + +// TestDefaultKeySplitterResultValidates closes the loop: whatever the +// shipped splitter produces must satisfy the contract the writer +// enforces on injected ones, or the default path fails at Finalize. +func TestDefaultKeySplitterResultValidates(t *testing.T) { + res, err := DefaultKeySplitter().Split(context.Background(), nil, []byte("0123456789abcdef"), + &policy.SimpleKasKey{ + KasUri: "https://kas.example.com", + PublicKey: &policy.SimpleKasPublicKey{ + Algorithm: policy.Algorithm_ALGORITHM_RSA_2048, + Kid: "k1", + Pem: splitterTestPEM, + }, + }) + require.NoError(t, err) + require.NoError(t, res.Validate()) +} From bd53671ddf4ca790f46e5624efa55041133c9a2a Mon Sep 17 00:00:00 2001 From: Dave Mihalcik Date: Wed, 16 Sep 2026 15:14:43 -0400 Subject: [PATCH 2/2] Update sdk/chunked_writer.go Co-authored-by: sujankota --- sdk/chunked_writer.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/sdk/chunked_writer.go b/sdk/chunked_writer.go index 2225c868c9..a74474ba54 100644 --- a/sdk/chunked_writer.go +++ b/sdk/chunked_writer.go @@ -792,6 +792,10 @@ func (w *chunkedWriter) WriteSegment(ctx context.Context, index int, data []byte if err != nil { return nil, fmt.Errorf("encrypt segment %d: %w", index, err) } + // SegmentGMAC reads the trailing AEAD tag, so hashing ciphertext alone is + // equivalent to hashing nonce||ciphertext -- which is why there is no + // concatenation here. An algorithm that MACs the whole segment (HS256) + // would need the nonce prepended back. sig, err := segmentIntegrity(ciphertext, w.dek, SegmentGMAC, w.useHex) if err != nil { return nil, fmt.Errorf("segment %d signature: %w", index, err)