Use poll instead of timeouts when waiting for response

untls_proxy.py

@@ -253,6 +253,8 @@ def proxy(sock, host):
 			print('Bad port number', file=sys.stderr)
 			sock.sendall(b'HTTP/1.0 400 Bad Request\r\n\r\nBad port number\n')
 			return
+	if remote_host[:1] == b'[' and remote_host[-1:] == b']': #IPv6
+		remote_host = remote_host[1:-1]
 	try:
 		remote_host = remote_host.decode('ascii')
 	except UnicodeDecodeError:
@@ -287,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(0)) == 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)
@@ -377,22 +382,45 @@ 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)
+	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):