diff --git a/sdk/chunked_options.go b/sdk/chunked_options.go new file mode 100644 index 0000000000..e80abbfa43 --- /dev/null +++ b/sdk/chunked_options.go @@ -0,0 +1,269 @@ +package sdk + +import ( + "errors" + "fmt" + "io" + + "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. + +// withChunkedArchiveWriterFactory overrides the ZIP archive writer +// factory used by the chunked Writer. 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 chunked Writer. 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 chunked +// Writer 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 = 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 + } +} + +// WithChunkedIntegrityAlgorithm sets the algorithm used for the +// manifest root signature. algo must be HS256 or GMAC. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedIntegrityAlgorithm(algo IntegrityAlgorithm) ChunkedWriterOption { + return func(c *ChunkedWriterConfig) error { + if algo != HS256 && algo != GMAC { + return fmt.Errorf("chunked: unsupported integrity algorithm %v", algo) + } + c.integrityAlgorithm = algo + return nil + } +} + +// WithChunkedKeySplitter overrides the key splitter used by the +// chunked Writer. 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 + } +} + +// WithChunkedSegmentIntegrityAlgorithm sets the algorithm used for +// per-segment integrity hashes. algo must be HS256 or GMAC. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedSegmentIntegrityAlgorithm(algo IntegrityAlgorithm) ChunkedWriterOption { + return func(c *ChunkedWriterConfig) error { + if algo != HS256 && algo != GMAC { + return fmt.Errorf("chunked: unsupported integrity algorithm %v", algo) + } + c.segmentIntegrityAlgorithm = algo + 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 = assertions + return nil + } +} + +// WithChunkedAttributes overrides the writer's initial attributes for +// this Finalize call. +// +// 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 = values + return nil + } +} + +// WithChunkedDefaultKASForFinalize overrides the writer's initial +// default KAS for this Finalize 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 + } +} + +// WithChunkedExcludeVersion omits the schemaVersion field from the +// produced manifest for compatibility with older readers. +// +// A reader treats a missing schemaVersion as "predates 4.3.0" and so +// expects hex-then-base64 signatures. Those are written during +// WriteSegment, before this option is seen, so on its own this option +// makes Finalize fail with [ErrChunkedVersionHexMismatch]. Pass +// [WithChunkedTargetMode] at construction instead; it sets both. +// +// Experimental: not part of the stable SDK API; may change or be removed. +func WithChunkedExcludeVersion() ChunkedFinalizeOption { + return func(c *ChunkedFinalizeConfig) error { + c.excludeVersion = true + return nil + } +} + +// WithChunkedTargetMode targets a specific TDF spec version, given as +// a semver string such as "4.2.2". +// +// Below 4.3.0 the writer emits the legacy wire format: segment, root, +// and assertion signatures are hex-encoded before base64 (yielding the +// doubly-encoded values pre-4.3.0 readers expect), and schemaVersion is +// omitted from the manifest, which is how those readers detect it. The +// two travel together -- a manifest carrying one without the other +// cannot be verified by any reader. +// +// An empty mode selects the current format. +// +// 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 { + c.keepSegments = indices + return nil + } +} diff --git a/sdk/chunked_test.go b/sdk/chunked_test.go new file mode 100644 index 0000000000..1b1490594e --- /dev/null +++ b/sdk/chunked_test.go @@ -0,0 +1,1243 @@ +package sdk + +import ( + "archive/zip" + "bytes" + "context" + "crypto/ecdsa" + "crypto/x509" + "encoding/pem" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "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) + }) + } +} + +// TestChunkedSegmentsNotStartingAtZero covers a writer whose lowest +// written index is not 0 -- 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. +// +// Either answer is acceptable: Finalize may refuse the write set, or it +// may produce a TDF that reads back. What it must not do is return +// success alongside bytes that are not a readable archive, because by +// then the upload has happened and the plaintext is gone. Today only +// segment 0 emits the payload's ZIP local file header (see +// zipstream.segmentWriter.WriteSegment), so the third case is what +// happens. +func TestChunkedSegmentsNotStartingAtZero(t *testing.T) { + ctx := context.Background() + s := newChunkedTestSDK(t) + w, kasBundle := newChunkedWriterForTest(ctx, t) + + chunks := map[int][]byte{5: []byte("hello-"), 6: []byte("world!")} + indices := []int{5, 6} + + encrypted := make(map[int][]byte, len(indices)) + for _, idx := range indices { + 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 + } + + fin, err := w.Finalize(ctx) + if err != nil { + // Refusing the write set is a valid outcome; nothing was + // published, so there is nothing further to check. + t.Logf("Finalize rejected a segment set starting at %d: %v", indices[0], err) + return + } + + // Finalize claimed success, so the bytes it told the caller to + // assemble have to be a TDF. + var body bytes.Buffer + for _, idx := range indices { + body.Write(encrypted[idx]) + } + body.Write(fin.Data) + + reader, err := s.LoadTDF(bytes.NewReader(body.Bytes()), + WithKasAllowlist([]string{kasBundle.url}), + ) + + // The specific defect: with no local file header the reader runs off + // the end of the buffer parsing the ZIP structure, so the container + // never opens. Anything else is some other test's business. + require.NotErrorIs(t, err, io.ErrUnexpectedEOF, + "Finalize succeeded but the assembled bytes are not a ZIP container") + if err != nil { + t.Logf("LoadTDF failed for an unrelated reason, not this test's subject: %v", err) + return + } + + plain, err := io.ReadAll(reader) + require.NoError(t, err, "Finalize succeeded, so the payload must decrypt") + assert.Equal(t, []byte("hello-world!"), plain) +} + +// 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. +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. +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) + } + 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 +} + +// 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: string(ocrypto.EC256Key), + KID: "ec-kid", + PEM: pubPEM, + URL: kasURL, + }, + }, + Splits: []Split{{Data: dek, KASURLs: []string{kasURL}}}, + } + + kaos, err := buildChunkedKeyAccessObjects(splits, []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) +} + +// 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 HS256 signature is base64(hex(hmac)): 32 HMAC bytes + // rendered as 64 hex characters. The 4.3.0 form is base64(hmac), + // which decodes to 32 bytes. + 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, 64, "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) +} + +// TestChunkedExcludeVersionRequiresLegacyMode verifies that omitting +// schemaVersion without the matching signature encoding is refused +// rather than silently producing an unverifiable TDF. +func TestChunkedExcludeVersionRequiresLegacyMode(t *testing.T) { + ctx := context.Background() + writer, _ := newChunkedWriterForTest(ctx, t) + + writeChunkedSegments(ctx, t, writer, [][]byte{[]byte("mismatch")}) + + _, err := writer.Finalize(ctx, WithChunkedExcludeVersion()) + require.ErrorIs(t, err, ErrChunkedVersionHexMismatch) +} + +// 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) + }) + } +} + +// TestChunkedIntegrityAlgorithmRejectsUnsupported verifies both +// integrity-algorithm options reject a value outside {HS256, GMAC} +// rather than letting it reach calculateSignature, which treats any +// unrecognized value as GMAC. +func TestChunkedIntegrityAlgorithmRejectsUnsupported(t *testing.T) { + ctx := context.Background() + kasBundle := newChunkedFakeKAS(t) + defer kasBundle.server.Close() + + const bogus IntegrityAlgorithm = 99 + + for _, tc := range []struct { + name string + opt ChunkedWriterOption + }{ + {"root", WithChunkedIntegrityAlgorithm(bogus)}, + {"segment", WithChunkedSegmentIntegrityAlgorithm(bogus)}, + } { + 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(), "unsupported integrity algorithm") + assert.Nil(t, writer) + }) + } +} + +// errArchiveWriteFailed is the injected archive failure used to drive +// WriteSegment's error paths. +var errArchiveWriteFailed = errors.New("archive write 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. WriteSegment used to publish the segment metadata before +// handing the bytes to the archive, so Finalize emitted a manifest +// covering a payload the archive had rejected and the caller never +// received: the reader then mapped every later segment at the wrong +// payload offset. +// +// 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. WriteSegment used to reserve the index up front +// and never release it, so a single transient failure made the index +// permanently unwritable and left the writer unable to 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 +} + +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) + 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. +func TestChunkedWriteSegmentCleansUpArchiveOnFailure(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 + })) + + _, err := writer.WriteSegment(ctx, 1, []byte("doomed")) + require.ErrorIs(t, err, errArchiveWriteFailed) + assert.Equal(t, []int{1}, archive.cleanedUp) + + _, err = writer.WriteSegment(ctx, 1, []byte("retry")) + require.NoError(t, err, "the archive's own record of the failed attempt must be rolled back, not just the sdk-level reservation") +} + +// 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") +} diff --git a/sdk/chunked_writer.go b/sdk/chunked_writer.go new file mode 100644 index 0000000000..93f264edcd --- /dev/null +++ b/sdk/chunked_writer.go @@ -0,0 +1,872 @@ +package sdk + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/json" + "errors" + "fmt" + "hash/crc32" + "io" + "slices" + "sort" + "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 usable from +// outside. archiveWriterFactory in particular returns a +// zipstream.SegmentWriter, which lives under internal/, so no external package +// could implement it even if the type were exported. + +// 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 } + +// defaultRand is the production entropy source used by the chunked writer when +// no other io.Reader is injected. +var defaultRand io.Reader = rand.Reader + +// segmentCipher encrypts a single payload segment. Implementations must be +// safe for concurrent use by segment writers. +type segmentCipher interface { + // EncryptInPlace returns (ciphertext, nonce, error). + 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 sized for +// a single starting segment (it grows as more segments arrive), with its clock +// plumbed to the caller-supplied clock. +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") + + // 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. + ErrChunkedCloseFailed = errors.New("chunked: archive close failed after finalize; 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") + + // ErrChunkedVersionHexMismatch is returned when Finalize is asked + // to omit schemaVersion on a writer that was not constructed in + // legacy signature mode. Use WithChunkedTargetMode to set both. + ErrChunkedVersionHexMismatch = errors.New("chunked: excluding schemaVersion requires a pre-4.3.0 target mode; use WithChunkedTargetMode") +) + +// 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, or + // ErrChunkedCloseFailed if the archive's Close fails after its + // Finalize already succeeded -- the writer is unusable at that + // point and every subsequent call returns the same error. + Finalize(ctx context.Context, opts ...ChunkedFinalizeOption) (*ChunkedFinalizeResult, error) + + // GetManifest returns the manifest for the TDF. Before Finalize + // this is a snapshot built from currently-written segments; after + // Finalize it is the manifest that was written. + 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. + 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. + 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. + 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 + // emitted segments. + 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 emitted + // segments. + TotalSize int64 +} + +// ChunkedWriterConfig captures the settings supplied at +// NewChunkedWriter time. Fields are unexported; use options. +// +// Experimental: not part of the stable SDK API; may change or be removed. +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 + + // integrityAlgorithm is the algorithm used for the root + // signature. Defaults to HS256. + integrityAlgorithm IntegrityAlgorithm + + // rand is the entropy source used to generate the DEK. Defaults + // to crypto/rand.Reader. + rand io.Reader + + // segmentIntegrityAlgorithm is the algorithm used for per-segment + // integrity hashes. Defaults to HS256. + segmentIntegrityAlgorithm IntegrityAlgorithm + + // 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. +// +// Experimental: not part of the stable SDK API; may change or be removed. +type ChunkedFinalizeConfig struct { + // assertions to sign and attach to the produced TDF. Each + // AssertionConfig must carry a SigningKey (or the writer's DEK + // will be used with HS256). + 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 "application/octet-stream". + 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 + +// 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 + + // closeFailed is true once archiveWriter.Close has failed after + // archiveWriter.Finalize already succeeded. The archive itself is + // terminally finalized at that point even though w.finalized was + // never set, so every method must refuse to proceed rather than + // retry against an archive that can only ever fail again. + closeFailed bool + + // 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 + + // integrityAlgorithm is used for the root signature. + integrityAlgorithm IntegrityAlgorithm + + // manifest holds the finalized manifest for post-Finalize + // GetManifest calls. + manifest *Manifest + + // mu guards writer state that spans WriteSegment and Finalize. + mu sync.RWMutex + + // segmentIntegrityAlgorithm is used for per-segment hashes. + segmentIntegrityAlgorithm IntegrityAlgorithm + + // segments records per-index Segment metadata (hash + sizes). + segments map[int]*Segment + + // 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. +// +// No SDK value is needed: everything the writer depends on — the key +// splitter, the archive and cipher factories, the entropy source — is +// supplied through options. +// +// 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{}, + integrityAlgorithm: HS256, + rand: defaultRand, + segmentIntegrityAlgorithm: HS256, + 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, + integrityAlgorithm: cfg.integrityAlgorithm, + segmentIntegrityAlgorithm: cfg.segmentIntegrityAlgorithm, + segments: make(map[int]*Segment), + 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.closeFailed { + return nil, ErrChunkedCloseFailed + } + if w.finalized { + return nil, ErrChunkedAlreadyFinalized + } + + // 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. Size stays negative until the archive + // accepts the write, so a reservation in flight does not count. + if seg, ok := w.segments[0]; !ok || seg.Size < 0 { + return nil, ErrChunkedMissingSegmentZero + } + + cfg, err := w.applyFinalizeOptions(opts) + if err != nil { + return nil, err + } + + manifest, totalPlaintext, totalEncrypted, 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) + } + finalBytes, err := w.archiveWriter.Finalize(ctx, manifestBytes) + if err != nil { + return nil, fmt.Errorf("finalize archive: %w", err) + } + 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. + // Mark the writer unusable rather than leaving it looking + // retryable. + w.closeFailed = true + return nil, fmt.Errorf("%w: %w", ErrChunkedCloseFailed, err) + } + + w.finalized = true + w.manifest = manifest + return &ChunkedFinalizeResult{ + Data: finalBytes, + EncryptedSize: totalEncrypted, + Manifest: manifest, + TotalSegments: len(manifest.Segments), + TotalSize: totalPlaintext, + }, 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.closeFailed { + return nil, ErrChunkedCloseFailed + } + if w.finalized && w.manifest != nil { + return cloneChunkedManifest(w.manifest), nil + } + cfg, err := w.applyFinalizeOptions(opts) + if err != nil { + return nil, err + } + 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.closeFailed { + w.mu.Unlock() + return nil, ErrChunkedCloseFailed + } + 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 Size negative: the segment does not count as + // written until its bytes are in the archive. + seg := &Segment{Size: -1} + w.segments[index] = seg + 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. + release := func() { + w.mu.Lock() + if cur, ok := w.segments[index]; ok && cur == seg && cur.Size < 0 { + delete(w.segments, index) + } + 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 + } + release() + 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's failure here is best-effort + // and does not change what this call returns. + _ = w.archiveWriter.CleanupSegment(index) + } + }() + + ciphertext, nonce, err := w.block.EncryptInPlace(data) + if err != nil { + return nil, fmt.Errorf("encrypt segment %d: %w", index, err) + } + sealed := make([]byte, 0, len(nonce)+len(ciphertext)) + sealed = append(sealed, nonce...) + sealed = append(sealed, ciphertext...) + sig, err := calculateSignature(sealed, w.dek, w.segmentIntegrityAlgorithm, w.useHex) + if err != nil { + return nil, fmt.Errorf("segment %d signature: %w", index, err) + } + hash := string(ocrypto.Base64Encode([]byte(sig))) + encryptedSize := int64(len(sealed)) + + crc := crc32.NewIEEE() + if _, err := crc.Write(nonce); err != nil { + return nil, err + } + if _, err := crc.Write(ciphertext); err != nil { + return nil, err + } + archiveWriteAttempted = true + header, err := w.archiveWriter.WriteSegment(ctx, index, uint64(encryptedSize), crc.Sum32()) + 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() + seg.EncryptedSize = encryptedSize + seg.Hash = hash + seg.Size = int64(len(data)) + w.mu.Unlock() + committed = true + + 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 +} + +// 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, + mimeType: "application/octet-stream", + } + for _, opt := range opts { + if err := opt(cfg); err != nil { + return nil, err + } + } + // Omitting schemaVersion is how a reader is told the TDF predates + // 4.3.0, and such a reader expects hex-then-base64 signatures. The + // segment signatures were already written by then, so the two + // settings cannot be reconciled here -- refuse rather than emit a + // TDF that no reader can verify. + if cfg.excludeVersion && !w.useHex { + return nil, ErrChunkedVersionHexMismatch + } + if len(cfg.attributes) == 0 && len(w.initialAttributes) > 0 { + cfg.attributes = w.initialAttributes + } + if cfg.defaultKAS == nil && w.initialDefaultKAS != nil { + cfg.defaultKAS = w.initialDefaultKAS + } + return cfg, nil +} + +// buildManifest composes the manifest from writer state, splits the +// DEK, wraps splits into KAOs, and computes the root signature. +func (w *chunkedWriter) buildManifest(ctx context.Context, cfg *ChunkedFinalizeConfig) (*Manifest, int64, int64, error) { + order, err := w.segmentOrderLocked(cfg.keepSegments) + if err != nil { + return nil, 0, 0, err + } + + splits, err := w.splitter.Split(ctx, cfg.attributes, w.dek, cfg.defaultKAS) + if err != nil { + return nil, 0, 0, err + } + policyBytes, err := buildChunkedPolicy(cfg.attributes) + if err != nil { + return nil, 0, 0, err + } + kaos, err := buildChunkedKeyAccessObjects(splits, policyBytes, cfg.encryptedMetadata) + if err != nil { + return nil, 0, 0, err + } + + encInfo := EncryptionInformation{ + KeyAccessObjs: kaos, + KeyAccessType: kSplitKeyType, + Policy: string(ocrypto.Base64Encode(policyBytes)), + Method: Method{ + Algorithm: kGCMCipherAlgorithm, + IsStreamable: true, + }, + IntegrityInformation: IntegrityInformation{ + SegmentHashAlgorithm: integrityAlgorithmString(w.segmentIntegrityAlgorithm), + Segments: make([]Segment, len(order)), + }, + } + + var aggregate bytes.Buffer + var totalPlaintext, totalEncrypted int64 + for i, idx := range order { + seg, ok := w.segments[idx] + if !ok || seg.Size < 0 { + return nil, 0, 0, fmt.Errorf("segment %d not written; cannot finalize", idx) + } + if seg.Hash == "" { + return nil, 0, 0, fmt.Errorf("segment %d has empty hash", idx) + } + encInfo.Segments[i] = *seg + totalPlaintext += seg.Size + totalEncrypted += seg.EncryptedSize + decoded, err := ocrypto.Base64Decode([]byte(seg.Hash)) + if err != nil { + return nil, 0, 0, fmt.Errorf("decode segment %d hash: %w", idx, err) + } + aggregate.Write(decoded) + } + if len(order) > 0 { + if first, ok := w.segments[order[0]]; ok { + encInfo.DefaultEncryptedSegSize = first.EncryptedSize + encInfo.DefaultSegmentSize = first.Size + } + } + + rootSig, err := calculateSignature(aggregate.Bytes(), w.dek, w.integrityAlgorithm, w.useHex) + if err != nil { + return nil, 0, 0, err + } + encInfo.RootSignature = RootSignature{ + Algorithm: integrityAlgorithmString(w.integrityAlgorithm), + 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, 0, 0, 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, totalPlaintext, totalEncrypted, 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. 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 := range w.segments { + written = append(written, idx) + } + sort.Ints(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 + } + if _, ok := w.segments[idx]; !ok { + 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, policyBytes []byte, metadata string) ([]KeyAccess, error) { + if splits == nil || len(splits.Splits) == 0 { + return nil, errors.New("no splits produced") + } + base64Policy := ocrypto.Base64Encode(policyBytes) + + var out []KeyAccess + 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 { + pk, ok := splits.KASPublicKeys[url] + if !ok { + continue + } + if pk.PEM == "" { + return nil, fmt.Errorf("splitID:[%s], kas:[%s]: %w", split.ID, url, errKasPubKeyMissing) + } + 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) + } + } + if len(out) == 0 { + return nil, errors.New("no valid key access objects generated") + } + 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 _, v := range values { + p.Body.DataAttributes = append(p.Body.DataAttributes, attributeObject{ + Attribute: v.GetFqn(), + }) + } + return json.Marshal(p) +} + +// cloneChunkedManifest returns a shallow-deep copy safe to hand out. +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/key_splitter.go b/sdk/key_splitter.go new file mode 100644 index 0000000000..a40fded1f3 --- /dev/null +++ b/sdk/key_splitter.go @@ -0,0 +1,169 @@ +package sdk + +import ( + "context" + "errors" + "fmt" + + "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 the split share (XOR of the DEK with the other shares). + Data []byte + + // ID uniquely identifies this split within a SplitResult. Empty + // when the result contains only one split (single-KAO TDF). + 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 +} + +// 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 as an exact + // ocrypto.KeyType string, e.g. "rsa:2048" or "ec:secp256r1" -- use + // PolicyAlgorithmToKeyType to derive it from a policy.Algorithm. + // A bare or unrecognized value falls through createKeyAccess's RSA + // default and can produce a KAO that cannot be decrypted. + Algorithm string + + // 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 +} + +// 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: 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 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. + alg := algorithmPolicyToString(defaultKAS.GetPublicKey().GetAlgorithm()) + if alg == "" { + return nil, fmt.Errorf("%w: kas %s advertises algorithm %v", + ErrSplitterUnsupportedAlgorithm, url, defaultKAS.GetPublicKey().GetAlgorithm()) + } + + 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 +} + +// algorithmPolicyToString maps a policy.Algorithm enum to the +// ocrypto.KeyType string form used when picking a wrap scheme. +// Unknown enums, including ALGORITHM_UNSPECIFIED, return the empty +// string; callers must reject that rather than pass it on, since +// createKeyAccess reads it as a request for RSA. singleKASSplitter.Split +// has that guard. +func algorithmPolicyToString(a policy.Algorithm) string { + if kt, err := PolicyAlgorithmToKeyType(a); err == nil { + return string(kt) + } + return "" +} diff --git a/sdk/key_splitter_test.go b/sdk/key_splitter_test.go new file mode 100644 index 0000000000..6fc55f1774 --- /dev/null +++ b/sdk/key_splitter_test.go @@ -0,0 +1,107 @@ +package sdk + +import ( + "context" + "testing" + + "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 string + }{ + {"rsa 2048", policy.Algorithm_ALGORITHM_RSA_2048, "rsa:2048"}, + {"ec p256", policy.Algorithm_ALGORITHM_EC_P256, "ec:secp256r1"}, + } { + 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) + }) + } +}