fix: accept a trailing empty final chunk in the count check

A plaintext of exactly k*chunkSize bytes has two valid encodings: k
chunks whose last one is full, or k full chunks followed by an empty
final chunk. Both decrypt to the same bytes, aesstream reads either (the
"full + empty final" row in TestChunkLayout pins it), and whole-object
Decrypt accepts either.

envelopePlaintextSize derived the expected chunk count with
chunkCountFor, which models only the first form, so a blob in the second
form declaring k+1 chunks failed with ErrSizeMismatch. DecryptRange and
PlaintextSize were stricter than Decrypt on the same object, which is
not a distinction this API means to draw.

Compare against the count aesstream derives from the ciphertext layout
instead, via a new exported ChunkCount. chunkCountFor stays as the
producer's rule for writing the header, matching the reference
implementation, and its doc now says so.

The check keeps its purpose: a size wrong by a whole chunk or more still
trips ErrSizeMismatch. A size short by only part of the final chunk now
accounts for the same number of chunks and passes, surfacing instead as
an authentication failure when that chunk is read -- the model the docs
already state for blobs carrying no chunk count. The two lengths are
indistinguishable without reading the chunks.

Also adds the exact-multiple boundary to the cross-implementation
vectors, which had no fixture for it: exact-multiple-go is 3 full chunks,
and the pinned foc-encryption decrypts it.

Assisted-by: Claude:claude-opus-5
Signed-off-by: Miroslav Bajtoš <oss@bajtos.net>

Author: bajtos

README.md

@@ -370,6 +370,13 @@ recipient   = [ {1: alg}, {4: kid, ...}, wrappedKey ]           # alg -31 or -5
 - **Body AAD** — the COSE `Enc_structure` over the protected header, identical
   for every chunk, so the algorithm, envelope type and any protected metadata
   are authenticated into the ciphertext.
+- **Chunking** — a producer writes `ceil(len / chunkSize)` chunks, minimum 1,
+  with the remainder in the final chunk; empty input is one empty chunk. A
+  decoder also accepts a stream that ends with an empty final chunk, so a
+  plaintext of exactly `k × chunkSize` bytes is a valid stream of either `k`
+  chunks or `k+1`. Only the ciphertext length distinguishes them
+  (`aesstream.ChunkCount`), which is why a declared count must never be checked
+  against a count re-derived from the plaintext length.
 - **Chunk count** (label −65791) — advisory metadata for range/seek consumers,
   emitted only when the plaintext length is known; not required to decrypt.
 

aesstream/spanreader.go

@@ -102,6 +102,29 @@ func DecryptedSize(ciphertextLen int64, chunkSize int) (int64, error) {
 	return plaintextLen, err
 }
 
+// ChunkCount returns how many chunks a complete ciphertext of ciphertextLen
+// bytes contains at the given chunk size. A zero chunkSize selects
+// DefaultChunkSize; any other value must be in [MinChunkSize, MaxChunkSize]
+// (else ErrChunkSize), the same rule as Config.ChunkSize. It returns
+// ErrCiphertextSize if ciphertextLen is not a structurally valid stream length.
+//
+// The count is not derivable from the plaintext length, which is why this is
+// worth asking for: the final chunk may be full, partial, or empty, so a
+// plaintext of exactly k*chunkSize bytes is a valid stream of either k chunks
+// (the last one full) or k+1 (the last one empty). Both are read the same way
+// and yield the same plaintext, and only the ciphertext length tells them apart.
+// A caller checking a stream against a separately recorded chunk count should
+// compare against this rather than against ceil(plaintextLen/chunkSize), which
+// only describes the first form.
+func ChunkCount(ciphertextLen int64, chunkSize int) (int64, error) {
+	chunkSize, err := resolveChunkSize(chunkSize)
+	if err != nil {
+		return 0, err
+	}
+	numChunks, _, _, err := chunkLayout(ciphertextLen, chunkSize)
+	return numChunks, err
+}
+
 // CiphertextRange returns the single contiguous ciphertext byte range
 // [start, start+n) that must be read to serve the plaintext range
 // [off, off+length) of a stream whose complete ciphertext is

aesstream/spanreader_test.go

@@ -497,3 +497,51 @@ func TestDecryptedSize_InvertsEncryptedSize(t *testing.T) {
 		require.Equalf(t, n, got, "DecryptedSize(EncryptedSize(%d))", n)
 	}
 }
+
+// TestChunkCount pins the exported chunk count over the boundary ciphertext
+// lengths, including the two encodings of an exact-multiple plaintext: the same
+// plaintext length is one chunk more when the stream ends with an empty final
+// chunk, which is why callers must not re-derive the count from a plaintext size.
+func TestChunkCount(t *testing.T) {
+	const cs = aesstream.MinChunkSize
+	enc := int64(cs) + aesstream.TagSize
+
+	cases := map[string]struct {
+		ciphertextLen int64
+		want          int64
+	}{
+		"empty stream (one tag-only chunk)": {aesstream.TagSize, 1},
+		"single partial chunk":              {aesstream.TagSize + 100, 1},
+		"single full chunk":                 {enc, 1},
+		"full + empty final":                {enc + aesstream.TagSize, 2},
+		"full + partial final":              {enc + aesstream.TagSize + 7, 2},
+		"two full chunks":                   {2 * enc, 2},
+	}
+	for name, c := range cases {
+		t.Run(name, func(t *testing.T) {
+			got, err := aesstream.ChunkCount(c.ciphertextLen, cs)
+			require.NoError(t, err)
+			require.Equal(t, c.want, got)
+		})
+	}
+}
+
+// TestChunkCountInvalid confirms ChunkCount reports the same errors as the rest
+// of the geometry API for a length or chunk size that cannot describe a stream.
+func TestChunkCountInvalid(t *testing.T) {
+	t.Run("ciphertext too short to hold a tag", func(t *testing.T) {
+		_, err := aesstream.ChunkCount(aesstream.TagSize-1, aesstream.MinChunkSize)
+		require.ErrorIs(t, err, aesstream.ErrCiphertextSize)
+	})
+
+	t.Run("chunk size out of range", func(t *testing.T) {
+		_, err := aesstream.ChunkCount(aesstream.TagSize, aesstream.MinChunkSize-1)
+		require.ErrorIs(t, err, aesstream.ErrChunkSize)
+	})
+
+	t.Run("zero chunk size selects the default", func(t *testing.T) {
+		got, err := aesstream.ChunkCount(aesstream.EncryptedSize(3*aesstream.DefaultChunkSize, 0), 0)
+		require.NoError(t, err)
+		require.Equal(t, int64(3), got)
+	})
+}

fee.go

@@ -410,6 +410,12 @@ func encryptStream(plaintext io.Reader, cek []byte, recipients []Recipient, opts
 
 // chunkCountFor reports how many STREAM chunks a plaintext of nPlain bytes
 // produces at the given chunk size. Empty input is one (empty) final chunk.
+//
+// This is the producer's rule, matching the reference implementation, and it is
+// only for writing the chunk-count header. It is not a decoder's rule: a stream
+// whose plaintext is an exact multiple of the chunk size may legitimately carry
+// one more chunk than this (an empty final chunk), so a count read off the wire
+// is checked against [aesstream.ChunkCount] instead.
 func chunkCountFor(nPlain, chunkSize int64) int64 {
 	if nPlain <= 0 {
 		return 1

range.go

@@ -9,7 +9,7 @@ import (
 	"github.com/filecoin-project/go-fee/cose"
 )
 
-// ErrSizeMismatch means the ciphertext length implied by the supplied blob size
+// ErrSizeMismatch means the number of chunks the supplied blob size accounts for
 // disagrees with the chunk count the envelope declares: the blob is truncated or
 // padded, or the size came from the wrong object.
 //
@@ -121,8 +121,11 @@ func (r *RangeReader) CiphertextSpan() (off, n int64) { return r.spanOff, r.span
 // are never read and so never checked, and — unlike whole-object [Decrypt] — a
 // range read cannot by itself detect that the stored object was truncated: an
 // understated blobSize describes a shorter object whose ranges decrypt cleanly.
-// The chunk-count check catches that when the envelope carries a chunk count, but
-// whole-object integrity is properly the job of the layer that supplied blobSize.
+// The chunk-count check catches an understatement of a whole chunk or more when
+// the envelope carries a count, but a size short by only part of the final chunk
+// accounts for the same number of chunks and so passes it, surfacing instead as
+// an authentication failure when that chunk is read. Whole-object integrity is
+// properly the job of the layer that supplied blobSize.
 func DecryptRange(blob io.ReaderAt, blobSize int64, unwrap RecipientUnwrapper, off, length int64) (*RangeReader, error) {
 	if unwrap == nil {
 		return nil, ErrNilUnwrapper
@@ -282,9 +285,16 @@ func spanRangeReader(blob io.ReaderAt, blobSize, headerLen int64, body bodyParam
 
 // envelopePlaintextSize is [plaintextSizeFrom] for a blob whose parameters came
 // from its envelope rather than from a caller's cache: it adds the cross-check of
-// the envelope's declared chunk count against the derived size, when one is
-// present, catching a blob size that describes a different object than the
-// envelope does.
+// the envelope's declared chunk count against the chunks the ciphertext actually
+// holds, when a count is present, catching a blob size that describes a different
+// object than the envelope does.
+//
+// The comparison is against [aesstream.ChunkCount] rather than a count re-derived
+// from the plaintext size, because a plaintext that is an exact multiple of the
+// chunk size seals to either k chunks or k+1 (with an empty final chunk) and both
+// decrypt identically. Deriving from the plaintext size would recognize only the
+// first form and refuse the second, which the whole-object [Decrypt] path reads
+// without complaint.
 func envelopePlaintextSize(env *cose.Envelope, blobSize, headerLen int64, chunkSize int) (int64, error) {
 	plainSize, err := plaintextSizeFrom(blobSize, headerLen, chunkSize)
 	if err != nil {
@@ -297,8 +307,12 @@ func envelopePlaintextSize(env *cose.Envelope, blobSize, headerLen int64, chunkS
 	if !ok {
 		return 0, fmt.Errorf("%w: chunk-count header is present but not an integer", ErrMalformedEnvelope)
 	}
-	if want := chunkCountFor(plainSize, int64(chunkSize)); declared != want {
-		return 0, fmt.Errorf("%w: envelope declares %d chunks, the blob size implies %d",
+	want, err := aesstream.ChunkCount(blobSize-headerLen, chunkSize)
+	if err != nil {
+		return 0, fmt.Errorf("fee: blob of %d ciphertext bytes: %w", blobSize-headerLen, err)
+	}
+	if declared != want {
+		return 0, fmt.Errorf("%w: envelope declares %d chunks, the ciphertext holds %d",
 			ErrSizeMismatch, declared, want)
 	}
 	return plainSize, nil

range_test.go

@@ -2,8 +2,11 @@ package fee_test
 
 import (
 	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
 	"crypto/ecdh"
 	"crypto/rand"
+	"encoding/binary"
 	"encoding/hex"
 	"errors"
 	"io"
@@ -88,6 +91,61 @@ func encryptWithCEK(t *testing.T, plaintext, cek []byte, recipients []fee.Recipi
 	return io.ReadAll(r)
 }
 
+// sealTrailingEmptyFinalChunk builds a FEE blob whose ciphertext is k full
+// chunks followed by an empty final chunk, declared as k+1 — the second valid
+// encoding of a plaintext that is an exact multiple of the chunk size.
+//
+// aesstream.Writer never emits this form (it flushes a full buffer as the final
+// chunk, giving k), so the body is sealed chunk by chunk here, under the
+// documented per-chunk nonce and the envelope's own Enc_structure as AAD.
+func sealTrailingEmptyFinalChunk(t *testing.T, plaintext, cek, baseNonce []byte, chunkSize int) []byte {
+	t.Helper()
+	require.Zero(t, len(plaintext)%chunkSize, "plaintext must be an exact multiple of the chunk size")
+	chunks := len(plaintext)/chunkSize + 1 // the trailing empty chunk is the extra one
+
+	env := &cose.Envelope{Headers: cose.Headers{
+		Protected: cose.Header{}.
+			Set(cose.HeaderLabelAlg, algChunkedStream).
+			Set(cose.HeaderLabelType, fee.EnvelopeType),
+		Unprotected: cose.Header{}.
+			Set(cose.HeaderLabelIV, baseNonce).
+			Set(labelChunkSize, int64(chunkSize)).
+			Set(labelChunkCount, int64(chunks)),
+	}}
+	aad, err := env.EncStructure(nil)
+	require.NoError(t, err)
+	header, err := env.Encode()
+	require.NoError(t, err)
+
+	block, err := aes.NewCipher(cek)
+	require.NoError(t, err)
+	aead, err := cipher.NewGCM(block)
+	require.NoError(t, err)
+
+	blob := bytes.Clone(header)
+	for i := range chunks {
+		last := i == chunks-1
+		var chunk []byte
+		if !last {
+			chunk = plaintext[i*chunkSize : (i+1)*chunkSize]
+		}
+		blob = aead.Seal(blob, streamNonce(baseNonce, i, last), chunk, aad)
+	}
+	return blob
+}
+
+// streamNonce builds the per-chunk GCM nonce of the FEE body cipher:
+// baseNonce(7) || chunkIndex(4, big-endian) || lastFlag(1).
+func streamNonce(baseNonce []byte, index int, last bool) []byte {
+	nonce := make([]byte, 0, aesstream.NonceSize)
+	nonce = append(nonce, baseNonce...)
+	nonce = binary.BigEndian.AppendUint32(nonce, uint32(index))
+	if last {
+		return append(nonce, 0x01)
+	}
+	return append(nonce, 0x00)
+}
+
 // hexBytes decodes a hex literal, for the handful of tests that pin behaviour
 // against exact wire bytes rather than an encrypted fixture.
 func hexBytes(t *testing.T, s string) []byte {
@@ -126,10 +184,24 @@ func decryptRange(t *testing.T, blob []byte, u fee.RecipientUnwrapper, off, leng
 // envelope, unwrap material and a byte range returns exactly the requested
 // plaintext — across the interesting geometries of a multi-chunk object: whole
 // object, within one chunk, spanning a boundary, exactly one aligned chunk, into
-// the short final chunk, and single bytes at each end.
+// the final chunk, and single bytes at each end. It runs the whole set against
+// both object shapes, since the final chunk being full rather than short is the
+// corner the declared chunk count is ambiguous at.
 func TestDecryptRangeRoundTrip(t *testing.T) {
-	const size = 3*rangeChunk + rangeChunk/2 // 3.5 chunks
-	f := newRangeFixture(t, size)
+	t.Run("partial final chunk", func(t *testing.T) {
+		testRangeRoundTrip(t, 3*rangeChunk+rangeChunk/2)
+	})
+
+	// An exact multiple of the chunk size: the final chunk is full rather than
+	// short, the geometry corner the wire format's chunk count is ambiguous at.
+	t.Run("exact multiple of the chunk size", func(t *testing.T) {
+		testRangeRoundTrip(t, 4*rangeChunk)
+	})
+}
+
+func testRangeRoundTrip(t *testing.T, size int64) {
+	t.Helper()
+	f := newRangeFixture(t, int(size), fee.WithContentLength(size))
 
 	cases := []struct {
 		name        string
@@ -144,7 +216,7 @@ func TestDecryptRangeRoundTrip(t *testing.T) {
 		{"across two boundaries", rangeChunk - 10, 2*rangeChunk + 20},
 		{"exactly one aligned chunk", rangeChunk, rangeChunk},
 		{"aligned start, unaligned end", 2 * rangeChunk, rangeChunk + 5},
-		{"whole short final chunk", 3 * rangeChunk, rangeChunk / 2},
+		{"from the last chunk's start", 3 * rangeChunk, rangeChunk / 2},
 		{"into final chunk", 3*rangeChunk - 5, 100},
 		{"ends exactly on boundary", rangeChunk / 2, rangeChunk / 2},
 		{"length past end clamps", size - 10, 1000},
@@ -156,7 +228,7 @@ func TestDecryptRangeRoundTrip(t *testing.T) {
 
 			wantLen := clampLen(size, tc.off, tc.length)
 			require.Equal(t, wantLen, r.Len(), "Len is the clamped range length")
-			require.Equal(t, int64(size), r.Size(), "Size is the whole object")
+			require.Equal(t, size, r.Size(), "Size is the whole object")
 			require.Equal(t, f.plaintext[tc.off:tc.off+wantLen], got)
 		})
 	}
@@ -603,6 +675,50 @@ func TestDecryptRangeChunkCountMismatch(t *testing.T) {
 	})
 }
 
+// TestDecryptRangeTrailingEmptyFinalChunk pins that the range path accepts the
+// second valid encoding of an exact-multiple plaintext: k full chunks followed by
+// an empty final chunk, declared as k+1. aesstream reads that layout and so does
+// whole-object decryption, so the range entry points must agree rather than
+// refusing a blob the rest of the package accepts.
+func TestDecryptRangeTrailingEmptyFinalChunk(t *testing.T) {
+	const size = 3 * rangeChunk
+	cek, plaintext := newCEK(t), patternBytes(size)
+	baseNonce := bytes.Repeat([]byte{0xa5}, aesstream.BaseNonceSize)
+	blob := sealTrailingEmptyFinalChunk(t, plaintext, cek, baseNonce, rangeChunk)
+	blobSize := int64(len(blob))
+
+	// The premise: whole-object decryption already reads this blob.
+	whole, err := fee.DecryptWithCEK(bytes.NewReader(blob), cek)
+	require.NoError(t, err)
+	got, err := io.ReadAll(whole)
+	require.NoError(t, err)
+	require.Equal(t, plaintext, got)
+
+	t.Run("PlaintextSize agrees", func(t *testing.T) {
+		n, err := fee.PlaintextSize(bytes.NewReader(blob), blobSize)
+		require.NoError(t, err)
+		require.Equal(t, int64(size), n)
+	})
+
+	t.Run("range across the last full chunk", func(t *testing.T) {
+		off, length := int64(2*rangeChunk-10), int64(20)
+		r, err := fee.DecryptRangeWithCEK(bytes.NewReader(blob), blobSize, cek, off, length)
+		require.NoError(t, err)
+		require.Equal(t, int64(size), r.Size())
+		got, err := io.ReadAll(r)
+		require.NoError(t, err)
+		require.Equal(t, plaintext[off:off+length], got)
+	})
+
+	t.Run("whole object as one range", func(t *testing.T) {
+		r, err := fee.DecryptRangeWithCEK(bytes.NewReader(blob), blobSize, cek, 0, size)
+		require.NoError(t, err)
+		got, err := io.ReadAll(r)
+		require.NoError(t, err)
+		require.Equal(t, plaintext, got)
+	})
+}
+
 // TestDecryptRangeWrongBlobSizeNoChunkCount pins the documented trust model for an
 // envelope with no chunk count: a wrong blob size cannot be caught at construction,
 // so it surfaces as a failure to authenticate when the affected chunks are read.

vectors/README.md

@@ -24,6 +24,7 @@ to it and this repo matches it (see [Wire format](#wire-format)).
 | `multi-chunk-ts` | TS seals → Go decrypts | tag 16 (COSE_Encrypt0) |
 | `multi-recipient-go` | Go seals → TS parses recipients + decrypts body | tag 96 (COSE_Encrypt) |
 | `multi-chunk-go` | Go seals → TS decrypts (extra multi-chunk coverage) | tag 16 |
+| `exact-multiple-go` | Go seals → TS decrypts (plaintext is exactly 3 chunks) | tag 16 |
 
 Each `testdata/<name>/` holds `blob.bin` (`envelope‖ciphertext`),
 `plaintext.bin`, and `meta.json`.
@@ -54,6 +55,12 @@ recipient   = [ {1: alg}, {4: kid, ...}, wrappedKey ]           # alg -31 or -5
 - **Body cipher** — chunked AES-256-GCM-STREAM, alg `-65793`. Per-chunk nonce is
   `baseNonce[7] ‖ chunkIndex[4, big-endian] ‖ lastFlag[1]` (`0x01` on the final
   chunk), tag 16 bytes.
+- **Chunking** — a producer writes `ceil(len / chunkSize)` chunks, minimum 1,
+  with the remainder in the final chunk; empty input is one empty chunk. Both
+  implementations follow that rule, so `exact-multiple-go` declares 3 chunks
+  rather than 3 full chunks plus an empty one. A *decoder* also accepts a
+  trailing empty final chunk, so the declared count is authoritative and must
+  not be re-derived from the plaintext length.
 - **Body AAD** — `Enc_structure = [ context, protected, "" ]`, the **same** for
   every chunk. `context` follows the envelope structure per RFC 9052 §5.3:
   `"Encrypt"` for a tag-96 envelope, `"Encrypt0"` for tag-16. AAD interop is

vectors/testdata/exact-multiple-go/blob.bin

@@ -0,0 +1,12 @@
+{
+  "name": "exact-multiple-go",
+  "producer": "go",
+  "description": "Plaintext of exactly 3 chunks (full final chunk) encrypted in Go; decrypts in foc-encryption (TS).",
+  "tag": 16,
+  "algorithm": -65793,
+  "typ": "application/vnd.foc-envelope+cose",
+  "chunk_size": 4096,
+  "chunk_count": 3,
+  "cek_hex": "1b73245d0995604fb9b7040935a46bc6db4e2b6ec9f907b8a528810925166867",
+  "base_nonce_hex": "baf5c82a5faa4c"
+}