Improve tests of types.py

minecraft/networking/types.py

@@ -83,14 +83,23 @@ class VarInt(Type):
     @staticmethod
     def read(file_object):
         number = 0
-        for i in range(5):
+        # Limit of 5 bytes, otherwise its possible to cause
+        # a DOS attack by sending VarInts that just keep
+        # going
+        bytes_encountered = 0
+        while True:
             byte = file_object.read(1)
             if len(byte) < 1:
                 raise EOFError("Unexpected end of message.")
+
             byte = ord(byte)
-            number |= (byte & 0x7F) << 7 * i
+            number |= (byte & 0x7F) << 7 * bytes_encountered
             if not byte & 0x80:
                 break
+
+            bytes_encountered += 1
+            if bytes_encountered > 5:
+                raise ValueError("Tried to read too long of a VarInt")
         return number
 
     @staticmethod
@@ -109,6 +118,7 @@ class VarInt(Type):
         for max_value, size in VARINT_SIZE_TABLE.items():
             if value < max_value:
                 return size
+        raise ValueError("Integer too large")
 
 
 # Maps (maximum integer value -> size of VarInt in bytes)
@@ -199,3 +209,7 @@ class UUID(Type):
     @staticmethod
     def read(file_object):
         return str(uuid.UUID(bytes=file_object.read(16)))
+
+    @staticmethod
+    def send(value, socket):
+        socket.send(uuid.UUID(value).bytes)

tests/test_packets.py

@@ -11,6 +11,30 @@ from minecraft.networking.packets import (
     PacketBuffer, ChatPacket, KeepAlivePacket, PacketListener)
 
 
+class PacketBufferTest(unittest.TestCase):
+    def test_basic_read_write(self):
+        message = b"hello"
+
+        packet_buffer = PacketBuffer()
+        packet_buffer.send(message)
+
+        packet_buffer.reset_cursor()
+        self.assertEqual(packet_buffer.read(), message)
+        packet_buffer.reset_cursor()
+        self.assertEqual(packet_buffer.recv(), message)
+
+        packet_buffer.reset()
+        self.assertNotEqual(packet_buffer.read(), message)
+
+    def test_get_writable(self):
+        message = b"hello"
+
+        packet_buffer = PacketBuffer()
+        packet_buffer.send(message)
+
+        self.assertEqual(packet_buffer.get_writable(), message)
+
+
 class PacketSerializatonTest(unittest.TestCase):
 
     def test_packet(self):

tests/test_serialization.py

@@ -1,10 +1,11 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 import unittest
+import uuid
 from minecraft.networking.types import (
     Type, Boolean, UnsignedByte, Byte, Short, UnsignedShort,
     Integer, VarInt, Long, Float, Double, ShortPrefixedByteArray,
-    VarIntPrefixedByteArray, String as StringType
+    VarIntPrefixedByteArray, UUID, String as StringType
 )
 from minecraft.networking.packets import PacketBuffer
 
@@ -22,6 +23,7 @@ TEST_DATA = {
     Double: [36.004002],
     ShortPrefixedByteArray: [bytes(245)],
     VarIntPrefixedByteArray: [bytes(1234)],
+    UUID: ["12345678-1234-5678-1234-567812345678"],
     StringType: ["hello world"]
 }
 
@@ -60,6 +62,15 @@ class SerializationTest(unittest.TestCase):
         self.assertEqual(VarInt.size(2), 1)
         self.assertEqual(VarInt.size(1250), 2)
 
+        with self.assertRaises(ValueError):
+            VarInt.size(2 ** 90)
+
+        with self.assertRaises(ValueError):
+            packet_buffer = PacketBuffer()
+            VarInt.send(2 ** 49, packet_buffer)
+            packet_buffer.reset_cursor()
+            VarInt.read(packet_buffer)
+
         packet_buffer = PacketBuffer()
         VarInt.send(50000, packet_buffer)
         packet_buffer.reset_cursor()