fix: bound zero-width collections and roundtrip map entries and polymorphic primitives

README.md

@@ -70,6 +70,9 @@ val bytes = format.encodeToByteArray(payload)
 Fields are written in the order you declare them, with no names or tags. So if
 you change a class, old strings may no longer decode.
 
+The encoded format only stays the same between patch versions. A new minor
+version may change it, so encode stored strings again after upgrading.
+
 ## EncodedFormat
 
 EncodedFormat serializes your object, optionally transforms the bytes, and

src/commonMain/kotlin/com/eignex/kencode/PackedContext.kt

@@ -214,5 +214,8 @@ internal class StructureLayout private constructor(
     }
 }
 
-/** Read cursor shared by an inline decoder and its parent. */
-internal class PackedPosition(var value: Int = 0)
+/** Read state shared by a decoder and its children. */
+internal class PackedPosition(var value: Int = 0) {
+    /** Collection elements declared so far in this decode. */
+    var collectionElements: Long = 0
+}

src/commonMain/kotlin/com/eignex/kencode/PackedDecoder.kt

@@ -10,6 +10,9 @@ import kotlinx.serialization.encoding.Decoder
 import kotlinx.serialization.modules.EmptySerializersModule
 import kotlinx.serialization.modules.SerializersModule
 
+// Elements that need no input bits (Unit, objects, single-entry enums) beyond one per input bit.
+private const val MAX_ZERO_WIDTH_ELEMENTS = 1 shl 20
+
 /** [PackedFormat] decoder. */
 @Suppress("UNCHECKED_CAST")
 @OptIn(ExperimentalSerializationApi::class)
@@ -105,6 +108,11 @@ class PackedDecoder internal constructor(
         collectionSize = size
         collectionElementCount = size.toLong() * elementsPerEntry
         require(collectionElementCount <= Int.MAX_VALUE) { "Map size $size exceeds the element index range" }
+        // Every other element owns at least one input bit, so this bounds what a short input can allocate.
+        cursor.collectionElements += collectionElementCount
+        require(cursor.collectionElements <= 8L * input.size + MAX_ZERO_WIDTH_ELEMENTS) {
+            "Collection size $size exceeds what the input can hold"
+        }
         collectionIndex = 0
         nullBitmapIndex = 0
         var valueCount = size
@@ -137,7 +145,8 @@ class PackedDecoder internal constructor(
     }
 
     // A size larger than the remaining bytes decodes element by element, so a bogus size hits
-    // EOF instead of exhausting memory. Zero-byte elements (Unit, packed booleans) take this path too.
+    // EOF instead of exhausting memory. Zero-byte elements (Unit, packed booleans) take this path too,
+    // bounded by readCollectionHeader.
     override fun decodeSequentially(): Boolean = !layout.isCollection || collectionElementCount <= input.size - position
 
     private val inClass: Boolean get() = inStructure && !layout.isCollection
@@ -172,7 +181,7 @@ class PackedDecoder internal constructor(
 
     override fun decodeString(): String {
         val len = readVarInt()
-        require(len >= 0 && position + len <= input.size) { "Unexpected EOF while decoding String" }
+        require(len >= 0 && len <= input.size - position) { "Unexpected EOF while decoding String" }
         val value = try {
             input.decodeToString(position, position + len, throwOnInvalidSequence = true)
         } catch (e: CharacterCodingException) {
@@ -296,6 +305,9 @@ class PackedDecoder internal constructor(
         val parentTracksNull = layout.isNullableCollection || !layout.isCollection
 
         val value = when {
+            // Polymorphic wrappers pass the subclass serializer directly, bypassing decodeSerializableValue.
+            inClass && descriptor.isContextualElement(index) -> deserializer.deserialize(child())
+
             !isStructural -> deserializer.deserialize(this)
 
             parentTracksNull && childDescriptor.isNullable ->

src/commonMain/kotlin/com/eignex/kencode/PackedEncoder.kt

@@ -4,6 +4,7 @@ import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.SerializationException
 import kotlinx.serialization.SerializationStrategy
 import kotlinx.serialization.descriptors.SerialDescriptor
+import kotlinx.serialization.descriptors.StructureKind
 import kotlinx.serialization.encoding.CompositeEncoder
 import kotlinx.serialization.encoding.Encoder
 import kotlinx.serialization.modules.EmptySerializersModule
@@ -45,6 +46,10 @@ class PackedEncoder internal constructor(
     private lateinit var collectionNullBits: BitOutput
     private lateinit var collectionValueBits: BitOutput
 
+    // A collection opened with beginStructure, not beginCollection, writes its own size at the end.
+    private var writesOwnSize = false
+    private var collectionElementCount = 0
+
     private val dataBuffer = ByteOutput()
 
     private fun child(headerCtx: HeaderContext? = null) =
@@ -67,7 +72,10 @@ class PackedEncoder internal constructor(
 
     override fun beginCollection(descriptor: SerialDescriptor, collectionSize: Int): CompositeEncoder {
         writeVarInt(collectionSize, getBuffer())
-        return child().also { it.initializeStructure(descriptor) }
+        return child().also {
+            it.initializeStructure(descriptor)
+            it.writesOwnSize = false
+        }
     }
 
     private fun initializeStructure(descriptor: SerialDescriptor) {
@@ -87,6 +95,8 @@ class PackedEncoder internal constructor(
         } else {
             collectionNullBits = BitOutput()
             collectionValueBits = BitOutput()
+            writesOwnSize = true
+            collectionElementCount = 0
         }
 
         dataBuffer.reset()
@@ -206,6 +216,10 @@ class PackedEncoder internal constructor(
         if (!inStructure) return
 
         if (layout.isCollection) {
+            if (writesOwnSize) {
+                val perEntry = if (descriptor.kind is StructureKind.MAP) 2 else 1
+                writeVarInt(collectionElementCount / perEntry, output)
+            }
             collectionNullBits.writeTo(output)
             collectionValueBits.writeTo(output)
         } else if (layout.isMergedKind) {
@@ -238,6 +252,7 @@ class PackedEncoder internal constructor(
 
     private fun markEncoded(index: Int) {
         if (inClass && index >= 0) encodedElements[index] = true
+        if (inStructure && layout.isCollection) collectionElementCount = maxOf(collectionElementCount, index + 1)
     }
 
     override fun encodeBooleanElement(descriptor: SerialDescriptor, index: Int, value: Boolean) {
@@ -306,7 +321,8 @@ class PackedEncoder internal constructor(
     ) {
         markEncoded(index)
         currentIndex = index
-        serializer.serialize(this, value)
+        // Polymorphic wrappers pass the subclass serializer directly, bypassing encodeSerializableValue.
+        serializer.serialize(if (inClass && descriptor.isContextualElement(index)) child() else this, value)
         currentIndex = -1
     }
 

src/commonTest/kotlin/com/eignex/kencode/MalformedInputTest.kt

@@ -150,4 +150,23 @@ class MalformedInputTest {
             PackedFormat.decodeFromByteArray(mapSerializer, PackedFormat.encodeToByteArray(mapSerializer, map)),
         )
     }
+
+    @Test
+    fun `huge collections of zero-width elements are rejected`() {
+        assertFailsWith<IllegalArgumentException> {
+            PackedFormat.decodeFromByteArray(ListSerializer(Unit.serializer()), bytes(0xFF, 0xFF, 0xFF, 0xFF, 0x07))
+        }
+        // Nesting must not multiply the limit: 1000 inner lists of 16384 Units each.
+        val nested = bytes(0xE8, 0x07) + ByteArray(3_000) { if (it % 3 == 2) 0x01 else 0x80.toByte() }
+        assertFailsWith<IllegalArgumentException> {
+            PackedFormat.decodeFromByteArray(ListSerializer(ListSerializer(Unit.serializer())), nested)
+        }
+    }
+
+    @Test
+    fun `string lengths that overflow the position are rejected`() {
+        assertFailsWith<IllegalArgumentException> {
+            PackedFormat.decodeFromByteArray(String.serializer(), bytes(0xFF, 0xFF, 0xFF, 0xFF, 0x07, 0x41))
+        }
+    }
 }

src/commonTest/kotlin/com/eignex/kencode/PackedStructureEdgeTest.kt

@@ -0,0 +1,52 @@
+package com.eignex.kencode
+
+import kotlinx.serialization.PolymorphicSerializer
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.builtins.MapEntrySerializer
+import kotlinx.serialization.builtins.serializer
+import kotlinx.serialization.modules.SerializersModule
+import kotlinx.serialization.modules.polymorphic
+import kotlinx.serialization.modules.subclass
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class PackedStructureEdgeTest {
+    @Serializable
+    data class EntryHolder(val entry: Map.Entry<String, Int>, val flag: Boolean)
+
+    @Serializable
+    data class Shape(val sides: Int)
+
+    @Test
+    fun `map entries roundtrip`() {
+        val serializer = MapEntrySerializer(String.serializer(), Int.serializer())
+        val entry = mapOf("a" to 1).entries.single()
+        val decoded = PackedFormat.decodeFromByteArray(serializer, PackedFormat.encodeToByteArray(serializer, entry))
+        assertEquals(entry.key to entry.value, decoded.key to decoded.value)
+
+        val holder = EntryHolder(mapOf("key" to -5).entries.single(), true)
+        val back = EncodedFormat.decodeFromString(
+            EntryHolder.serializer(),
+            EncodedFormat.encodeToString(EntryHolder.serializer(), holder),
+        )
+        assertEquals(holder.entry.key to holder.entry.value, back.entry.key to back.entry.value)
+        assertEquals(holder.flag, back.flag)
+    }
+
+    @Test
+    fun `polymorphic primitive subclasses roundtrip`() {
+        val format = PackedFormat {
+            serializersModule = SerializersModule {
+                polymorphic(Any::class) {
+                    subclass(Boolean::class, Boolean.serializer())
+                    subclass(Int::class, Int.serializer())
+                    subclass(Shape::class)
+                }
+            }
+        }
+        val serializer = PolymorphicSerializer(Any::class)
+        for (value in listOf<Any>(true, false, 42, Shape(3))) {
+            assertEquals(value, format.decodeFromByteArray(serializer, format.encodeToByteArray(serializer, value)))
+        }
+    }
+}