Skip to content

Commit 9491a0c

Browse files
CopilotCopilot
authored
Fix malformed batch queue growth vulnerability in Unbatcher (#4113)
* fix: handle malformed batches in Unbatcher to prevent queue growth attack - Add Unbatcher.Clear() to drain queued batches and return writers to pool - Make GetNextMessage() throw when message size prefix exceeds remaining bytes, clearing all batches before throwing - Wrap batch processing loops in NetworkServer and NetworkClient with try-catch to handle malformed batch exceptions and disconnect - Disconnect (not just log) when batches remain after processing loop - Add tests for malformed batch detection and Clear() * refactor: use Array.Empty<byte>() instead of new byte[0] in Unbatcher.Clear * fix: use explicit ArraySegment<byte> in Unbatcher.Clear for cross-version Unity compatibility * Use #if UNITY_2021_3_OR_NEWER for implicit ArraySegment conversion in Unbatcher.Clear() --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
1 parent 97bae68 commit 9491a0c

4 files changed

Lines changed: 203 additions & 66 deletions

File tree

‎Assets/Mirror/Core/Batching/Unbatcher.cs‎

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,23 @@ public class Unbatcher
1818

1919
public int BatchesCount => batches.Count;
2020

21+
// clear all batches and return writers to pool.
22+
// called when a malformed batch is detected or connection disconnects.
23+
public void Clear()
24+
{
25+
while (batches.Count > 0)
26+
{
27+
NetworkWriterPooled writer = batches.Dequeue();
28+
NetworkWriterPool.Return(writer);
29+
}
30+
// reset reader so it doesn't point to returned batch
31+
#if UNITY_2021_3_OR_NEWER
32+
reader.SetBuffer(Array.Empty<byte>());
33+
#else
34+
reader.SetBuffer(new ArraySegment<byte>(Array.Empty<byte>()));
35+
#endif
36+
}
37+
2138
// NetworkReader is only created once,
2239
// then pointed to the first batch.
2340
readonly NetworkReader reader = new NetworkReader(new byte[0]);
@@ -117,9 +134,14 @@ public bool GetNextMessage(out ArraySegment<byte> message, out double remoteTime
117134
// see Batcher.AddMessage comments for explanation.
118135
int size = (int)Compression.DecompressVarUInt(reader);
119136

120-
// validate size prefix, in case attackers send malicious data
137+
// validate size prefix, in case attackers send malicious data.
138+
// a size larger than remaining bytes means the batch is malformed.
139+
// clear all batches and throw so the caller can disconnect.
121140
if (reader.Remaining < size)
122-
return false;
141+
{
142+
Clear();
143+
throw new InvalidOperationException($"GetNextMessage: malformed batch with message size {size} > remaining {reader.Remaining}. All batches cleared.");
144+
}
123145

124146
// return the message of size
125147
message = reader.ReadBytesSegment(size);

‎Assets/Mirror/Core/NetworkClient.cs‎

Lines changed: 51 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -354,53 +354,68 @@ internal static void OnTransportData(ArraySegment<byte> data, int channelId)
354354
// would only be processed when OnTransportData is called
355355
// the next time.
356356
// => consider moving processing to NetworkEarlyUpdate.
357-
while (!isLoadingScene &&
358-
unbatcher.GetNextMessage(out ArraySegment<byte> message, out double remoteTimestamp))
357+
//
358+
// GetNextMessage may throw for malformed batches
359+
// (invalid varint, message size > remaining bytes).
360+
// catch and disconnect to prevent queue growth attacks.
361+
try
359362
{
360-
using (NetworkReaderPooled reader = NetworkReaderPool.Get(message))
363+
while (!isLoadingScene &&
364+
unbatcher.GetNextMessage(out ArraySegment<byte> message, out double remoteTimestamp))
361365
{
362-
// enough to read at least header size?
363-
if (reader.Remaining >= NetworkMessages.IdSize)
366+
using (NetworkReaderPooled reader = NetworkReaderPool.Get(message))
364367
{
365-
// make remoteTimeStamp available to the user
366-
connection.remoteTimeStamp = remoteTimestamp;
368+
// enough to read at least header size?
369+
if (reader.Remaining >= NetworkMessages.IdSize)
370+
{
371+
// make remoteTimeStamp available to the user
372+
connection.remoteTimeStamp = remoteTimestamp;
367373

368-
// handle message
369-
if (!UnpackAndInvoke(reader, channelId))
374+
// handle message
375+
if (!UnpackAndInvoke(reader, channelId))
376+
{
377+
// warn, disconnect and return if failed
378+
// -> warning because attackers might send random data
379+
// -> messages in a batch aren't length prefixed.
380+
// failing to read one would cause undefined
381+
// behaviour for every message afterwards.
382+
// so we need to disconnect.
383+
// -> return to avoid the below unbatches.count error.
384+
// we already disconnected and handled it.
385+
if (exceptionsDisconnect)
386+
{
387+
Debug.LogError($"NetworkClient: failed to unpack and invoke message. Disconnecting.");
388+
connection.Disconnect();
389+
}
390+
else
391+
Debug.LogWarning($"NetworkClient: failed to unpack and invoke message.");
392+
393+
return;
394+
}
395+
}
396+
// otherwise disconnect
397+
else
370398
{
371-
// warn, disconnect and return if failed
372-
// -> warning because attackers might send random data
373-
// -> messages in a batch aren't length prefixed.
374-
// failing to read one would cause undefined
375-
// behaviour for every message afterwards.
376-
// so we need to disconnect.
377-
// -> return to avoid the below unbatches.count error.
378-
// we already disconnected and handled it.
379399
if (exceptionsDisconnect)
380400
{
381-
Debug.LogError($"NetworkClient: failed to unpack and invoke message. Disconnecting.");
401+
Debug.LogError($"NetworkClient: received Message was too short (messages should start with message id). Disconnecting.");
382402
connection.Disconnect();
383403
}
384404
else
385-
Debug.LogWarning($"NetworkClient: failed to unpack and invoke message.");
386-
405+
Debug.LogWarning("NetworkClient: received Message was too short (messages should start with message id)");
387406
return;
388407
}
389408
}
390-
// otherwise disconnect
391-
else
392-
{
393-
if (exceptionsDisconnect)
394-
{
395-
Debug.LogError($"NetworkClient: received Message was too short (messages should start with message id). Disconnecting.");
396-
connection.Disconnect();
397-
}
398-
else
399-
Debug.LogWarning("NetworkClient: received Message was too short (messages should start with message id)");
400-
return;
401-
}
402409
}
403410
}
411+
catch (Exception e)
412+
{
413+
// malformed batch: invalid varint, message size > remaining, etc.
414+
// unbatcher already cleared batches when it detected the error.
415+
Debug.LogError($"NetworkClient: failed to parse batch: {e.Message}. Disconnecting.");
416+
connection.Disconnect();
417+
return;
418+
}
404419

405420
// if we weren't interrupted by a scene change,
406421
// then all batched messages should have been processed now.
@@ -421,6 +436,10 @@ internal static void OnTransportData(ArraySegment<byte> data, int channelId)
421436
if (!isLoadingScene && unbatcher.BatchesCount > 0)
422437
{
423438
Debug.LogError($"Still had {unbatcher.BatchesCount} batches remaining after processing, even though processing was not interrupted by a scene change. This should never happen, as it would cause ever growing batches.\nPossible reasons:\n* A message didn't deserialize as much as it serialized\n*There was no message handler for a message id, so the reader wasn't read until the end.");
439+
440+
// disconnect and clear to prevent memory leak / queue growth
441+
unbatcher.Clear();
442+
connection.Disconnect();
424443
}
425444
}
426445
else Debug.LogError("Skipped Data message handling because connection is null.");

‎Assets/Mirror/Core/NetworkServer.cs‎

Lines changed: 51 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -951,54 +951,69 @@ internal static void OnTransportData(int connectionId, ArraySegment<byte> data,
951951
// would only be processed when OnTransportData is called
952952
// the next time.
953953
// => consider moving processing to NetworkEarlyUpdate.
954-
while (!isLoadingScene &&
955-
connection.unbatcher.GetNextMessage(out ArraySegment<byte> message, out double remoteTimestamp))
954+
//
955+
// GetNextMessage may throw for malformed batches
956+
// (invalid varint, message size > remaining bytes).
957+
// catch and disconnect to prevent queue growth attacks.
958+
try
956959
{
957-
using (NetworkReaderPooled reader = NetworkReaderPool.Get(message))
960+
while (!isLoadingScene &&
961+
connection.unbatcher.GetNextMessage(out ArraySegment<byte> message, out double remoteTimestamp))
958962
{
959-
// enough to read at least header size?
960-
if (reader.Remaining >= NetworkMessages.IdSize)
963+
using (NetworkReaderPooled reader = NetworkReaderPool.Get(message))
961964
{
962-
// make remoteTimeStamp available to the user
963-
connection.remoteTimeStamp = remoteTimestamp;
965+
// enough to read at least header size?
966+
if (reader.Remaining >= NetworkMessages.IdSize)
967+
{
968+
// make remoteTimeStamp available to the user
969+
connection.remoteTimeStamp = remoteTimestamp;
964970

965-
// handle message
966-
if (!UnpackAndInvoke(connection, reader, channelId))
971+
// handle message
972+
if (!UnpackAndInvoke(connection, reader, channelId))
973+
{
974+
// warn, disconnect and return if failed
975+
// -> warning because attackers might send random data
976+
// -> messages in a batch aren't length prefixed.
977+
// failing to read one would cause undefined
978+
// behaviour for every message afterwards.
979+
// so we need to disconnect.
980+
// -> return to avoid the below unbatches.count error.
981+
// we already disconnected and handled it.
982+
if (exceptionsDisconnect)
983+
{
984+
Debug.LogError($"NetworkServer: failed to unpack and invoke message. Disconnecting {connectionId}.");
985+
connection.Disconnect();
986+
}
987+
else
988+
Debug.LogWarning($"NetworkServer: failed to unpack and invoke message from connectionId:{connectionId}.");
989+
990+
return;
991+
}
992+
}
993+
// otherwise disconnect
994+
else
967995
{
968-
// warn, disconnect and return if failed
969-
// -> warning because attackers might send random data
970-
// -> messages in a batch aren't length prefixed.
971-
// failing to read one would cause undefined
972-
// behaviour for every message afterwards.
973-
// so we need to disconnect.
974-
// -> return to avoid the below unbatches.count error.
975-
// we already disconnected and handled it.
976996
if (exceptionsDisconnect)
977997
{
978-
Debug.LogError($"NetworkServer: failed to unpack and invoke message. Disconnecting {connectionId}.");
998+
Debug.LogError($"NetworkServer: received message from connectionId:{connectionId} was too short (messages should start with message id). Disconnecting.");
979999
connection.Disconnect();
9801000
}
9811001
else
982-
Debug.LogWarning($"NetworkServer: failed to unpack and invoke message from connectionId:{connectionId}.");
1002+
Debug.LogWarning($"NetworkServer: received message from connectionId:{connectionId} was too short (messages should start with message id).");
9831003

9841004
return;
9851005
}
9861006
}
987-
// otherwise disconnect
988-
else
989-
{
990-
if (exceptionsDisconnect)
991-
{
992-
Debug.LogError($"NetworkServer: received message from connectionId:{connectionId} was too short (messages should start with message id). Disconnecting.");
993-
connection.Disconnect();
994-
}
995-
else
996-
Debug.LogWarning($"NetworkServer: received message from connectionId:{connectionId} was too short (messages should start with message id).");
997-
998-
return;
999-
}
10001007
}
10011008
}
1009+
catch (Exception e)
1010+
{
1011+
// malformed batch: invalid varint, message size > remaining, etc.
1012+
// unbatcher already cleared batches when it detected the error.
1013+
Debug.LogError($"NetworkServer: failed to parse batch from connectionId:{connectionId}: {e.Message}. Disconnecting.");
1014+
connection.Disconnect();
1015+
return;
1016+
}
10021017

10031018
// if we weren't interrupted by a scene change,
10041019
// then all batched messages should have been processed now.
@@ -1018,6 +1033,10 @@ internal static void OnTransportData(int connectionId, ArraySegment<byte> data,
10181033
if (!isLoadingScene && connection.unbatcher.BatchesCount > 0)
10191034
{
10201035
Debug.LogError($"Still had {connection.unbatcher.BatchesCount} batches remaining after processing, even though processing was not interrupted by a scene change. This should never happen, as it would cause ever growing batches.\nPossible reasons:\n* A message didn't deserialize as much as it serialized\n*There was no message handler for a message id, so the reader wasn't read until the end.");
1036+
1037+
// disconnect and clear to prevent memory leak / queue growth
1038+
connection.unbatcher.Clear();
1039+
connection.Disconnect();
10211040
}
10221041
}
10231042
else Debug.LogError($"HandleData Unknown connectionId:{connectionId}");

‎Assets/Mirror/Tests/Editor/Batching/UnbatcherTests.cs‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,5 +133,82 @@ public void RetireBatchAndTryNewBatch()
133133
Assert.That(reader.ReadByte(), Is.EqualTo(0x04));
134134
Assert.That(remoteTimeStamp, Is.EqualTo(TimeStamp + 1));
135135
}
136+
137+
// malformed batch: message size prefix larger than remaining bytes.
138+
// GetNextMessage should throw and clear all batches.
139+
[Test]
140+
public void GetNextMessage_MalformedBatch_SizeLargerThanRemaining()
141+
{
142+
// craft a batch with timestamp + varint size prefix claiming 1,000,000 bytes
143+
// but only a few bytes actually remaining
144+
NetworkWriter writer = new NetworkWriter();
145+
writer.WriteDouble(TimeStamp);
146+
Compression.CompressVarUInt(writer, 1000000); // claim 1M bytes
147+
writer.WriteByte(0xFF); // but only 1 byte of data
148+
byte[] batch = writer.ToArray();
149+
150+
unbatcher.AddBatch(new ArraySegment<byte>(batch));
151+
Assert.That(unbatcher.BatchesCount, Is.EqualTo(1));
152+
153+
// GetNextMessage should throw InvalidOperationException
154+
Assert.Throws<InvalidOperationException>(() =>
155+
{
156+
unbatcher.GetNextMessage(out _, out _);
157+
});
158+
159+
// all batches should be cleared
160+
Assert.That(unbatcher.BatchesCount, Is.EqualTo(0));
161+
}
162+
163+
// malformed batch should not allow queue growth when repeated
164+
[Test]
165+
public void GetNextMessage_MalformedBatch_NoQueueGrowth()
166+
{
167+
// add multiple malformed batches
168+
for (int i = 0; i < 10; i++)
169+
{
170+
NetworkWriter writer = new NetworkWriter();
171+
writer.WriteDouble(TimeStamp);
172+
Compression.CompressVarUInt(writer, 1000000);
173+
writer.WriteByte(0xFF);
174+
byte[] batch = writer.ToArray();
175+
176+
unbatcher.AddBatch(new ArraySegment<byte>(batch));
177+
178+
// each GetNextMessage should throw and clear
179+
Assert.Throws<InvalidOperationException>(() =>
180+
{
181+
unbatcher.GetNextMessage(out _, out _);
182+
});
183+
184+
// batches should be cleared after each malformed batch
185+
Assert.That(unbatcher.BatchesCount, Is.EqualTo(0));
186+
}
187+
}
188+
189+
[Test]
190+
public void Clear()
191+
{
192+
// add batches
193+
byte[] batch1 = BatcherTests.MakeBatch(TimeStamp, new byte[] {0x01, 0x02});
194+
byte[] batch2 = BatcherTests.MakeBatch(TimeStamp + 1, new byte[] {0x03, 0x04});
195+
unbatcher.AddBatch(new ArraySegment<byte>(batch1));
196+
unbatcher.AddBatch(new ArraySegment<byte>(batch2));
197+
Assert.That(unbatcher.BatchesCount, Is.EqualTo(2));
198+
199+
// clear should remove all batches
200+
unbatcher.Clear();
201+
Assert.That(unbatcher.BatchesCount, Is.EqualTo(0));
202+
203+
// should be able to add and read new batches after clear
204+
byte[] batch3 = BatcherTests.MakeBatch(TimeStamp + 2, new byte[] {0x05, 0x06});
205+
unbatcher.AddBatch(new ArraySegment<byte>(batch3));
206+
bool result = unbatcher.GetNextMessage(out ArraySegment<byte> message, out double remoteTimeStamp);
207+
Assert.That(result, Is.True);
208+
NetworkReader reader = new NetworkReader(message);
209+
Assert.That(reader.ReadByte(), Is.EqualTo(0x05));
210+
Assert.That(reader.ReadByte(), Is.EqualTo(0x06));
211+
Assert.That(remoteTimeStamp, Is.EqualTo(TimeStamp + 2));
212+
}
136213
}
137214
}

0 commit comments

Comments
 (0)