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..a74474ba54 --- /dev/null +++ b/sdk/chunked_writer.go @@ -0,0 +1,1197 @@ +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) + } + // 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) + } + 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()) +}