CVE-2025-59089

@@ -149,6 +149,7 @@ def __await_reply(self, pr, rsocks, wsocks, timeout):
                     if self.sock_type(sock) == socket.SOCK_STREAM:
                         # Remove broken TCP socket from readers
                         rsocks.remove(sock)
+                        read_buffers.pop(sock)
                 else:
                     if reply is not None:
                         return reply
@@ -174,7 +175,7 @@ def __handle_recv(self, sock, read_buffers):
         if self.sock_type(sock) == socket.SOCK_DGRAM:
             # For UDP sockets, recv() returns an entire datagram
             # package. KDC sends one datagram as reply.
-            reply = sock.recv(1048576)
+            reply = sock.recv(self.MAX_LENGTH)
             # If we proxy over UDP, we will be missing the 4-byte
             # length prefix. So add it.
             reply = struct.pack("!I", len(reply)) + reply
@@ -186,30 +187,38 @@ def __handle_recv(self, sock, read_buffers):
         if buf is None:
             read_buffers[sock] = buf = io.BytesIO()
 
-        part = sock.recv(1048576)
-        if not part:
-            # EOF received.  Return any incomplete data we have on the theory
-            # that a decode error is more apparent than silent failure.  The
-            # client will fail faster, at least.
-            read_buffers.pop(sock)
-            reply = buf.getvalue()
-            return reply
+        part = sock.recv(self.MAX_LENGTH)
+        if part:
+            # Data received, accumulate it in a buffer.
+            buf.write(part)
 
-        # Data received, accumulate it in a buffer.
-        buf.write(part)
+            reply = buf.getbuffer()
+            if len(reply) < 4:
+                # We don't have the length yet.
+                return None
 
-        reply = buf.getvalue()
-        if len(reply) < 4:
-            # We don't have the length yet.
-            return None
+            # Got enough data to check if we have the full package.
+            (length, ) = struct.unpack("!I", reply[0:4])
+            length += 4  # add prefix length
 
-        # Got enough data to check if we have the full package.
-        (length, ) = struct.unpack("!I", reply[0:4])
-        if length + 4 == len(reply):
-            read_buffers.pop(sock)
-            return reply
+            if length > self.MAX_LENGTH:
+                raise ValueError('Message length exceeds the maximum length '
+                                 'for a Kerberos message (%i > %i)'
+                                 % (length, self.MAX_LENGTH))
 
-        return None
+            if len(reply) > length:
+                raise ValueError('Message length exceeds its expected length '
+                                 '(%i > %i)' % (len(reply), length))
+
+            if len(reply) < length:
+                return None
+
+        # Else (if part is None), EOF was received.  Return any incomplete data
+        # we have on the theory that a decode error is more apparent than
+        # silent failure.  The client will fail faster, at least.
+
+        read_buffers.pop(sock)
+        return buf.getvalue()
 
     def __filter_addr(self, addr):
         if addr[0] not in (socket.AF_INET, socket.AF_INET6):

@@ -20,6 +20,8 @@
 # THE SOFTWARE.
 
 import os
+import socket
+import struct
 import unittest
 from base64 import b64decode
 try:
@@ -122,6 +124,74 @@ def test_no_server(self):
                                                      kpasswd=True)
         self.assertEqual(response.status_code, 503)
 
+    @mock.patch("socket.getaddrinfo", return_value=addrinfo)
+    @mock.patch("socket.socket")
+    def test_tcp_message_length_exceeds_max(self, m_socket, m_getaddrinfo):
+        # Test that TCP messages with length > MAX_LENGTH raise ValueError
+        # Create a message claiming to be larger than MAX_LENGTH
+        max_len = self.app.MAX_LENGTH
+        # Length prefix claiming message is larger than allowed
+        oversized_length = max_len + 1
+        malicious_msg = struct.pack("!I", oversized_length)
+
+        # Mock socket to return the malicious length prefix
+        mock_sock = m_socket.return_value
+        mock_sock.recv.return_value = malicious_msg
+        mock_sock.getsockopt.return_value = socket.SOCK_STREAM
+
+        # Manually call the receive method to test it
+        read_buffers = {}
+        with self.assertRaises(ValueError) as cm:
+            self.app._Application__handle_recv(mock_sock, read_buffers)
+
+        self.assertIn("exceeds the maximum length", str(cm.exception))
+        self.assertIn(str(max_len), str(cm.exception))
+
+    @mock.patch("socket.getaddrinfo", return_value=addrinfo)
+    @mock.patch("socket.socket")
+    def test_tcp_message_data_exceeds_expected_length(
+        self, m_socket, m_getaddrinfo
+    ):
+        # Test that receiving more data than expected raises ValueError
+        # Create a message with length = 100 but send more data
+        expected_length = 100
+        length_prefix = struct.pack("!I", expected_length)
+        # Send more data than the length prefix indicates
+        extra_data = b"X" * (expected_length + 10)
+        malicious_msg = length_prefix + extra_data
+
+        mock_sock = m_socket.return_value
+        mock_sock.recv.return_value = malicious_msg
+        mock_sock.getsockopt.return_value = socket.SOCK_STREAM
+
+        read_buffers = {}
+        with self.assertRaises(ValueError) as cm:
+            self.app._Application__handle_recv(mock_sock, read_buffers)
+
+        self.assertIn("exceeds its expected length", str(cm.exception))
+
+    @mock.patch("socket.getaddrinfo", return_value=addrinfo)
+    @mock.patch("socket.socket")
+    def test_tcp_eof_returns_buffered_data(self, m_socket, m_getaddrinfo):
+        # Test that EOF returns any buffered data
+        initial_data = b"\x00\x00\x00\x10"  # Length = 16
+        mock_sock = m_socket.return_value
+        mock_sock.getsockopt.return_value = socket.SOCK_STREAM
+
+        # First recv returns some data, second returns empty (EOF)
+        mock_sock.recv.side_effect = [initial_data, b""]
+
+        read_buffers = {}
+        # First call buffers the data
+        result = self.app._Application__handle_recv(mock_sock, read_buffers)
+        self.assertIsNone(result)  # Not complete yet
+
+        # Second call gets EOF and returns buffered data
+        result = self.app._Application__handle_recv(mock_sock, read_buffers)
+        self.assertEqual(result, initial_data)
+        # Buffer should be cleaned up
+        self.assertNotIn(mock_sock, read_buffers)
+
 
 def decode(data):
     data = data.replace(b'\\n', b'')