feat(sdk): add ChunkedWriter for out-of-order segment TDF creation#3782
feat(sdk): add ChunkedWriter for out-of-order segment TDF creation#3782ronelliott wants to merge 1 commit into
Conversation
DSPX-2604 Introduces ChunkedWriter, an SDK entry point that produces a TDF from segments supplied one at a time. Callers encrypt and upload each segment independently — off-thread or in parallel — then call Finalize to seal the archive. Contrast with SDK.CreateTDF, which requires the full plaintext buffered up front. The writer exposes injection seams (Clock, SegmentCipher, ArchiveWriterFactory, KeySplitter) so tests can substitute deterministic or in-memory implementations without touching global state. The default splitter is single-KAS only; callers with attribute-based multi-KAS grants inject their own via WithChunkedKeySplitter. No reader is included — produced TDFs round-trip through the mainline SDK.LoadTDF path, verified end-to-end by TestChunkedRoundTrip. Files: chunked_writer.go — ChunkedWriter interface + concrete writer chunked_options.go — WithChunked* options clock.go — Clock interface + defaults segment.go — SegmentCipher interface + AES-256-GCM default archive_writer.go — ArchiveWriterFactory + ZIP64 default key_splitter.go — KeySplitter interface + single-KAS default chunked_test.go — round-trip, keep-segments, assertion-rejection Signed-off-by: Ron Elliott <ron.elliott@virtru.com>
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
github-actions
Bot
added
comp:sdk
📝 WalkthroughWalkthroughAdds a configurable chunked TDF writer that encrypts indexed segments, streams ZIP output, builds manifests with KAS-wrapped keys, supports finalize-time controls, and includes end-to-end tests using a fake KAS. ChangesChunked TDF writing
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested labels: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@sdk/chunked_test.go`:
- Around line 116-124: Extend the chunked-segment tests around
writeChunkedSegments to add a round-trip case that calls WriteSegment with
indices in out-of-order sequence such as 2, 0, 1, appends each returned TDFData
in that write order, and verifies LoadTDF reconstructs plaintext ordered by
segment index. Keep the existing in-order helper behavior unchanged.
In `@sdk/chunked_writer.go`:
- Around line 438-455: Refactor buildManifest and its Finalize/GetManifest
callers so the long-running splitter.Split and related KAO construction do not
execute while w.mu is held. Snapshot segment metadata, dek, and attributes under
the appropriate lock, release the lock for splitting and wrapping, then
reacquire it only to commit w.manifest and w.finalized while preserving
consistency and existing finalize behavior.
- Around line 747-762: Update chunkedWrapKeyWithPublicKey so empty or
unrecognized pk.Algorithm values return an explicit unsupported-algorithm error
instead of falling through to chunkedWrapKeyWithRSA. Preserve the existing KEM
and EC dispatch, and only invoke RSA wrapping when the algorithm is explicitly
recognized as an RSA type.
- Around line 41-65: Align ChunkedWriter with sdk/experimental/tdf.Writer
instead of maintaining a smaller duplicate API: either delegate to or deprecate
the experimental implementation, or document ChunkedWriter as its intended
replacement. Add assertion handling to ChunkedFinalize and GetManifest using the
existing ChunkedFinalizeConfig assertion fields, and preserve attribute-based
K/O configuration plus contiguous-prefix trimming through WithSegments rather
than returning ErrChunkedAssertionsUnsupported.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: f266eba3-88cf-43d2-994a-43ac5a66402a
📒 Files selected for processing (7)
sdk/archive_writer.gosdk/chunked_options.gosdk/chunked_test.gosdk/chunked_writer.gosdk/clock.gosdk/key_splitter.gosdk/segment.go
| 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) | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Cover out-of-order segment arrival.
This helper hardcodes index order to write order, so the suite never verifies the documented WriteSegment contract for independently supplied, out-of-order segments. Add a round-trip case that writes indices such as 2, 0, 1, appends each returned TDFData in write order, then verifies LoadTDF returns plaintext ordered by segment index. Based on provided upstream contract, indices need not arrive in order and the PR promises independently supplied segments.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@sdk/chunked_test.go` around lines 116 - 124, Extend the chunked-segment tests
around writeChunkedSegments to add a round-trip case that calls WriteSegment
with indices in out-of-order sequence such as 2, 0, 1, appends each returned
TDFData in that write order, and verifies LoadTDF reconstructs plaintext ordered
by segment index. Keep the existing in-order helper behavior unchanged.
| // 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. | ||
| 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 (central directory | ||
| // + end-of-central-directory record) that must be appended after | ||
| // every segment's TDFData. | ||
| 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 (local header + nonce + ciphertext). | ||
| // Callers upload or buffer those bytes; Finalize does not | ||
| // re-emit them. Indices need not arrive in order and need not be | ||
| // contiguous. | ||
| WriteSegment(ctx context.Context, index int, data []byte) (*ChunkedSegmentResult, error) | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoff
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether sdk/experimental/tdf exists in this checkout and compare its writer API surface.
fd -t d experimental sdk
fd -t f '.go$' sdk/experimental/tdf 2>/dev/null | xargs -r ast-grep run --pattern 'func $NAME($$$) $$$' --lang goRepository: opentdf/platform
Length of output: 50373
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== chunked writer relevant symbols =="
ast-grep outline sdk/chunked_writer.go --view compact 2>/dev/null || true
rg -n "type Chunked|func New|func .*Chunked|Finalize|WriteSegment|ErrChunkedAssertionsUnsupported|Segment|Manifest" sdk/chunked_writer.go
echo
echo "== experimental tdf relevant symbols (named/exported and documentation snippets) =="
rg -n "type .*[Ww]riter|func .*[Ww]riter|NewWriter|WriteSegment|Finalize|Segment|LoadTDF|assert|Assert" sdk/experimental/tdf -g '*.go' | head -300
echo
echo "== file sizes =="
wc -l sdk/chunked_writer.go sdk/experimental/tdf/*.go 2>/dev/null | head -60Repository: opentdf/platform
Length of output: 36816
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== chunked_final options and unsupported assertions path =="
sed -n '150,210p' sdk/chunked_writer.go
sed -n '280,305p' sdk/chunked_writer.go
sed -n '450,530p' sdk/chunked_writer.go
rg -n "ErrChunkedAssertionsUnsupported|FinalizeOption\\(.*Chunked|ChunkedAssertions|WithChunked" sdk/chunked_writer.go
echo
echo "== experimental writer New/Finalize/API and assertions path =="
sed -n '1,140p' sdk/experimental/tdf/writer.go
sed -n '230,350p' sdk/experimental/tdf/writer.go
sed -n '430,520p' sdk/experimental/tdf/options.go
sed -n '264,400p' sdk/experimental/tdf/writer_test.go
sed -n '400,550p' sdk/experimental/tdf/writer_test.goRepository: opentdf/platform
Length of output: 28135
Coordinate this chunked writer with the existing experimental TDF writer.
sdk/experimental/tdf.Writer already supports streamed/out-of-order segment writing, NewWriter/WriteSegment/Finalize, attribute-based K/O configuration, and contiguous-prefix segment trimming via WithSegments. This stable ChunkedWriter duplicates that API while explicitly unsupported by ErrChunkedAssertionsUnsupported despite ChunkedFinalizeConfig containing assertions fields. Either build from/deprecate the experimental writer or clearly document the intended replacement path and add the missing assertion support so the stable API is not a smaller fork of the same capability.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@sdk/chunked_writer.go` around lines 41 - 65, Align ChunkedWriter with
sdk/experimental/tdf.Writer instead of maintaining a smaller duplicate API:
either delegate to or deprecate the experimental implementation, or document
ChunkedWriter as its intended replacement. Add assertion handling to
ChunkedFinalize and GetManifest using the existing ChunkedFinalizeConfig
assertion fields, and preserve attribute-based K/O configuration plus
contiguous-prefix trimming through WithSegments rather than returning
ErrChunkedAssertionsUnsupported.
| 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 | ||
| } |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
Holding w.mu across KeySplitter.Split serializes concurrent WriteSegment calls during Finalize/GetManifest.
buildManifest is invoked while Finalize holds w.mu.Lock() (full exclusive lock for the whole function) or GetManifest holds w.mu.RLock(), and it calls w.splitter.Split(ctx, ...) (line 444) inside that critical section. A real KeySplitter implementation resolving attribute-based multi-KAS grants will likely make network calls to the policy/KAS services. Holding the writer's lock for that duration blocks every concurrent WriteSegment call — undermining the stated goal of writing segments "off-thread or in parallel" (see the ChunkedWriter doc comment).
Consider snapshotting the needed state (segment metadata, dek, attributes) under the lock, releasing it, performing Split/KAO wrapping unlocked, then re-acquiring briefly only to commit w.manifest/w.finalized.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@sdk/chunked_writer.go` around lines 438 - 455, Refactor buildManifest and its
Finalize/GetManifest callers so the long-running splitter.Split and related KAO
construction do not execute while w.mu is held. Snapshot segment metadata, dek,
and attributes under the appropriate lock, release the lock for splitting and
wrapping, then reacquire it only to commit w.manifest and w.finalized while
preserving consistency and existing finalize behavior.
| // chunkedWrapKeyWithPublicKey dispatches to the right KAS wrapping | ||
| // scheme based on the algorithm advertised by the splitter. | ||
| func chunkedWrapKeyWithPublicKey(symKey []byte, pk KASPublicKey) (string, string, string, error) { | ||
| if pk.PEM == "" { | ||
| return "", "", "", fmt.Errorf("public key PEM is empty for kas %s", pk.URL) | ||
| } | ||
| ktype := ocrypto.KeyType(pk.Algorithm) | ||
| switch { | ||
| case ocrypto.IsKEMKeyType(ktype): | ||
| return chunkedWrapKeyWithKEM(ktype, pk.PEM, symKey) | ||
| case ocrypto.IsECKeyType(ktype): | ||
| return chunkedWrapKeyWithEC(ktype, pk.PEM, symKey) | ||
| default: | ||
| return chunkedWrapKeyWithRSA(pk.PEM, symKey) | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Default branch silently attempts RSA wrap for unsupported/empty algorithm types.
key_splitter.go's algorithmPolicyToString comment states unknown algorithms return "" and "the caller treats that as 'unsupported algorithm'." But chunkedWrapKeyWithPublicKey's switch has no case for an empty/unrecognized ktype — it silently falls to chunkedWrapKeyWithRSA, which will attempt to parse a non-RSA PEM as an RSA key, producing a confusing low-level parse error instead of a clear "unsupported algorithm" error.
🐛 Proposed fix
func chunkedWrapKeyWithPublicKey(symKey []byte, pk KASPublicKey) (string, string, string, error) {
if pk.PEM == "" {
return "", "", "", fmt.Errorf("public key PEM is empty for kas %s", pk.URL)
}
ktype := ocrypto.KeyType(pk.Algorithm)
switch {
+ case ktype == "":
+ return "", "", "", fmt.Errorf("unsupported or unrecognized algorithm for kas %s", pk.URL)
case ocrypto.IsKEMKeyType(ktype):
return chunkedWrapKeyWithKEM(ktype, pk.PEM, symKey)
case ocrypto.IsECKeyType(ktype):
return chunkedWrapKeyWithEC(ktype, pk.PEM, symKey)
default:
return chunkedWrapKeyWithRSA(pk.PEM, symKey)
}
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // chunkedWrapKeyWithPublicKey dispatches to the right KAS wrapping | |
| // scheme based on the algorithm advertised by the splitter. | |
| func chunkedWrapKeyWithPublicKey(symKey []byte, pk KASPublicKey) (string, string, string, error) { | |
| if pk.PEM == "" { | |
| return "", "", "", fmt.Errorf("public key PEM is empty for kas %s", pk.URL) | |
| } | |
| ktype := ocrypto.KeyType(pk.Algorithm) | |
| switch { | |
| case ocrypto.IsKEMKeyType(ktype): | |
| return chunkedWrapKeyWithKEM(ktype, pk.PEM, symKey) | |
| case ocrypto.IsECKeyType(ktype): | |
| return chunkedWrapKeyWithEC(ktype, pk.PEM, symKey) | |
| default: | |
| return chunkedWrapKeyWithRSA(pk.PEM, symKey) | |
| } | |
| } | |
| // chunkedWrapKeyWithPublicKey dispatches to the right KAS wrapping | |
| // scheme based on the algorithm advertised by the splitter. | |
| func chunkedWrapKeyWithPublicKey(symKey []byte, pk KASPublicKey) (string, string, string, error) { | |
| if pk.PEM == "" { | |
| return "", "", "", fmt.Errorf("public key PEM is empty for kas %s", pk.URL) | |
| } | |
| ktype := ocrypto.KeyType(pk.Algorithm) | |
| switch { | |
| case ktype == "": | |
| return "", "", "", fmt.Errorf("unsupported or unrecognized algorithm for kas %s", pk.URL) | |
| case ocrypto.IsKEMKeyType(ktype): | |
| return chunkedWrapKeyWithKEM(ktype, pk.PEM, symKey) | |
| case ocrypto.IsECKeyType(ktype): | |
| return chunkedWrapKeyWithEC(ktype, pk.PEM, symKey) | |
| default: | |
| return chunkedWrapKeyWithRSA(pk.PEM, symKey) | |
| } | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@sdk/chunked_writer.go` around lines 747 - 762, Update
chunkedWrapKeyWithPublicKey so empty or unrecognized pk.Algorithm values return
an explicit unsupported-algorithm error instead of falling through to
chunkedWrapKeyWithRSA. Preserve the existing KEM and EC dispatch, and only
invoke RSA wrapping when the algorithm is explicitly recognized as an RSA type.
1 participant
Proposed Changes
Adds ChunkedWriter, an SDK entry point that produces a TDF from segments
supplied one at a time. Callers encrypt and upload each segment
independently — off-thread or in parallel — then call Finalize to seal
the archive. Contrast with
SDK.CreateTDF, which requires the fullplaintext buffered up front.
The writer exposes injection seams (
Clock,SegmentCipher,ArchiveWriterFactory,KeySplitter) so tests can substitutedeterministic or in-memory implementations without touching global
state. The default splitter is single-KAS only; callers with
attribute-based multi-KAS grants inject their own via
WithChunkedKeySplitter.No reader is included — produced TDFs round-trip through the mainline
SDK.LoadTDFpath.New files (all under
sdk/):chunked_writer.go— ChunkedWriter interface + concrete writerchunked_options.go—WithChunked*optionsclock.go— Clock interface + defaultssegment.go— SegmentCipher interface + AES-256-GCM defaultarchive_writer.go— ArchiveWriterFactory + ZIP64 defaultkey_splitter.go— KeySplitter interface + single-KAS defaultchunked_test.go— round-trip, keep-segments, assertion-rejectionJira: DSPX-2604
Checklist
make fmtcleansdkpackage tests pass with-raceTesting Instructions
Verifies three scenarios end-to-end:
TestChunkedRoundTrip— write 3 segments, decrypt via mainlineSDK.LoadTDFagainst in-process RSA-2048 fake KAS.TestChunkedKeepSegments— write 3, keep prefix[0, 1], decrypt only retained.TestChunkedFinalizeRejectsAssertions— assertions option returns sentinel error (unsupported).Summary by CodeRabbit