From 0ac0006a631ec6150d4aaaddb6ba2caef57be349 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Juhani=20Krekel=C3=A4?= <juhani@krekelä.fi>
Date: Sun, 17 Oct 2021 18:58:58 +0300
Subject: [PATCH] Use poll instead of timeouts when waiting for response

---
 untls_proxy.py | 38 ++++++++++++++++++++------------------
 1 file changed, 20 insertions(+), 18 deletions(-)

diff --git a/untls_proxy.py b/untls_proxy.py
index c6d847e..b825270 100755
--- a/untls_proxy.py
+++ b/untls_proxy.py
@@ -289,23 +289,26 @@ def proxy(sock, host):
 		# Keep sending request body, if any, until we get a response from remote
 		poll = select.poll()
 		poll.register(remote_sock, select.POLLIN)
-		sock.settimeout(1)
-		while len(poll.poll(1000)) == 0:
-			try:
-				data = sock.recv(1024)
-			except ConnectionResetError:
-				break
-			except socket.timeout:
-				continue
-			if data == b'': break
-			# Save the part we've sent already in case we need to re-send request
-			request_data.extend(data)
-			try:
-				remote_sock.sendall(data)
-			except (ConnectionResetError, BrokenPipeError):
-				print('Remote hung up', file=sys.stderr)
-				return
-		sock.settimeout(None)
+		poll.register(sock, select.POLLIN)
+		no_response = True
+		while no_response:
+			for fd, _ in poll.poll():
+				if fd == remote_sock.fileno():
+					no_response = False
+					break
+				else:
+					try:
+						data = sock.recv(1024)
+					except ConnectionResetError:
+						break
+					if data == b'': break
+					# Save the part we've sent already in case we need to re-send request
+					request_data.extend(data)
+					try:
+						remote_sock.sendall(data)
+					except (ConnectionResetError, BrokenPipeError):
+						print('Remote hung up', file=sys.stderr)
+						return
 
 		# Get response headers
 		remote_sock.settimeout(10)
@@ -382,7 +385,6 @@ def proxy(sock, host):
 	print('', file=sys.stderr)
 	sock.settimeout(60)
 	remote_sock.settimeout(60)
-	poll.register(sock, select.POLLIN)
 	last_transfer = time.monotonic()
 	while True:
 		events = poll.poll(60_000)