feat: add aggregate decode limits

packages/protons-runtime/src/codec.ts

@@ -65,6 +65,13 @@ export interface DecodeOptions<T> {
    * Runtime-specified limits for lengths of repeated/map fields
    */
   limits?: Limits<T>
+
+  /**
+   * Maximum wire elements at each repeated/map field path across the entire
+   * decode or stream operation. Uses the same shape as `limits`; map entries
+   * count even when duplicate keys overwrite earlier entries.
+   */
+  totalLimits?: Limits<T>
 }
 
 export interface DecodeFunction<T> {

packages/protons-runtime/src/codecs/message.ts

@@ -1,10 +1,32 @@
 import { createCodec, CODEC_TYPES } from '../codec.ts'
+import { UnsupportedDecodeLimitsError } from '../index.ts'
 import type { EncodeFunction, DecodeFunction, Codec, StreamFunction } from '../codec.ts'
 
 export interface Factory<A, T> {
   new (obj: A): T
 }
 
-export function message <D, E> (encode: EncodeFunction<E>, decode: DecodeFunction<D>, stream: StreamFunction<D>): Codec<D, E> {
-  return createCodec('message', CODEC_TYPES.LENGTH_DELIMITED, encode, decode, stream)
+export interface MessageCodecOptions {
+  totalLimits?: boolean
+}
+
+export function message <D, E> (encode: EncodeFunction<E>, decode: DecodeFunction<D>, stream: StreamFunction<D>, options: MessageCodecOptions = {}): Codec<D, E> {
+  const codec = createCodec('message', CODEC_TYPES.LENGTH_DELIMITED, encode, decode, stream)
+
+  if (options.totalLimits !== true) {
+    codec.decode = (r, length, opts) => {
+      if (opts?.totalLimits != null) {
+        throw new UnsupportedDecodeLimitsError('Regenerate this codec to use totalLimits')
+      }
+      return decode(r, length, opts)
+    }
+    codec.stream = function * (r, length, prefix, opts) {
+      if (opts?.totalLimits != null) {
+        throw new UnsupportedDecodeLimitsError('Regenerate this codec to use totalLimits')
+      }
+      yield * stream(r, length, prefix, opts)
+    }
+  }
+
+  return codec
 }

packages/protons-runtime/src/decode-context.ts

@@ -0,0 +1,53 @@
+import { MaxLengthError, MaxSizeError } from './index.ts'
+
+const contextKey = Symbol('decodeContext')
+
+interface InternalDecodeOptions {
+  totalLimits?: unknown
+  [contextKey]?: DecodeContext
+}
+
+/** Shared counters for a field path during one decode or stream operation. */
+export class DecodeContext {
+  private readonly counts = new Map<string, number>()
+  private readonly children = new Map<string, DecodeContext>()
+
+  private constructor () {}
+
+  static get (opts: InternalDecodeOptions): DecodeContext | undefined {
+    if (opts.totalLimits == null) {
+      return
+    }
+
+    return opts[contextKey] ?? new DecodeContext()
+  }
+
+  static child<T extends InternalDecodeOptions> (context: DecodeContext | undefined, field: string, opts: T): T {
+    if (context == null || opts.totalLimits == null) {
+      return opts
+    }
+
+    let child = context.children.get(field)
+
+    if (child == null) {
+      child = new DecodeContext()
+      context.children.set(field, child)
+    }
+
+    return { ...opts, [contextKey]: child }
+  }
+
+  consume (field: string, limit: number | undefined, ErrorType: typeof MaxLengthError | typeof MaxSizeError = MaxLengthError): void {
+    if (limit == null) {
+      return
+    }
+
+    const count = this.counts.get(field) ?? 0
+
+    if (count >= limit) {
+      throw new ErrorType(`Decode error - field "${field}" exceeded its total element limit`)
+    }
+
+    this.counts.set(field, count + 1)
+  }
+}

packages/protons-runtime/src/index.ts

@@ -45,6 +45,7 @@ export type {
 } from './codec.ts'
 
 export { enumeration } from './codecs/enum.ts'
+export { DecodeContext } from './decode-context.ts'
 export { message } from './codecs/message.ts'
 export { createReader as reader } from './utils/reader.ts'
 export { createWriter as writer } from './utils/writer.ts'
@@ -396,6 +397,13 @@ export class MaxSizeError extends Error {
   public name = 'MaxSizeError'
 }
 
+/**
+ * Thrown when aggregate limits are requested from a codec that cannot enforce them
+ */
+export class UnsupportedDecodeLimitsError extends Error {
+  public name = 'UnsupportedDecodeLimitsError'
+}
+
 export class ParseError extends Error {
   /**
    * This will be removed in a future release

packages/protons-runtime/src/stream.ts

@@ -1,8 +1,8 @@
 import { createReader } from './utils/reader.ts'
-import type { Codec } from './codec.ts'
+import type { Codec, DecodeOptions } from './codec.ts'
 import type { Uint8ArrayList } from 'uint8arraylist'
 
-export function * streamMessage <D, E> (buf: Uint8Array | Uint8ArrayList, codec: Pick<Codec<D, E>, 'stream'>, opts?: any): Generator<any> {
+export function * streamMessage <D, E> (buf: Uint8Array | Uint8ArrayList, codec: Pick<Codec<D, E>, 'stream'>, opts?: DecodeOptions<D>): Generator<any> {
   const reader = createReader(buf)
 
   yield * codec.stream(reader, undefined, '.', opts)

packages/protons/src/fields/array-field.ts

@@ -148,6 +148,7 @@ export class ArrayField extends Field {
     parent.addImport('protons-runtime', 'MaxLengthError')
 
     let limit = `
+              decodeContext?.consume('${this.name}', opts.totalLimits?.${this.name})
               if (opts.limits?.${this.name} != null && obj.${this.name}.length === opts.limits.${this.name}) {
                 throw new MaxLengthError('Decode error - repeated field "${this.name}" had too many elements')
               }
@@ -188,6 +189,7 @@ export class ArrayField extends Field {
     parent.addImport('protons-runtime', 'MaxLengthError')
 
     let limit = `
+              decodeContext?.consume('${this.name}', opts.totalLimits?.${this.name})
               if (opts.limits?.${this.name} != null && obj.${this.name} === opts.limits.${this.name}) {
                 throw new MaxLengthError('Streaming decode error - repeated field "${this.name}" had too many elements')
               }

packages/protons/src/fields/map-field.ts

@@ -77,6 +77,7 @@ export class MapField extends Field {
     const entryType = parent.findType(this.entryType)
 
     let limit = `
+              decodeContext?.consume('${this.name}', opts.totalLimits?.${this.name}, MaxSizeError)
               if (opts.limits?.${this.name} != null && obj.${this.name}.size === opts.limits.${this.name}) {
                 throw new MaxSizeError('Decode error - map field "${this.name}" had too many elements')
               }
@@ -98,9 +99,11 @@ export class MapField extends Field {
   }
 
   getStreamingDecoder (parent: Parent): string {
+    parent.addImport('protons-runtime', 'MaxSizeError')
     parent.addImport('protons-runtime', 'MaxLengthError')
 
     let limit = `
+              decodeContext?.consume('${this.name}', opts.totalLimits?.${this.name}, MaxSizeError)
               if (opts.limits?.${this.name} != null && obj.${this.name} === opts.limits.${this.name}) {
                 throw new MaxLengthError('Decode error - map field "${this.name}" had too many elements')
               }

packages/protons/src/index.ts

@@ -170,6 +170,25 @@
  * })
  * ```
  *
+ * #### Limiting totals across nested messages
+ *
+ * `limits` applies separately to each container. To bound the total number of
+ * decoded elements at a field path, use `totalLimits` with the same shape:
+ *
+ * ```TypeScript
+ * const message = MyMessage.decode(buf, {
+ *   limits: { messages$value: { repeatedField: 5 } },
+ *   totalLimits: { messages$value: { repeatedField: 10 } }
+ * })
+ * ```
+ * Each nested list may contain up to five elements, and all those lists together
+ * may contain up to ten. Totals include repeated occurrences of singular parent
+ * messages and duplicate map entries, even when earlier values are overwritten.
+ * Limits are checked before decoding the next element. The same options work
+ * with `.stream()` and can be reused: each operation starts with fresh counters.
+ * Regenerate existing codecs to use this option. Older generated codecs throw
+ * `UnsupportedDecodeLimitsError` when passed `totalLimits`.
+ *
  * ### Overriding 64 bit types
  *
  * By default 64 bit types are implemented as [BigInt](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/BigInt)s.

packages/protons/src/types/message.ts

@@ -176,19 +176,22 @@ export class Message implements Type {
     let opts = ''
 
     if (field instanceof MessageField) {
-      opts = `, {
-${indent}                limits: opts.limits?.${field.name}
-${indent}              }`
+      opts = `, DecodeContext.child(decodeContext, '${field.name}', {
+${indent}                limits: opts.limits?.${field.name},
+${indent}                totalLimits: opts.totalLimits?.${field.name}
+${indent}              })`
     } else if (field instanceof ArrayField) {
-      opts = `, {
-${indent}                limits: opts.limits?.${field.name}$
-${indent}              }`
+      opts = `, DecodeContext.child(decodeContext, '${field.name}', {
+${indent}                limits: opts.limits?.${field.name}$,
+${indent}                totalLimits: opts.totalLimits?.${field.name}$
+${indent}              })`
     } else if (field instanceof MapField) {
-      opts = `, {
+      opts = `, DecodeContext.child(decodeContext, '${field.name}', {
 ${indent}                limits: {
 ${indent}                  value: opts.limits?.${field.name}$value
-                }
-${indent}              }`
+                },
+${indent}                totalLimits: opts.totalLimits?.${field.name}$value == null ? undefined : { value: opts.totalLimits.${field.name}$value }
+${indent}              })`
     }
 
     return `${this.jsType.decode}.codec().decode(${reader}, ${reader}.uint32()${opts})`
@@ -198,13 +201,15 @@ ${indent}              }`
     let opts = ''
 
     if (field instanceof MessageField) {
-      opts = `, {
-${indent}              limits: opts.limits?.${field.name}
-${indent}            }`
+      opts = `, DecodeContext.child(decodeContext, '${field.name}', {
+${indent}              limits: opts.limits?.${field.name},
+${indent}              totalLimits: opts.totalLimits?.${field.name}
+${indent}            })`
     } else if (field instanceof ArrayField) {
-      opts = `, {
-${indent}              limits: opts.limits?.${field.name}$
-${indent}            }`
+      opts = `, DecodeContext.child(decodeContext, '${field.name}', {
+${indent}              limits: opts.limits?.${field.name}$,
+${indent}              totalLimits: opts.totalLimits?.${field.name}$
+${indent}            })`
 
       return `for (const evt of ${this.jsType.decode}.codec().stream(${reader}, ${reader}.uint32(), ${prefix}${opts})) {
 ${indent}              yield {
@@ -213,11 +218,12 @@ ${indent}                index: obj.${field.name}
 ${indent}              }
 ${indent}            }`
     } else if (field instanceof MapField) {
-      opts = `, {
+      opts = `, DecodeContext.child(decodeContext, '${field.name}', {
 ${indent}              limits: {
 ${indent}                value: opts.limits?.${field.name}$value
-                }
-${indent}            }`
+${indent}              },
+${indent}              totalLimits: opts.totalLimits?.${field.name}$value == null ? undefined : { value: opts.totalLimits.${field.name}$value }
+${indent}            })`
     }
 
     return `yield * ${this.jsType.decode}.codec().stream(${reader}, ${reader}.uint32(), ${prefix}${opts})`
@@ -302,6 +308,13 @@ export interface ${this.pbType}Input {
     const decodeOpts = this.fields.some(field => field instanceof ArrayField || field instanceof MapField || field instanceof MessageField)
       ? ', opts = {}'
       : ''
+    const decodeContext = decodeOpts === ''
+      ? ''
+      : '\n        const decodeContext = DecodeContext.get(opts)'
+
+    if (decodeContext !== '') {
+      parent.addImport('protons-runtime', 'DecodeContext')
+    }
 
     interfaceCodecDef = `
   let _codec: Codec<${this.pbType}, ${this.pbType}Input>
@@ -316,7 +329,7 @@ export interface ${this.pbType}Input {
         if (opts.lengthDelimited !== false) {
           w.ldelim()
         }
-      }, (r, length${decodeOpts}) => {
+      }, (r, length${decodeOpts}) => {${decodeContext}
         const obj: any = {${this.createDefaultObject()}}
 
         const end = length == null ? r.len : r.pos + length
@@ -333,7 +346,7 @@ export interface ${this.pbType}Input {
         }
 ${enforceOneOfDecoding === '' ? '' : `${enforceOneOfDecoding}\n`}
         return obj
-      }, function * (r, length, prefix${decodeOpts}) {
+      }, function * (r, length, prefix${decodeOpts}) {${decodeContext}
         ${this.createLimitObject()}const end = length == null ? r.len : r.pos + length
 
         if (prefix !== '.') {
@@ -362,7 +375,7 @@ ${enforceOneOfDecoding === '' ? '' : `${enforceOneOfDecoding}\n`}
             message: '${this.jsType.decode}'
           }
         }
-      })
+      }, { totalLimits: true })
     }
 
     return _codec

packages/protons/test/fixtures/basic.ts

@@ -105,7 +105,7 @@ export namespace Basic {
             message: 'Basic'
           }
         }
-      })
+      }, { totalLimits: true })
     }
 
     return _codec
@@ -197,7 +197,7 @@ export namespace Empty {
             message: 'Empty'
           }
         }
-      })
+      }, { totalLimits: true })
     }
 
     return _codec