From 5f0a735a1df109ea601388c27d7fa892b081f0af Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Juhani=20Krekel=C3=A4?= <juhani@krekelä.fi>
Date: Sun, 17 Oct 2021 18:48:08 +0300
Subject: [PATCH] Keep transmitting request body even after remote server
 answers

---
 untls_proxy.py | 46 +++++++++++++++++++++++++++++++++++-----------
 1 file changed, 35 insertions(+), 11 deletions(-)

diff --git a/untls_proxy.py b/untls_proxy.py
index 41e63d6..c6d847e 100755
--- a/untls_proxy.py
+++ b/untls_proxy.py
@@ -290,7 +290,7 @@ def proxy(sock, host):
 		poll = select.poll()
 		poll.register(remote_sock, select.POLLIN)
 		sock.settimeout(1)
-		while len(poll.poll(0)) == 0:
+		while len(poll.poll(1000)) == 0:
 			try:
 				data = sock.recv(1024)
 			except ConnectionResetError:
@@ -379,22 +379,46 @@ def proxy(sock, host):
 	del request_data
 
 	# TODO: Un-https links
-	# TODO: Keep sending request body, if any
 	print('', file=sys.stderr)
 	sock.settimeout(60)
 	remote_sock.settimeout(60)
+	poll.register(sock, select.POLLIN)
+	last_transfer = time.monotonic()
 	while True:
-		try:
-			data = remote_sock.recv(1024)
-		except (ConnectionResetError, socket.timeout):
+		events = poll.poll(60_000)
+		if len(events) == 0 and time.monotonic() - last_transfer > 60:
+			remote_sock.close()
 			return
-		if data == b'': break
-		try:
-			sock.sendall(data)
-		except (ConnectionResetError, BrokenPipeError, socket.timeout):
-			break
 
-	remote_sock.close()
+		for fd, _ in events:
+			if fd == sock.fileno():
+				try:
+					data = sock.recv(1024)
+				except (ConnectionResetError):
+					return
+				if data != b'':
+					try:
+						remote_sock.sendall(data)
+					except (ConnectionResetError, BrokenPipeError):
+						return
+					except socket.timeout:
+						pass
+
+			else:
+				try:
+					data = remote_sock.recv(1024)
+				except (ConnectionResetError, socket.timeout):
+					return
+				if data == b'':
+					remote_sock.close()
+					return
+				try:
+					sock.sendall(data)
+				except (ConnectionResetError, BrokenPipeError, socket.timeout):
+					remote_sock.close()
+					return
+
+			last_transfer = time.monotonic()
 
 class ProxyThread(threading.Thread):
 	def __init__(self, sock, host):