Add size method to ReadableStream returned by encrypt and decrypt

Author: FiloSottile

README.md

@@ -112,10 +112,12 @@ const e = new Encrypter()
 e.setScryptWorkFactor(12)
 e.setPassphrase("light-original-energy-average-wish-blind-vendor-pencil-illness-scorpion")
 const encryptedStream = await e.encrypt(file.stream())
+console.log(encryptedStream.size(file.size))
 
 const d = new Decrypter()
 d.addPassphrase("light-original-energy-average-wish-blind-vendor-pencil-illness-scorpion")
 const decryptedStream = await d.decrypt(encryptedStream)
+console.log(decryptedStream.size(encryptedStream.size(file.size)))
 
 console.log(await new Response(decryptedStream).text())
 ```

lib/format.ts

@@ -61,7 +61,9 @@ async function parseNextStanza(hdr: LineReader): Promise<{ s: Stanza, next?: nev
 }
 
 export async function parseHeader(header: ReadableStream<Uint8Array>): Promise<{
-    stanzas: Stanza[], MAC: Uint8Array, headerNoMAC: Uint8Array, rest: ReadableStream<Uint8Array>,
+    stanzas: Stanza[], MAC: Uint8Array,
+    headerNoMAC: Uint8Array, headerSize: number,
+    rest: ReadableStream<Uint8Array>,
 }> {
     const hdr = new LineReader(header)
     const versionLine = await hdr.readLine()
@@ -83,7 +85,7 @@ export async function parseHeader(header: ReadableStream<Uint8Array>): Promise<{
         const MAC = base64nopad.decode(macLine.slice(4))
         const { rest, transcript } = hdr.close()
         const headerNoMAC = transcript.slice(0, transcript.length - 1 - macLine.length + 3)
-        return { stanzas, headerNoMAC, MAC, rest: prepend(header, rest) }
+        return { stanzas, headerNoMAC, MAC, headerSize: transcript.length, rest: prepend(header, rest) }
     }
 }
 

lib/index.ts

@@ -4,7 +4,7 @@ import { sha256 } from "@noble/hashes/sha256"
 import { randomBytes } from "@noble/hashes/utils"
 import { ScryptIdentity, ScryptRecipient, X25519Identity, X25519Recipient } from "./recipients.js"
 import { encodeHeader, encodeHeaderNoMAC, parseHeader, Stanza } from "./format.js"
-import { decryptSTREAM, encryptSTREAM } from "./stream.js"
+import { ciphertextSize, decryptSTREAM, encryptSTREAM, plaintextSize } from "./stream.js"
 import { readAll, stream, read, readAllString, prepend } from "./io.js"
 
 export * as armor from "./armor.js"
@@ -64,6 +64,27 @@ export interface Recipient {
     wrapFileKey(fileKey: Uint8Array): Stanza[] | Promise<Stanza[]>;
 }
 
+/**
+ * A ReadableStream that can also report the expected output size based on the
+ * input size.
+ */
+export interface ReadableStreamWithSize extends ReadableStream<Uint8Array> {
+    /**
+     * Calculate the expected plaintext size from the given ciphertext size, or
+     * vice versa.
+     *
+     * @param sourceSize - The size of the ciphertext or plaintext to calculate
+     * the expected output size for.
+     *
+     * @returns The expected plaintext size if the input is a ciphertext, or the
+     * expected ciphertext size if the input is a plaintext.
+     *
+     * @throws Only if the input is a ciphertext, and the ciphertext size is
+     * too small or invalid. There are no invalid plaintext sizes.
+     */
+    size(sourceSize: number): number
+}
+
 export { generateIdentity, identityToRecipient } from "./recipients.js"
 
 /**
@@ -139,11 +160,11 @@ export class Encrypter {
      * encoded as UTF-8.
      *
      * @returns A promise that resolves to the encrypted file as a Uint8Array,
-     * or as a ReadableStream if the input was a stream.
+     * or as a {@link ReadableStreamWithSize} if the input was a stream.
      */
     async encrypt(file: Uint8Array | string): Promise<Uint8Array>
-    async encrypt(file: ReadableStream<Uint8Array>): Promise<ReadableStream<Uint8Array>>
-    async encrypt(file: Uint8Array | string | ReadableStream<Uint8Array>): Promise<Uint8Array | ReadableStream<Uint8Array>> {
+    async encrypt(file: ReadableStream<Uint8Array>): Promise<ReadableStreamWithSize>
+    async encrypt(file: Uint8Array | string | ReadableStream<Uint8Array>): Promise<Uint8Array | ReadableStreamWithSize> {
         const fileKey = randomBytes(16)
         const stanzas: Stanza[] = []
 
@@ -163,9 +184,13 @@ export class Encrypter {
         const streamKey = hkdf(sha256, fileKey, nonce, "payload", 32)
         const encrypter = encryptSTREAM(streamKey)
 
-        if (file instanceof ReadableStream) return prepend(file.pipeThrough(encrypter), header, nonce)
-        if (typeof file === "string") file = new TextEncoder().encode(file)
-        return await readAll(prepend(stream(file).pipeThrough(encrypter), header, nonce))
+        if (!(file instanceof ReadableStream)) {
+            if (typeof file === "string") file = new TextEncoder().encode(file)
+            return await readAll(prepend(stream(file).pipeThrough(encrypter), header, nonce))
+        }
+        return Object.assign(prepend(file.pipeThrough(encrypter), header, nonce), {
+            size: (size: number): number => ciphertextSize(size) + header.length + nonce.length
+        })
     }
 }
 
@@ -230,22 +255,25 @@ export class Decrypter {
      * Ignored if the input is a stream.
      *
      * @returns A promise that resolves to the decrypted file, or to a
-     * ReadableStream of the decrypted file if the input was a stream.
-     * The header is processed before the promise resolves.
+     * {@link ReadableStreamWithSize} of the decrypted file if the input was a
+     * stream. The header is processed before the promise resolves.
      */
     async decrypt(file: Uint8Array, outputFormat?: "uint8array"): Promise<Uint8Array>
     async decrypt(file: Uint8Array, outputFormat: "text"): Promise<string>
-    async decrypt(file: ReadableStream<Uint8Array>): Promise<ReadableStream<Uint8Array>>
-    async decrypt(file: Uint8Array | ReadableStream<Uint8Array>, outputFormat?: "text" | "uint8array"): Promise<string | Uint8Array | ReadableStream<Uint8Array>> {
+    async decrypt(file: ReadableStream<Uint8Array>): Promise<ReadableStreamWithSize>
+    async decrypt(file: Uint8Array | ReadableStream<Uint8Array>, outputFormat?: "text" | "uint8array"): Promise<string | Uint8Array | ReadableStreamWithSize> {
         const s = file instanceof ReadableStream ? file : stream(file)
-        const { fileKey, rest } = await this.decryptHeaderInternal(s)
+        const { fileKey, headerSize, rest } = await this.decryptHeaderInternal(s)
         const { data: nonce, rest: payload } = await read(rest, 16)
 
         const streamKey = hkdf(sha256, fileKey, nonce, "payload", 32)
         const decrypter = decryptSTREAM(streamKey)
         const out = payload.pipeThrough(decrypter)
 
-        if (file instanceof ReadableStream) return out
+        const outWithSize = Object.assign(out, {
+            size: (size: number): number => plaintextSize(size - headerSize - nonce.length)
+        })
+        if (file instanceof ReadableStream) return outWithSize
         if (outputFormat === "text") return await readAllString(out)
         return await readAll(out)
     }
@@ -267,7 +295,7 @@ export class Decrypter {
         return (await this.decryptHeaderInternal(stream(header))).fileKey
     }
 
-    private async decryptHeaderInternal(file: ReadableStream<Uint8Array>): Promise<{ fileKey: Uint8Array, rest: ReadableStream<Uint8Array> }> {
+    private async decryptHeaderInternal(file: ReadableStream<Uint8Array>): Promise<{ fileKey: Uint8Array, headerSize: number, rest: ReadableStream<Uint8Array> }> {
         const h = await parseHeader(file)
         const fileKey = await this.unwrapFileKey(h.stanzas)
         if (fileKey === null) throw Error("no identity matched any of the file's recipients")
@@ -276,7 +304,7 @@ export class Decrypter {
         const mac = hmac(sha256, hmacKey, h.headerNoMAC)
         if (!compareBytes(h.MAC, mac)) throw Error("invalid header HMAC")
 
-        return { fileKey: fileKey, rest: h.rest }
+        return { fileKey, headerSize: h.headerSize, rest: h.rest }
     }
 
     private async unwrapFileKey(stanzas: Stanza[]): Promise<Uint8Array | null> {

lib/stream.ts

@@ -48,6 +48,24 @@ export function decryptSTREAM(key: Uint8Array): TransformStream<Uint8Array, Uint
     })
 }
 
+export function plaintextSize(ciphertextSize: number): number {
+    if (ciphertextSize < chacha20poly1305Overhead) {
+        throw Error("ciphertext is too small")
+    }
+    if (ciphertextSize === chacha20poly1305Overhead) {
+        return 0 // Empty plaintext.
+    }
+    const fullChunks = Math.floor(ciphertextSize / chunkSizeWithOverhead)
+    const lastChunk = ciphertextSize % chunkSizeWithOverhead
+    if (0 < lastChunk && lastChunk <= chacha20poly1305Overhead) {
+        throw Error("ciphertext size is invalid")
+    }
+    let size = ciphertextSize
+    size -= fullChunks * chacha20poly1305Overhead
+    size -= lastChunk > 0 ? chacha20poly1305Overhead : 0
+    return size
+}
+
 export function encryptSTREAM(key: Uint8Array): TransformStream<Uint8Array, Uint8Array> {
     const streamNonce = new Uint8Array(12)
     const incNonce = () => {
@@ -84,3 +102,8 @@ export function encryptSTREAM(key: Uint8Array): TransformStream<Uint8Array, Uint
         },
     })
 }
+
+export function ciphertextSize(plaintextSize: number): number {
+    const chunks = Math.max(1, Math.ceil(plaintextSize / chunkSize))
+    return plaintextSize + chacha20poly1305Overhead * chunks
+}

tests/index.test.ts

@@ -140,9 +140,14 @@ describe("AgeEncrypter", function () {
         const source = randomBytesStream(size, chunk)
         const sourceHash = new HashingTransformStream()
         const encrypted = await e.encrypt(source.pipeThrough(sourceHash))
-        const decrypted = await d.decrypt(encrypted)
+        const expectedCiphertextSize = encrypted.size(size)
+        const ciphertextSize = new SizeTransformStream()
+        const encryptedWithSize = encrypted.pipeThrough(ciphertextSize)
+        const decrypted = await d.decrypt(encryptedWithSize)
+        assert.equal(size, decrypted.size(expectedCiphertextSize))
         const decryptedHash = new HashingTransformStream()
         await readAll(decrypted.pipeThrough(decryptedHash))
+        assert.equal(ciphertextSize.size, expectedCiphertextSize)
         assert.deepEqual(sourceHash.digest, decryptedHash.digest)
     })
     it("should throw when using multiple passphrases", function () {
@@ -214,3 +219,15 @@ class HashingTransformStream extends TransformStream<Uint8Array, Uint8Array> {
         })
     }
 }
+
+class SizeTransformStream extends TransformStream<Uint8Array, Uint8Array> {
+    size = 0
+    constructor() {
+        super({
+            transform: (chunk, controller) => {
+                this.size += chunk.length
+                controller.enqueue(chunk)
+            }
+        })
+    }
+}

tests/stream.test.ts

@@ -0,0 +1,20 @@
+import { describe, it, assert, expect } from "vitest"
+import { plaintextSize, ciphertextSize } from "../lib/stream.js"
+
+describe("stream", function () {
+    it.for([
+        0, 1, 15, 16, 17, 500,
+        64 * 1024 - 1, 64 * 1024, 64 * 1024 + 1,
+        64 * 1024 * 2 - 1, 64 * 1024 * 2, 64 * 1024 * 2 + 1
+    ])("should round-trip plaintext size and ciphertext size", function (ps) {
+        assert.equal(ps, plaintextSize(ciphertextSize(ps)),
+            `plaintextSize(ciphertextSize(${ps})) should return ${ps}`)
+    })
+    it.for([
+        0, 1, 15,
+        64 * 1024 + 16 + 1, 64 * 1024 + 16 + 15,
+        64 * 1024 * 2 + 16 * 2 + 1, 64 * 1024 * 2 + 16 * 2 + 15,
+    ])("should throw for invalid chiphertext size", function (cs) {
+        expect(() => plaintextSize(cs)).to.throw()
+    })
+})