|
 |
eff8e2 |
From 0472f5b12beb845515b6ef8a339b7223189112fd Mon Sep 17 00:00:00 2001
|
|
|
eff8e2 |
From: Robbie Harwood <rharwood@redhat.com>
|
|
|
eff8e2 |
Date: Thu, 29 Aug 2019 11:13:41 -0400
|
|
|
eff8e2 |
Subject: [PATCH] Always buffer TCP data in __handle_recv()
|
|
|
eff8e2 |
|
|
|
eff8e2 |
Refactor __handle_recv() to always create a BytesIO() object for TCP
|
|
|
eff8e2 |
data. Linearize control flow for ease of debugging. Always apply
|
|
|
eff8e2 |
length checks so that we don't have to wait for EOF in the multiple-recv
|
|
|
eff8e2 |
case.
|
|
|
eff8e2 |
|
|
|
eff8e2 |
Fixes a bug where we wouldn't return any data because we never received
|
|
|
eff8e2 |
the EOF, or didn't receive it fast enough.
|
|
|
eff8e2 |
|
|
|
eff8e2 |
Signed-off-by: Robbie Harwood <rharwood@redhat.com>
|
|
|
eff8e2 |
(cherry picked from commit 7e2b1ab27b843c220fe301b74bab01ed61b0f59a)
|
|
|
eff8e2 |
---
|
|
|
eff8e2 |
kdcproxy/__init__.py | 54 +++++++++++++++++++++++++-------------------
|
|
|
eff8e2 |
1 file changed, 31 insertions(+), 23 deletions(-)
|
|
|
eff8e2 |
|
|
|
eff8e2 |
diff --git a/kdcproxy/__init__.py b/kdcproxy/__init__.py
|
|
|
eff8e2 |
index ab6ed8a..67680c5 100644
|
|
|
eff8e2 |
--- a/kdcproxy/__init__.py
|
|
|
eff8e2 |
+++ b/kdcproxy/__init__.py
|
|
|
eff8e2 |
@@ -128,29 +128,37 @@ class Application:
|
|
|
eff8e2 |
# length prefix. So add it.
|
|
|
eff8e2 |
reply = struct.pack("!I", len(reply)) + reply
|
|
|
eff8e2 |
return reply
|
|
|
eff8e2 |
- else:
|
|
|
eff8e2 |
- # TCP is a different story. The reply must be buffered
|
|
|
eff8e2 |
- # until the full answer is accumulated.
|
|
|
eff8e2 |
- buf = read_buffers.get(sock)
|
|
|
eff8e2 |
- part = sock.recv(1048576)
|
|
|
eff8e2 |
- if buf is None:
|
|
|
eff8e2 |
- if len(part) > 4:
|
|
|
eff8e2 |
- # got enough data in the initial package. Now check
|
|
|
eff8e2 |
- # if we got the full package in the first run.
|
|
|
eff8e2 |
- (length, ) = struct.unpack("!I", part[0:4])
|
|
|
eff8e2 |
- if length + 4 == len(part):
|
|
|
eff8e2 |
- return part
|
|
|
eff8e2 |
- read_buffers[sock] = buf = io.BytesIO()
|
|
|
eff8e2 |
-
|
|
|
eff8e2 |
- if part:
|
|
|
eff8e2 |
- # data received, accumulate it in a buffer
|
|
|
eff8e2 |
- buf.write(part)
|
|
|
eff8e2 |
- return None
|
|
|
eff8e2 |
- else:
|
|
|
eff8e2 |
- # EOF received
|
|
|
eff8e2 |
- read_buffers.pop(sock)
|
|
|
eff8e2 |
- reply = buf.getvalue()
|
|
|
eff8e2 |
- return reply
|
|
|
eff8e2 |
+
|
|
|
eff8e2 |
+ # TCP is a different story. The reply must be buffered until the full
|
|
|
eff8e2 |
+ # answer is accumulated.
|
|
|
eff8e2 |
+ buf = read_buffers.get(sock)
|
|
|
eff8e2 |
+ if buf is None:
|
|
|
eff8e2 |
+ read_buffers[sock] = buf = io.BytesIO()
|
|
|
eff8e2 |
+
|
|
|
eff8e2 |
+ part = sock.recv(1048576)
|
|
|
eff8e2 |
+ if not part:
|
|
|
eff8e2 |
+ # EOF received. Return any incomplete data we have on the theory
|
|
|
eff8e2 |
+ # that a decode error is more apparent than silent failure. The
|
|
|
eff8e2 |
+ # client will fail faster, at least.
|
|
|
eff8e2 |
+ read_buffers.pop(sock)
|
|
|
eff8e2 |
+ reply = buf.getvalue()
|
|
|
eff8e2 |
+ return reply
|
|
|
eff8e2 |
+
|
|
|
eff8e2 |
+ # Data received, accumulate it in a buffer.
|
|
|
eff8e2 |
+ buf.write(part)
|
|
|
eff8e2 |
+
|
|
|
eff8e2 |
+ reply = buf.getvalue()
|
|
|
eff8e2 |
+ if len(reply) < 4:
|
|
|
eff8e2 |
+ # We don't have the length yet.
|
|
|
eff8e2 |
+ return None
|
|
|
eff8e2 |
+
|
|
|
eff8e2 |
+ # Got enough data to check if we have the full package.
|
|
|
eff8e2 |
+ (length, ) = struct.unpack("!I", reply[0:4])
|
|
|
eff8e2 |
+ if length + 4 == len(reply):
|
|
|
eff8e2 |
+ read_buffers.pop(sock)
|
|
|
eff8e2 |
+ return reply
|
|
|
eff8e2 |
+
|
|
|
eff8e2 |
+ return None
|
|
|
eff8e2 |
|
|
|
eff8e2 |
def __filter_addr(self, addr):
|
|
|
eff8e2 |
if addr[0] not in (socket.AF_INET, socket.AF_INET6):
|