Skip to content

Commit 3e35e57

Browse files
BFergersonBFergerson
authored
fix: limit large arrays/lists/maps (#107)
* fix: limit large arrays/lists/maps * fix: handle large primitive arrays * fix: only write exceeded on more than max array size * fix: handle large maps * fix: copy map/set/collection to avoid CME * fix: serialize nulls * refactor: MAX_COLLECTION_SIZE_EXCEEDED fix: support sets
1 parent e1da1a1 commit 3e35e57

File tree

3 files changed

+128
-18
lines changed

3 files changed

+128
-18
lines changed

‎services/src/main/kotlin/spp/probe/services/common/ContextReceiver.kt

+10-1
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ import spp.protocol.instrument.meter.MetricValueType
3737
import java.util.*
3838
import java.util.concurrent.ConcurrentHashMap
3939
import java.util.concurrent.Executors
40+
import kotlin.collections.HashSet
41+
import kotlin.collections.LinkedHashSet
4042

4143
@Suppress("unused")
4244
object ContextReceiver {
@@ -80,7 +82,14 @@ object ContextReceiver {
8082
instrumentId: String, key: String, value: Any?, type: String,
8183
variableMap: MutableMap<String?, MutableMap<String, Pair<String, Any?>>>
8284
) {
83-
variableMap.computeIfAbsent(instrumentId) { HashMap() }[key] = Pair(type, value)
85+
var copyValue = value
86+
when (copyValue) {
87+
is Map<*, *> -> copyValue = copyValue.toMap()
88+
is LinkedHashSet<*> -> copyValue = LinkedHashSet(copyValue)
89+
is HashSet<*> -> copyValue = HashSet(copyValue)
90+
is Collection<*> -> copyValue = copyValue.toList()
91+
}
92+
variableMap.computeIfAbsent(instrumentId) { HashMap() }[key] = Pair(type, copyValue)
8493
}
8594

8695
@JvmStatic

‎services/src/main/kotlin/spp/probe/services/common/ModelSerializer.kt

+1
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ enum class ModelSerializer {
3939

4040
private val gson: Gson = GsonBuilder().disableHtmlEscaping().create()
4141
val extendedGson: Gson = GsonBuilder()
42+
.serializeNulls()
4243
.setJsogPolicy(JsogPolicy.DEFAULT.withJsogAlwaysEnabled())
4344
.registerTypeAdapterFactory(CappedTypeAdapterFactory(5))
4445
.registerTypeAdapterFactory(object : TypeAdapterFactory {

‎services/src/main/kotlin/spp/probe/services/common/serialize/CappedTypeAdapterFactory.kt

+117-17
Original file line numberDiff line numberDiff line change
@@ -58,24 +58,57 @@ class CappedTypeAdapterFactory(val maxDepth: Int) : TypeAdapterFactory {
5858
val objSize = instrumentation!!.getObjectSize(value)
5959
if (objSize <= maxMemorySize) {
6060
JsogRegistry.get().userData["depth"] = (JsogRegistry.get().userData["depth"] as Int) + 1
61-
try {
62-
ModelSerializer.INSTANCE.extendedGson.getDelegateAdapter(
63-
this@CappedTypeAdapterFactory, type
64-
).write(jsonWriter, value)
65-
} catch (e: Exception) {
66-
jsonWriter.beginObject()
67-
jsonWriter.name("@skip")
68-
jsonWriter.value("EXCEPTION_OCCURRED")
69-
jsonWriter.name("@class")
70-
jsonWriter.value(value.javaClass.name)
71-
jsonWriter.name("@size")
72-
jsonWriter.value(objSize.toString())
73-
jsonWriter.name("@cause")
74-
jsonWriter.value(e.message)
75-
jsonWriter.name("@id")
76-
jsonWriter.value(Integer.toHexString(System.identityHashCode(value)))
77-
jsonWriter.endObject()
61+
62+
if (value is Collection<*>) {
63+
writeCollection(jsonWriter, value.iterator(), value.size, objSize)
64+
} else if (value is Map<*, *> && value.size > maxArraySize) {
65+
jsonWriter.beginArray()
66+
value.onEachIndexed { i, entry ->
67+
if (i >= maxArraySize) return@onEachIndexed
68+
jsonWriter.beginObject()
69+
jsonWriter.name(entry.key.toString())
70+
if (entry.value == null) {
71+
jsonWriter.nullValue()
72+
} else {
73+
when (entry.value) {
74+
is Boolean -> jsonWriter.value(entry.value as Boolean)
75+
is Number -> jsonWriter.value(entry.value as Number)
76+
is Char -> jsonWriter.value(entry.value.toString())
77+
is String -> jsonWriter.value(entry.value as String)
78+
else -> doWrite(jsonWriter, entry.value, objSize)
79+
}
80+
}
81+
jsonWriter.endObject()
82+
}
83+
84+
if (value.size > maxArraySize) {
85+
jsonWriter.beginObject()
86+
jsonWriter.name("@skip")
87+
jsonWriter.value("MAX_COLLECTION_SIZE_EXCEEDED")
88+
jsonWriter.name("@skip[size]")
89+
jsonWriter.value(value.size)
90+
jsonWriter.name("@skip[max]")
91+
jsonWriter.value(maxArraySize)
92+
jsonWriter.endObject()
93+
}
94+
jsonWriter.endArray()
95+
} else if (value!!::class.java.isArray) {
96+
when (value) {
97+
is BooleanArray -> writeCollection(jsonWriter, value.iterator(), value.size, objSize)
98+
is ByteArray -> writeCollection(jsonWriter, value.iterator(), value.size, objSize)
99+
is CharArray -> writeCollection(jsonWriter, value.iterator(), value.size, objSize)
100+
is ShortArray -> writeCollection(jsonWriter, value.iterator(), value.size, objSize)
101+
is IntArray -> writeCollection(jsonWriter, value.iterator(), value.size, objSize)
102+
is LongArray -> writeCollection(jsonWriter, value.iterator(), value.size, objSize)
103+
is FloatArray -> writeCollection(jsonWriter, value.iterator(), value.size, objSize)
104+
is DoubleArray -> writeCollection(jsonWriter, value.iterator(), value.size, objSize)
105+
is Array<*> -> writeCollection(jsonWriter, value.iterator(), value.size, objSize)
106+
else -> throw IllegalArgumentException("Unsupported array type: " + value.javaClass.name)
107+
}
108+
} else {
109+
doWrite(jsonWriter, value as T, value.javaClass as Class<T>, objSize)
78110
}
111+
79112
JsogRegistry.get().userData["depth"] = (JsogRegistry.get().userData["depth"] as Int) - 1
80113
} else {
81114
jsonWriter.beginObject()
@@ -91,13 +124,80 @@ class CappedTypeAdapterFactory(val maxDepth: Int) : TypeAdapterFactory {
91124
}
92125
}
93126

127+
private fun writeCollection(jsonWriter: JsonWriter, value: Iterator<*>, arrSize: Int, objSize: Long) {
128+
jsonWriter.beginArray()
129+
value.withIndex().forEach { (i, it) ->
130+
if (i >= maxArraySize) return@forEach
131+
if (it == null) {
132+
jsonWriter.nullValue()
133+
} else {
134+
doWrite(jsonWriter, it as T, it::class.java as Class<T>, objSize)
135+
}
136+
}
137+
138+
if (arrSize > maxArraySize) {
139+
jsonWriter.beginObject()
140+
jsonWriter.name("@skip")
141+
jsonWriter.value("MAX_COLLECTION_SIZE_EXCEEDED")
142+
jsonWriter.name("@skip[size]")
143+
jsonWriter.value(arrSize)
144+
jsonWriter.name("@skip[max]")
145+
jsonWriter.value(maxArraySize)
146+
jsonWriter.endObject()
147+
}
148+
jsonWriter.endArray()
149+
}
150+
151+
private fun doWrite(jsonWriter: JsonWriter, value: Any?, objSize: Long) {
152+
try {
153+
ModelSerializer.INSTANCE.extendedGson.getDelegateAdapter(
154+
this@CappedTypeAdapterFactory, TypeToken.get(Any::class.java)
155+
).write(jsonWriter, value)
156+
} catch (e: Exception) {
157+
jsonWriter.beginObject()
158+
jsonWriter.name("@skip")
159+
jsonWriter.value("EXCEPTION_OCCURRED")
160+
jsonWriter.name("@class")
161+
jsonWriter.value(value!!::class.java.name)
162+
jsonWriter.name("@size")
163+
jsonWriter.value(objSize.toString())
164+
jsonWriter.name("@cause")
165+
jsonWriter.value(e.message)
166+
jsonWriter.name("@id")
167+
jsonWriter.value(Integer.toHexString(System.identityHashCode(value)))
168+
jsonWriter.endObject()
169+
}
170+
}
171+
172+
private fun <T> doWrite(jsonWriter: JsonWriter, value: T?, javaClazz: Class<T>, objSize: Long) {
173+
try {
174+
ModelSerializer.INSTANCE.extendedGson.getDelegateAdapter(
175+
this@CappedTypeAdapterFactory, TypeToken.get(javaClazz)
176+
).write(jsonWriter, value)
177+
} catch (e: Exception) {
178+
jsonWriter.beginObject()
179+
jsonWriter.name("@skip")
180+
jsonWriter.value("EXCEPTION_OCCURRED")
181+
jsonWriter.name("@class")
182+
jsonWriter.value(javaClazz.name)
183+
jsonWriter.name("@size")
184+
jsonWriter.value(objSize.toString())
185+
jsonWriter.name("@cause")
186+
jsonWriter.value(e.message)
187+
jsonWriter.name("@id")
188+
jsonWriter.value(Integer.toHexString(System.identityHashCode(value)))
189+
jsonWriter.endObject()
190+
}
191+
}
192+
94193
override fun read(jsonReader: JsonReader): T? = null
95194
}
96195
}
97196

98197
companion object {
99198
private var instrumentation: Instrumentation? = null
100199
private var maxMemorySize: Long = -1
200+
private var maxArraySize: Int = 100
101201

102202
@JvmStatic
103203
fun setInstrumentation(instrumentation: Instrumentation) {

0 commit comments

Comments
 (0)