Add 1 min timeout for no new data transferred

untls_proxy.py

@@ -227,10 +227,10 @@ def proxy(sock, host):
 	del username
 	del password
 
-	# Remove headers that don't need forwarding
-	headers = dict((key, value) for key, value in headers.items() if not key.startswith(b'proxy-'))
+	# Remove headers that don't need forwarding or are overwritten
+	headers = dict((key, value) for key, value in headers.items() if not key.startswith(b'proxy-') and not key in (b'connection', b'keep-alive'))
 
-	# TODO: connection: close
+	headers[b'connection'] = b' close'
 
 	# Split url into its constituents
 	fields = url.split(b'://', 1)
@@ -245,7 +245,7 @@ def proxy(sock, host):
 	elif b':' not in remote_host:
 		port = None
 	else:
-		remote_host, port = remote_host.rsplit(b':')
+		remote_host, port = remote_host.rsplit(b':', 1)
 		try:
 			port = int(port)
 			if port < 1 or port > 0xffff: raise ValueError
@@ -376,18 +376,20 @@ def proxy(sock, host):
 
 	del request_data
 
-	# TODO: Timeout
 	# TODO: Un-https links
+	# TODO: Keep sending request body, if any
 	print('', file=sys.stderr)
+	sock.settimeout(60)
+	remote_sock.settimeout(60)
 	while True:
 		try:
 			data = remote_sock.recv(1024)
-		except ConnectionResetError:
+		except (ConnectionResetError, socket.timeout):
 			return
 		if data == b'': break
 		try:
 			sock.sendall(data)
-		except (ConnectionResetError, BrokenPipeError):
+		except (ConnectionResetError, BrokenPipeError, socket.timeout):
 			break
 
 	remote_sock.close()