From 3aeb7b39cf48a6b65ebbeeacc86d9f962a32e957 Mon Sep 17 00:00:00 2001
From: Izalia Mae <izalia@barkshark.xyz>
Date: Tue, 3 May 2022 08:00:56 -0400
Subject: [PATCH] a

---
 bin/pa-socket-bridge.py  | 159 ++++++++++++++++++++++++++-------------
 bin/proton-ge-updater.py |  12 ++-
 2 files changed, 116 insertions(+), 55 deletions(-)

diff --git a/bin/pa-socket-bridge.py b/bin/pa-socket-bridge.py
index 4a8efa7..a22f49a 100755
--- a/bin/pa-socket-bridge.py
+++ b/bin/pa-socket-bridge.py
@@ -1,75 +1,126 @@
 #!/usr/bin/env python3
-# Pulseaudio Socket Bridge v1.0.1
-# by Zoey Mae (izalia@barkshark.xyz)
-#
-# Creates a unix socket to listen on and sends any connections to the specified
-# Pulseaudio server via socat. Socat will try to reconnect and will be restarted
-# if closed for any reason
-#
-# Usage:
-# pa-socket-bridge.py [server ip]
-
-import os, signal, sys, time
-from pathlib import Path
-from subprocess import Popen
+import asyncio
+import os
+import signal
+import sys
+import traceback
 
 
 try:
-	server = sys.argv[1]
+	remote_address = sys.argv[1]
 except IndexError:
-	server = "127.0.0.1"
+	raise ValueError('No remote address specified')
 
-sock_path = Path(f'/var/run/user/{os.getuid()}/pulse/native')
-command = ['socat', f'UNIX-LISTEN:{sock_path},fork,reuseaddr,mode=777', f'TCP:{server}:4713,retry']
-shutdown = False
-
-if not sock_path.parent.exists():
-	sock_path.parent.mkdir(parents=True)
-
-if sock_path.exists():
-	print('Socket already exists. Is Pulseaudio or another bridge running?')
-	sys.exit()
+try:
+	remote_port = int(sys.argv[2])
+except IndexError:
+	remote_port = 4713
 
 
-def signal_handler(func, *args, **kwargs):
-	handler = lambda signum, frame: func(signum, frame, *args, **kwargs)
-
-	signal.signal(signal.SIGHUP, handler)
-	signal.signal(signal.SIGINT, handler)
-	signal.signal(signal.SIGQUIT, handler)
-	signal.signal(signal.SIGTERM, handler)
+running = False
+socket_path = f'/var/run/user/{os.getuid()}/pulse/native'
+tasks = []
+signals_to_handle = [
+	signal.SIGHUP,
+	signal.SIGINT,
+	signal.SIGTERM,
+	signal.SIGQUIT
+]
 
 
-def handle_shutdown(signum, frame):
-	global shutdown
-	shutdown = True
+async def proxy(reader, writer):
+	while True:
+		if not running:
+			return
+
+		try:
+			data = await reader.read(512)
+			
+			if not data:
+				break
+
+			writer.write(data)
+			await writer.drain()
+
+		except ConnectionResetError:
+			break
+
+	writer.close()
+
+
+async def handle_connection(source_reader, source_writer):
+	if not running:
+		loop.stop()
+		return
+
+	target_reader, target_writer = await asyncio.open_connection(remote_address, remote_port)
+
+	try:
+		address = source_writer.get_extra_info('peername')[0]
+	except IndexError:
+		address = 'unix socket'
+
+	print('New connection from', address)
+
+	await asyncio.gather(*[proxy(source_reader, target_writer), proxy(target_reader, source_writer)])
+
+	print('Closed connection to', address)
+
+
+def signal_handler(callback=None):
+	for sig in signals_to_handle:
+		if callback:
+			loop.add_signal_handler(sig, callback)
+
+		else:
+			loop.remove_signal_handler(sig)
+
+
+def stop(*args):
+	global running
+
+	if not running:
+		return
+
+	running = False
+	signal_handler()
+
+	for server in servers:
+		server.close()
+
+	for task in tasks:
+		task.cancel()
+
+		while True:
+			if task.done():
+				break
+
+	loop.stop()
 
 
 if __name__ == '__main__':
-	print('Starting Pulseaudio socket bridge')
+	if os.path.exists(socket_path):
+		os.remove(socket_path)
 
-	signal_handler(handle_shutdown)
+	loop = asyncio.new_event_loop()
 
-	proc = Popen(command)
-	timer = 0.0
+	signal_handler(stop)
 
-	while not shutdown:
-		if proc.poll() != None:
-			print('Restarting process')
-			proc = Popen(command)
+	servers = [
+		asyncio.start_server(handle_connection, 'localhost', 4713),
+		asyncio.start_unix_server(handle_connection, path=socket_path)
+	]
 
-		time.sleep(0.25)
+	for server in servers:
+		tasks.append(loop.create_task(server))
 
-	print('Terminating process')
-	proc.terminate()
+	running = True
 
-	while proc.poll() == None:
-		time.sleep(0.25)
-		timer += 0.25
+	try:
+		loop.run_forever()
+	except KeyboardInterrupt:
+		pass
 
-		if timer >= 5.0:
-			proc.kill()
-			print('Socat failed to terminate. Killing process.')
-			break
+	stop()
 
-	print('Bye! :3')
+	loop.close()
diff --git a/bin/proton-ge-updater.py b/bin/proton-ge-updater.py
index e2da768..402ebf1 100755
--- a/bin/proton-ge-updater.py
+++ b/bin/proton-ge-updater.py
@@ -8,9 +8,19 @@ from urllib.request import Request, urlopen, urlretrieve
 
 
 url = 'https://api.github.com/repos/GloriousEggroll/proton-ge-custom/releases/latest'
-protondir = Path('~/.steam/compatibilitytools.d/').expanduser()
+compatdirs = [
+	'~/.steam/compatibilitytools.d',
+	'~/.steam/debian-installation/compatibilitytools.d',
+	'~/.steam/steam/compatibilitytools.d'
+]
 
 
+for directory in compatdirs:
+	protondir = Path(directory).expanduser()
+
+	if protondir.exists():
+		break
+
 if not protondir.exists():
 	protondir.mkdir(parents=True, exist_ok=True)