fix(sdk): DSPX-4590 default per-segment sizes when a writer omits them

sdk/manifest.go

@@ -1,5 +1,15 @@
 package sdk
 
+import "fmt"
+
+// Segment describes one chunk of the payload.
+//
+// Size and EncryptedSize are optional in the wire format: a writer may omit
+// either key whenever its value equals the corresponding manifest-level
+// default. Since JSON can't distinguish an omitted key from an explicit 0,
+// a zero EncryptedSize always means "omitted" -- ciphertext is never
+// legitimately zero bytes -- but a zero Size is ambiguous; see
+// resolveSegmentSizes for how it's disambiguated.
 type Segment struct {
 	Hash          string `json:"hash"`
 	Size          int64  `json:"segmentSize"`
@@ -19,6 +29,50 @@ type IntegrityInformation struct {
 	Segments                []Segment `json:"segments"`
 }
 
+// resolveSegmentSizes returns the plaintext and ciphertext sizes of seg in
+// bytes, substituting the manifest-level default for whichever field the
+// writer omitted, and validates that the resolved pair is internally
+// consistent.
+//
+// EncryptedSize is never ambiguous on its own: ciphertext is never
+// legitimately zero-length (there is always at least a nonce and a tag), so
+// a raw 0 always means the key was left out because it equals
+// DefaultEncryptedSegSize.
+//
+// Size is ambiguous on its own: web-sdk, for example, omits Size and
+// EncryptedSize independently of each other, each time its own value equals
+// the manifest-level default -- so a zero Size doesn't necessarily mean
+// EncryptedSize was omitted too. Disambiguate by comparing the resolved
+// EncryptedSize against its own default instead.
+//
+// AES-GCM frames every segment with a fixed-size nonce and tag, so the
+// plaintext size is pinned by the ciphertext size regardless of which
+// fields the manifest declared explicitly. A resolved pair that disagrees
+// with that framing is rejected here -- whether the inconsistency came from
+// the per-segment fields or the manifest-level defaults -- rather than left
+// for a caller to discover downstream.
+func (i IntegrityInformation) resolveSegmentSizes(seg Segment) (int64, int64, error) {
+	encryptedSize := seg.EncryptedSize
+	if encryptedSize == 0 {
+		encryptedSize = i.DefaultEncryptedSegSize
+	}
+
+	size := seg.Size
+	if size == 0 && encryptedSize == i.DefaultEncryptedSegSize {
+		size = i.DefaultSegmentSize
+	}
+
+	if size < 0 || encryptedSize <= 0 {
+		return 0, 0, fmt.Errorf("%w: segmentSize=%d encryptedSegmentSize=%d", ErrSegSizeUnresolved, size, encryptedSize)
+	}
+
+	if encryptedSize < gcmIvSize+aesBlockSize || size != encryptedSize-(gcmIvSize+aesBlockSize) {
+		return 0, 0, fmt.Errorf("%w: segment declares size %d with encrypted size %d", ErrSegSizeMismatch, size, encryptedSize)
+	}
+
+	return size, encryptedSize, nil
+}
+
 type KeyAccess struct {
 	KeyType            string      `json:"type"`
 	KasURL             string      `json:"url"`

sdk/tdf.go

@@ -879,7 +879,14 @@ func (s SDK) LoadTDF(reader io.ReadSeeker, opts ...TDFReaderOption) (*Reader, er
 
 	var payloadSize int64
 	for _, seg := range manifestObj.Segments {
-		payloadSize += seg.Size
+		// Sizes the writer left to the manifest-level default have to be
+		// filled in here too: without it the payload looks shorter than it
+		// is, and every read bounded by payloadSize comes up short.
+		size, _, err := manifestObj.resolveSegmentSizes(seg)
+		if err != nil {
+			return nil, err
+		}
+		payloadSize += size
 	}
 
 	return &Reader{
@@ -956,18 +963,28 @@ func (r *Reader) WriteTo(writer io.Writer) (int64, error) {
 	var payloadReadOffset int64
 	var decryptedDataOffset int64
 	for _, seg := range r.manifest.Segments {
-		if decryptedDataOffset+seg.Size < r.cursor {
-			decryptedDataOffset += seg.Size
-			payloadReadOffset += seg.EncryptedSize
+		// resolveSegmentSizes rejects a declared Size that disagrees with
+		// EncryptedSize; without that check here too, decryptedDataOffset
+		// could run ahead of the actual decrypted length, panicking on
+		// writeBuf[offset:] below once a later segment's slice runs shorter
+		// than expected.
+		segSize, encryptedSegSize, err := r.manifest.resolveSegmentSizes(seg)
+		if err != nil {
+			return totalBytes, err
+		}
+
+		if decryptedDataOffset+segSize < r.cursor {
+			decryptedDataOffset += segSize
+			payloadReadOffset += encryptedSegSize
 			continue
 		}
 
-		readBuf, err := r.tdfReader.ReadPayload(payloadReadOffset, seg.EncryptedSize)
+		readBuf, err := r.tdfReader.ReadPayload(payloadReadOffset, encryptedSegSize)
 		if err != nil {
 			return totalBytes, fmt.Errorf("TDFReader.ReadPayload failed: %w", err)
 		}
 
-		if int64(len(readBuf)) != seg.EncryptedSize {
+		if int64(len(readBuf)) != encryptedSegSize {
 			return totalBytes, ErrSegSizeMismatch
 		}
 
@@ -1006,9 +1023,9 @@ func (r *Reader) WriteTo(writer io.Writer) (int64, error) {
 			return totalBytes, errWriteFailed
 		}
 
-		payloadReadOffset += seg.EncryptedSize
+		payloadReadOffset += encryptedSegSize
 		r.cursor += int64(n)
-		decryptedDataOffset += seg.Size
+		decryptedDataOffset += segSize
 	}
 
 	return totalBytes, nil
@@ -1059,24 +1076,23 @@ func (r *Reader) ReadAt(buf []byte, offset int64) (int, error) { //nolint:funlen
 		// Segment.Size positions every plaintext offset derived below --
 		// including for the segments this request skips over -- but nothing
 		// authenticates it: the root signature aggregates only Segment.Hash.
-		// AES-GCM frames each segment with a fixed-size nonce and tag, so the
-		// plaintext size is pinned by the ciphertext size, and ReadPayload
-		// below checks EncryptedSize against the bytes actually present. This
-		// is the per-segment form of the check doPayloadKeyUnwrap already
-		// applies to the manifest defaults. Deriving Size from EncryptedSize
-		// rather than the reverse keeps the arithmetic from overflowing.
-		if seg.EncryptedSize < gcmIvSize+aesBlockSize || seg.Size != seg.EncryptedSize-(gcmIvSize+aesBlockSize) {
-			return 0, fmt.Errorf("%w: segment declares size %d with encrypted size %d",
-				ErrSegSizeMismatch, seg.Size, seg.EncryptedSize)
+		// resolveSegmentSizes pins the plaintext size to the ciphertext size
+		// (AES-GCM frames each segment with a fixed-size nonce and tag), and
+		// ReadPayload below checks EncryptedSize against the bytes actually
+		// present. This is the per-segment form of the check
+		// doPayloadKeyUnwrap already applies to the manifest defaults.
+		segSize, encryptedSegSize, err := r.manifest.resolveSegmentSizes(seg)
+		if err != nil {
+			return 0, err
 		}
 
-		segEnd := segStart + seg.Size
+		segEnd := segStart + segSize
 
 		// Wholly before the request. The comparison is <= rather than < so
 		// that a request starting exactly on a segment boundary, or a
 		// zero-length request, does not pull in the preceding segment.
 		if segEnd <= offset {
-			payloadReadOffset += seg.EncryptedSize
+			payloadReadOffset += encryptedSegSize
 			segStart = segEnd
 			continue
 		}
@@ -1090,12 +1106,12 @@ func (r *Reader) ReadAt(buf []byte, offset int64) (int, error) { //nolint:funlen
 			startIndex = offset - segStart
 		}
 
-		readBuf, err := r.tdfReader.ReadPayload(payloadReadOffset, seg.EncryptedSize)
+		readBuf, err := r.tdfReader.ReadPayload(payloadReadOffset, encryptedSegSize)
 		if err != nil {
 			return 0, fmt.Errorf("TDFReader.ReadPayload failed: %w", err)
 		}
 
-		if int64(len(readBuf)) != seg.EncryptedSize {
+		if int64(len(readBuf)) != encryptedSegSize {
 			return 0, ErrSegSizeMismatch
 		}
 
@@ -1128,7 +1144,7 @@ func (r *Reader) ReadAt(buf []byte, offset int64) (int, error) { //nolint:funlen
 			return 0, errWriteFailed
 		}
 
-		payloadReadOffset += seg.EncryptedSize
+		payloadReadOffset += encryptedSegSize
 		segStart = segEnd
 	}
 
@@ -1553,7 +1569,7 @@ func calculateSignature(data []byte, secret []byte, alg IntegrityAlgorithm, isLe
 		return string(hmac), nil
 	}
 	if kGMACPayloadLength > len(data) {
-		return "", errors.New("fail to create gmac signature")
+		return "", fmt.Errorf("%w: ciphertext length=%d", ErrGMACSignatureFailed, len(data))
 	}
 
 	if isLegacyTDF {

sdk/tdf_helpers_test.go

@@ -43,6 +43,23 @@ func TestIntegrityAlgorithmStringMatchesCalculateSignature(t *testing.T) {
 	}
 }
 
+// GMAC signs by returning the ciphertext's trailing kGMACPayloadLength bytes
+// verbatim; a ciphertext shorter than that has no tag to return and must be
+// rejected rather than silently truncated or padded.
+func TestCalculateSignatureGMACShortCiphertext(t *testing.T) {
+	key := make([]byte, kKeySize)
+	_, err := rand.Read(key)
+	require.NoError(t, err)
+
+	data := make([]byte, kGMACPayloadLength-1)
+	_, err = rand.Read(data)
+	require.NoError(t, err)
+
+	_, err = calculateSignature(data, key, GMAC, false)
+	require.ErrorIs(t, err, ErrGMACSignatureFailed)
+	require.ErrorIs(t, err, ErrTampered)
+}
+
 func TestCreatePolicyBinding(t *testing.T) {
 	symKey := make([]byte, kKeySize)
 	_, err := rand.Read(symKey)

sdk/tdf_readat_test.go

@@ -217,49 +217,69 @@ func TestReaderReadAtNonUniformEdges(t *testing.T) {
 
 // TestReaderReadAtDeclaredSizeMismatch checks that a manifest whose declared
 // segment sizes disagree with the AES-GCM framing is rejected rather than
-// trusted.
+// trusted, on both ReadAt and WriteTo.
 //
 // Nothing authenticates Segment.Size: the root signature aggregates only each
-// segment's Hash, and the schema types the size as a bare number. ReadAt derives
-// every plaintext offset from those sizes, so an altered Size shifts the mapping
-// -- and because a segment before the requested offset is skipped rather than
-// decrypted, checking the length that comes back from Decrypt is not enough on
-// its own to catch it.
+// segment's Hash, and the schema types the size as a bare number. ReadAt and
+// WriteTo both derive plaintext offsets from those sizes, so an altered Size
+// shifts the mapping -- and because a segment before the requested offset is
+// skipped rather than decrypted, checking the length that comes back from
+// Decrypt is not enough on its own to catch it. For WriteTo specifically, the
+// same drift would otherwise let decryptedDataOffset run ahead of the actual
+// decrypted length, panicking on writeBuf[offset:] once a later segment's
+// slice comes up short.
 func TestReaderReadAtDeclaredSizeMismatch(t *testing.T) {
 	for _, tc := range []struct {
-		name   string
-		mutate func(segments []Segment)
+		name    string
+		mutate  func(segments []Segment)
+		wantErr error
 	}{
 		// Understating the first segment shifts every later segment down by
 		// five bytes. The read below starts past that segment, so it is skipped
 		// and never decrypted.
-		{"understated", func(segments []Segment) { segments[0].Size = 5 }},
-		{"overstated", func(segments []Segment) { segments[0].Size = 40 }},
+		{"understated", func(segments []Segment) { segments[0].Size = 5 }, ErrSegSizeMismatch},
+		{"overstated", func(segments []Segment) { segments[0].Size = 40 }, ErrSegSizeMismatch},
 		// Sizes that sum back to something plausible: payloadSize is the sum of
 		// every Size, so a pair that overflows to a small positive total gets
-		// past the range check on offset and reaches the segment walk.
+		// past the range check on offset and reaches the segment walk. A
+		// negative declared Size is caught by resolveSegmentSizes itself,
+		// before the arithmetic consistency check below it ever runs.
 		{"negative", func(segments []Segment) {
 			segments[0].Size = math.MinInt64 + 1
 			segments[1].Size = math.MinInt64 + 7
-		}},
+		}, ErrSegSizeUnresolved},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			reader, _ := newNonUniformReader(t, []int{10, 10, 10})
 			tc.mutate(reader.manifest.Segments)
 
-			// Mirror what LoadTDF derives from the tampered manifest.
+			// LoadTDF would itself reject this manifest via resolveSegmentSizes
+			// before a Reader ever exists, so payloadSize is set by hand here to
+			// unit-test ReadAt/WriteTo against the tampered manifest directly,
+			// bypassing that earlier rejection.
 			var payloadSize int64
 			for _, seg := range reader.manifest.Segments {
 				payloadSize += seg.Size
 			}
 			reader.payloadSize = payloadSize
 
-			// The request spans the tampered segment and the one after it, so a
-			// reader that trusted Size would report a full 20 bytes of shifted
-			// plaintext rather than an error.
-			n, err := reader.ReadAt(make([]byte, 20), 5)
-			require.ErrorIs(t, err, ErrSegSizeMismatch)
-			assert.Zero(t, n)
+			t.Run("ReadAt", func(t *testing.T) {
+				// The request spans the tampered segment and the one after it, so a
+				// reader that trusted Size would report a full 20 bytes of shifted
+				// plaintext rather than an error.
+				n, err := reader.ReadAt(make([]byte, 20), 5)
+				require.ErrorIs(t, err, tc.wantErr)
+				assert.Zero(t, n)
+			})
+
+			t.Run("WriteTo", func(t *testing.T) {
+				// The tampered segment is first, so WriteTo hits the same error on
+				// its very first iteration, before writing any bytes.
+				var out bytes.Buffer
+				n, err := reader.WriteTo(&out)
+				require.ErrorIs(t, err, tc.wantErr)
+				assert.Zero(t, n)
+			})
 		})
 	}
 }