From ae9160df51ecc2edba15659cede1c205f01e29ce Mon Sep 17 00:00:00 2001 From: Jen Basch Date: Tue, 25 Aug 2026 08:28:28 -0700 Subject: [PATCH] Enforce max collection sizes when decoding pkl-binary (#1831) --- .../java/org/pkl/core/PklBinaryDecoder.java | 16 +++++++++ .../org/pkl/core/util/CollectionUtils.java | 15 +++++--- .../pklbinary/AbstractPklBinaryDecoder.java | 36 ++++++++++++++----- .../org/pkl/core/PklBinaryDecoderTest.kt | 31 ++++++++++++++++ 4 files changed, 85 insertions(+), 13 deletions(-) diff --git a/pkl-core/src/main/java/org/pkl/core/PklBinaryDecoder.java b/pkl-core/src/main/java/org/pkl/core/PklBinaryDecoder.java index 9a0e095b2..3e237e899 100644 --- a/pkl-core/src/main/java/org/pkl/core/PklBinaryDecoder.java +++ b/pkl-core/src/main/java/org/pkl/core/PklBinaryDecoder.java @@ -41,16 +41,32 @@ public class PklBinaryDecoder extends AbstractPklBinaryDecoder { super(unpacker); } + private PklBinaryDecoder(MessageUnpacker unpacker, int collectionSizeLimit) { + super(unpacker, collectionSizeLimit); + } + /** Decode a value from the supplied byte array. */ public static Object decode(byte[] bytes) { return new PklBinaryDecoder(MessagePack.newDefaultUnpacker(bytes)).decode(); } + /** Decode a value from the supplied byte array. */ + public static Object decode(byte[] bytes, int collectionSizeLimit) { + return new PklBinaryDecoder(MessagePack.newDefaultUnpacker(bytes), collectionSizeLimit) + .decode(); + } + /** Decode a value from the supplied {@link InputStream}. */ public static Object decode(InputStream inputStream) { return new PklBinaryDecoder(MessagePack.newDefaultUnpacker(inputStream)).decode(); } + /** Decode a value from the supplied {@link InputStream}. */ + public static Object decode(InputStream inputStream, int collectionSizeLimit) { + return new PklBinaryDecoder(MessagePack.newDefaultUnpacker(inputStream), collectionSizeLimit) + .decode(); + } + @Override protected RuntimeException doFail(Exception cause, long offset, List path) { return new RuntimeException( diff --git a/pkl-core/src/main/java/org/pkl/core/util/CollectionUtils.java b/pkl-core/src/main/java/org/pkl/core/util/CollectionUtils.java index 879072dd8..a793f4599 100644 --- a/pkl-core/src/main/java/org/pkl/core/util/CollectionUtils.java +++ b/pkl-core/src/main/java/org/pkl/core/util/CollectionUtils.java @@ -32,27 +32,32 @@ public final class CollectionUtils { @TruffleBoundary public static HashSet newHashSet(int expectedSize) { - return new HashSet<>((int) (expectedSize / LOAD_FACTOR) + 1, LOAD_FACTOR); + return new HashSet<>(calculateHashMapCapacity(expectedSize), LOAD_FACTOR); } @TruffleBoundary public static LinkedHashSet newLinkedHashSet(int expectedSize) { - return new LinkedHashSet<>((int) (expectedSize / LOAD_FACTOR) + 1, LOAD_FACTOR); + return new LinkedHashSet<>(calculateHashMapCapacity(expectedSize), LOAD_FACTOR); } @TruffleBoundary public static HashMap newHashMap(int expectedSize) { - return new HashMap<>((int) (expectedSize / LOAD_FACTOR) + 1, LOAD_FACTOR); + return new HashMap<>(calculateHashMapCapacity(expectedSize), LOAD_FACTOR); } @TruffleBoundary public static LinkedHashMap newLinkedHashMap(int expectedSize) { - return new LinkedHashMap<>((int) (expectedSize / LOAD_FACTOR) + 1, LOAD_FACTOR); + return new LinkedHashMap<>(calculateHashMapCapacity(expectedSize), LOAD_FACTOR); } @TruffleBoundary public static ConcurrentHashMap newConcurrentHashMap(int expectedSize) { - return new ConcurrentHashMap<>((int) (expectedSize / LOAD_FACTOR) + 1, LOAD_FACTOR); + return new ConcurrentHashMap<>(calculateHashMapCapacity(expectedSize), LOAD_FACTOR); + } + + private static int calculateHashMapCapacity(int expectedSize) { + // this is exactly how java.util.HashMap does it + return (int) Math.ceil(expectedSize / (double) LOAD_FACTOR); } } diff --git a/pkl-core/src/main/java/org/pkl/core/util/pklbinary/AbstractPklBinaryDecoder.java b/pkl-core/src/main/java/org/pkl/core/util/pklbinary/AbstractPklBinaryDecoder.java index 5bd94c766..edd1a1faf 100644 --- a/pkl-core/src/main/java/org/pkl/core/util/pklbinary/AbstractPklBinaryDecoder.java +++ b/pkl-core/src/main/java/org/pkl/core/util/pklbinary/AbstractPklBinaryDecoder.java @@ -1,5 +1,5 @@ /* - * Copyright © 2025 Apple Inc. and the Pkl project authors. All rights reserved. + * Copyright © 2025-2026 Apple Inc. and the Pkl project authors. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -40,10 +40,19 @@ import org.pkl.core.util.LateInit; */ public abstract class AbstractPklBinaryDecoder { private final MessageUnpacker unpacker; + private final int collectionSizeLimit; @LateInit protected Deque currPath; + private static final int defaultCollectionSizeLimit = 65535; // 16k + protected AbstractPklBinaryDecoder(MessageUnpacker unpacker) { this.unpacker = unpacker; + this.collectionSizeLimit = defaultCollectionSizeLimit; + } + + protected AbstractPklBinaryDecoder(MessageUnpacker unpacker, int collectionSizeLimit) { + this.unpacker = unpacker; + this.collectionSizeLimit = collectionSizeLimit; } protected static class DecodeException extends RuntimeException { @@ -176,6 +185,13 @@ public abstract class AbstractPklBinaryDecoder { }; } + private void checkCollectionLength(int length, String collectionType) { + if (length <= collectionSizeLimit) return; + throw new DecodeException( + "Unable to decode %s of length %d, exceeded maximum collection size of %d", + collectionType, length, collectionSizeLimit); + } + private Object decodeObject(int len) throws IOException { assertLength(PklBinaryCode.OBJECT, len, 3); currPath.push("'object"); @@ -201,7 +217,7 @@ public abstract class AbstractPklBinaryDecoder { private Object decodeMap(int len) throws IOException { assertLength(PklBinaryCode.MAP, len, 1); currPath.push("'map"); - var result = doDecodeMap(new MapDecodeIterator(unpacker.unpackMapHeader())); + var result = doDecodeMap(new MapDecodeIterator(unpacker.unpackMapHeader(), "map")); unpacker.skipValue(len - 2); currPath.pop(); return result; @@ -210,7 +226,7 @@ public abstract class AbstractPklBinaryDecoder { private Object decodeMapping(int len) throws IOException { assertLength(PklBinaryCode.MAPPING, len, 1); currPath.push("'mapping"); - var result = doDecodeMapping(new MapDecodeIterator(unpacker.unpackMapHeader())); + var result = doDecodeMapping(new MapDecodeIterator(unpacker.unpackMapHeader(), "mapping")); unpacker.skipValue(len - 2); currPath.pop(); return result; @@ -219,7 +235,7 @@ public abstract class AbstractPklBinaryDecoder { private Object decodeList(int len) throws IOException { assertLength(PklBinaryCode.LIST, len, 1); currPath.push("'list"); - var result = doDecodeList(new CollectionDecodeIterator(unpacker.unpackArrayHeader())); + var result = doDecodeList(new CollectionDecodeIterator(unpacker.unpackArrayHeader(), "list")); unpacker.skipValue(len - 2); currPath.pop(); return result; @@ -228,7 +244,8 @@ public abstract class AbstractPklBinaryDecoder { private Object decodeListing(int len) throws IOException { assertLength(PklBinaryCode.LISTING, len, 1); currPath.push("'listing"); - var result = doDecodeListing(new CollectionDecodeIterator(unpacker.unpackArrayHeader())); + var result = + doDecodeListing(new CollectionDecodeIterator(unpacker.unpackArrayHeader(), "listing")); unpacker.skipValue(len - 2); currPath.pop(); return result; @@ -237,7 +254,7 @@ public abstract class AbstractPklBinaryDecoder { private Object decodeSet(int len) throws IOException { assertLength(PklBinaryCode.SET, len, 1); currPath.push("'set"); - var result = doDecodeSet(new CollectionDecodeIterator(unpacker.unpackArrayHeader())); + var result = doDecodeSet(new CollectionDecodeIterator(unpacker.unpackArrayHeader(), "set")); currPath.pop(); unpacker.skipValue(len - 2); return result; @@ -403,6 +420,7 @@ public abstract class AbstractPklBinaryDecoder { protected class ObjectDecodeIterator extends DecodeIterator { ObjectDecodeIterator(int size) { super(size); + checkCollectionLength(size, "object"); } @Override @@ -441,8 +459,9 @@ public abstract class AbstractPklBinaryDecoder { } protected class CollectionDecodeIterator extends DecodeIterator { - CollectionDecodeIterator(int size) { + CollectionDecodeIterator(int size, String collectionType) { super(size); + checkCollectionLength(size, collectionType); } @Override @@ -452,8 +471,9 @@ public abstract class AbstractPklBinaryDecoder { } protected class MapDecodeIterator extends DecodeIterator> { - MapDecodeIterator(int size) { + MapDecodeIterator(int size, String collectionType) { super(size); + checkCollectionLength(size, collectionType); } @Override diff --git a/pkl-core/src/test/kotlin/org/pkl/core/PklBinaryDecoderTest.kt b/pkl-core/src/test/kotlin/org/pkl/core/PklBinaryDecoderTest.kt index c7eaa46f5..664f58395 100644 --- a/pkl-core/src/test/kotlin/org/pkl/core/PklBinaryDecoderTest.kt +++ b/pkl-core/src/test/kotlin/org/pkl/core/PklBinaryDecoderTest.kt @@ -230,4 +230,35 @@ class PklBinaryDecoderTest { "Unexpected blank typealias module URI", ) } + + @Test + fun `decode collections too large`() { + assertExceptionCauseMessageContains( + byteArrayOf(0x93.toByte(), PklBinaryCode.OBJECT.code) + + strFoo + + strFoo + + byteArrayOf(0xdd.toByte(), 0, 1, 0, 0), + "Unable to decode object of length 65536, exceeded maximum collection size", + ) + assertExceptionCauseMessageContains( + byteArrayOf(0x93.toByte(), PklBinaryCode.MAP.code, 0xdf.toByte(), 0, 1, 0, 0), + "Unable to decode map of length 65536, exceeded maximum collection size", + ) + assertExceptionCauseMessageContains( + byteArrayOf(0x93.toByte(), PklBinaryCode.MAPPING.code, 0xdf.toByte(), 0, 1, 0, 0), + "Unable to decode mapping of length 65536, exceeded maximum collection size", + ) + assertExceptionCauseMessageContains( + byteArrayOf(0x93.toByte(), PklBinaryCode.LIST.code, 0xdd.toByte(), 0, 1, 0, 0), + "Unable to decode list of length 65536, exceeded maximum collection size", + ) + assertExceptionCauseMessageContains( + byteArrayOf(0x93.toByte(), PklBinaryCode.LISTING.code, 0xdd.toByte(), 0, 1, 0, 0), + "Unable to decode listing of length 65536, exceeded maximum collection size", + ) + assertExceptionCauseMessageContains( + byteArrayOf(0x93.toByte(), PklBinaryCode.SET.code, 0xdd.toByte(), 0, 1, 0, 0), + "Unable to decode set of length 65536, exceeded maximum collection size", + ) + } }