fix: expose chunk address in chunksplitter

src/chunk/splitter.ts

@@ -52,8 +52,8 @@ export class ChunkBuilder {
   }
 }
 
-/** A sealed chunk awaiting upload, with its encryption key if encrypted. */
-export type ChunkEntry = { chunk: ChunkBuilder; key?: Uint8Array }
+/** A sealed chunk awaiting upload, with its address and its encryption key if encrypted. */
+export type ChunkEntry = { chunk: ChunkBuilder; address?: Reference; key?: Uint8Array }
 
 type PendingEntry = { entry: ChunkEntry; ref: Uint8Array; span: bigint }
 
@@ -161,11 +161,12 @@ export class ChunkSplitter {
       const ref = new Uint8Array(64)
       ref.set(address.toUint8Array())
       ref.set(key, 32)
-      this.pending[level]!.push({ entry: { chunk: this.chunks[level]!, key }, ref, span: originalSpan })
+      this.pending[level]!.push({ entry: { chunk: this.chunks[level]!, address, key }, ref, span: originalSpan })
     } else {
+      const address = this.chunks[level]!.hash()
       this.pending[level]!.push({
-        entry: { chunk: this.chunks[level]! },
-        ref: this.chunks[level]!.hash().toUint8Array(),
+        entry: { chunk: this.chunks[level]!, address },
+        ref: address.toUint8Array(),
         span: originalSpan,
       })
     }
@@ -184,8 +185,8 @@ export class ChunkSplitter {
     const parities = await this.onBatch(entries)
     if (parities.length > 0) {
       this.hasParity[level] = true
-      for (const { chunk } of parities) {
-        this.chunks[level]!.writer.write(new Uint8ArrayReader(chunk.hash().toUint8Array()))
+      for (const { chunk, address } of parities) {
+        this.chunks[level]!.writer.write(new Uint8ArrayReader((address ?? chunk.hash()).toUint8Array()))
       }
     }
   }

src/erasure-coding/batch.ts

@@ -1,4 +1,5 @@
 import { numberToUint64, uint64ToNumber } from '../bytes/encoding.js'
+import { Reference } from '../bytes/reference.js'
 import { ChunkBuilder, ChunkEntry } from '../chunk/splitter.js'
 import { encryptData, encryptSpan } from '../encryption/stream-cipher.js'
 import { getParities } from './levels.js'
@@ -38,11 +39,11 @@ export function makeIntermediateChunkHandler(level: number): (chunk: ChunkBuilde
 export function makeErasureBatch(
   level: number,
   encrypted: boolean,
-  onChunk: (chunk: ChunkBuilder, key?: Uint8Array) => Promise<void>,
+  onChunk: (chunk: ChunkBuilder, key?: Uint8Array, address?: Reference) => Promise<void>,
 ): (batch: ChunkEntry[]) => Promise<ChunkEntry[]> {
   return async (batch: ChunkEntry[]): Promise<ChunkEntry[]> => {
-    for (const { chunk, key } of batch) {
-      await onChunk(chunk, key)
+    for (const { chunk, key, address } of batch) {
+      await onChunk(chunk, key, address)
     }
     if (level <= 0) return []
 
@@ -57,8 +58,9 @@ export function makeErasureBatch(
       const parityEntries: ChunkEntry[] = []
       for (const bytes of parityShards) {
         const parityChunk = chunkFromBytes(bytes)
-        await onChunk(parityChunk)
-        parityEntries.push({ chunk: parityChunk })
+        const address = parityChunk.hash()
+        await onChunk(parityChunk, undefined, address)
+        parityEntries.push({ chunk: parityChunk, address })
       }
 
       return parityEntries
@@ -72,8 +74,9 @@ export function makeErasureBatch(
     const parityEntries: ChunkEntry[] = []
     for (const bytes of parityShards) {
       const chunk = chunkFromBytes(bytes)
-      await onChunk(chunk)
-      parityEntries.push({ chunk })
+      const address = chunk.hash()
+      await onChunk(chunk, undefined, address)
+      parityEntries.push({ chunk, address })
     }
 
     return parityEntries

src/mantaray/node.ts

@@ -441,7 +441,7 @@ export class MantarayNode {
    * callers can also create dispersed replicas from it.
    */
   async saveRecursively(
-    onChunk: (chunk: ChunkBuilder, key?: Uint8Array) => Promise<void>,
+    onChunk: (chunk: ChunkBuilder, key?: Uint8Array, address?: Reference) => Promise<void>,
   ): Promise<{ reference: Uint8Array; rootChunk: ChunkBuilder; encryptionKey?: Uint8Array }> {
     for (const fork of this.forks.values()) {
       if (fork.node.persisted) {
@@ -452,8 +452,8 @@ export class MantarayNode {
     }
 
     const onBatch = async (batch: ChunkEntry[]): Promise<ChunkEntry[]> => {
-      for (const { chunk, key } of batch) {
-        await onChunk(chunk, key)
+      for (const { chunk, key, address } of batch) {
+        await onChunk(chunk, key, address)
       }
 
       return []
@@ -465,15 +465,16 @@ export class MantarayNode {
 
     if (this.encrypt) {
       const { address, key } = rootChunk.encryptedHash()
-      await onChunk(rootChunk, key)
+      await onChunk(rootChunk, key, address)
       this.selfAddress = Bytes.concat(address.toUint8Array(), key)
       this.persisted = true
 
       return { reference: this.selfAddress, rootChunk, encryptionKey: key }
     }
 
-    await onChunk(rootChunk)
-    this.selfAddress = rootChunk.hash().toUint8Array()
+    const address = rootChunk.hash()
+    await onChunk(rootChunk, undefined, address)
+    this.selfAddress = address.toUint8Array()
     this.persisted = true
 
     return { reference: this.selfAddress, rootChunk }

test/chunk-addresses.test.ts

@@ -0,0 +1,76 @@
+import { describe, expect, it } from 'vitest'
+import { Reference } from '../src/bytes/reference.js'
+import { ChunkBuilder, ChunkEntry, ChunkSplitter } from '../src/chunk/splitter.js'
+import { makeErasureBatch, makeIntermediateChunkHandler } from '../src/erasure-coding/batch.js'
+import { getMaxShards } from '../src/erasure-coding/levels.js'
+import { MantarayNode } from '../src/mantaray/node.js'
+
+type Seen = { chunk: ChunkBuilder; key?: Uint8Array; address?: Reference }
+
+function expectedAddress({ chunk, key }: Seen): string {
+  return (key ? chunk.encryptedHash(key).address : chunk.hash()).toHex()
+}
+
+async function splitWithErasure(data: Uint8Array, level: number, encrypted: boolean): Promise<Seen[]> {
+  const seen: Seen[] = []
+  const onChunk = async (chunk: ChunkBuilder, key?: Uint8Array, address?: Reference) => {
+    seen.push({ chunk, key, address })
+  }
+  const splitter = new ChunkSplitter(
+    makeErasureBatch(level, encrypted, onChunk),
+    getMaxShards(level, encrypted),
+    encrypted,
+    makeIntermediateChunkHandler(level),
+  )
+  await splitter.append(data)
+  await splitter.finalize()
+
+  return seen
+}
+
+describe('chunk addresses passed to consumers', () => {
+  it('ChunkSplitter attaches the correct address to every batch entry', async () => {
+    const entries: ChunkEntry[] = []
+    const splitter = new ChunkSplitter(async batch => {
+      entries.push(...batch)
+
+      return []
+    })
+    await splitter.append(new Uint8Array(4096 * 300).fill(3))
+    await splitter.finalize()
+
+    expect(entries.length).toBeGreaterThan(300)
+    for (const entry of entries) {
+      expect(entry.address?.toHex()).toBe(expectedAddress(entry))
+    }
+  })
+
+  it.each([false, true])(
+    'makeErasureBatch passes correct addresses for data and parity chunks (encrypted: %s)',
+    async encrypted => {
+      const seen = await splitWithErasure(new Uint8Array(4096 * 200).fill(5), 2, encrypted)
+
+      const parities = seen.filter(s => !s.key)
+      expect(parities.length).toBeGreaterThan(0)
+      for (const s of seen) {
+        expect(s.address?.toHex()).toBe(expectedAddress(s))
+      }
+    },
+  )
+
+  it.each([false, true])('MantarayNode.saveRecursively passes correct addresses (encrypted: %s)', async encrypted => {
+    const node = new MantarayNode({ encrypt: encrypted })
+    node.addFork('a.txt', new Uint8Array(32).fill(1))
+    node.addFork('b/c.txt', new Uint8Array(32).fill(2))
+
+    const seen: Seen[] = []
+    await node.saveRecursively(async (chunk, key, address) => {
+      seen.push({ chunk, key, address })
+    })
+
+    expect(seen.length).toBeGreaterThan(0)
+    for (const s of seen) {
+      expect(s.address?.toHex()).toBe(expectedAddress(s))
+    }
+  })
+})