From 40414cb9a5b94940df291c17be351a31c8dfc836 Mon Sep 17 00:00:00 2001
From: Fierelier <fierelier@gmail.com>
Date: Wed, 14 Apr 2021 22:03:53 +0200
Subject: [PATCH] Outbound queue for client

---
 fstream-client.py | 51 ++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 42 insertions(+), 9 deletions(-)

diff --git a/fstream-client.py b/fstream-client.py
index 856b4de..d80deab 100644
--- a/fstream-client.py
+++ b/fstream-client.py
@@ -24,6 +24,7 @@ import threading
 import queue
 
 bufferSize = 1000 # buffer size in bytes
+maxAccumulatedData = 50*1000*1000 # How much data can be in an outbound thread's queue at maximum before the connection is closed?
 timeout = 15 # timeout in seconds
 connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 
@@ -37,6 +38,28 @@ class stdoutThread(threading.Thread):
 			data = self.queue.get(timeout=timeout)
 			sys.stdout.buffer.write(data)
 
+class stdinThread(threading.Thread):
+	def __init__(self,connection):
+		threading.Thread.__init__(self)
+		self.queue = queue.Queue()
+		self.connection = connection
+	
+	def run(self):
+		try:
+			while True:
+				accumulatedData = self.queue.qsize() * bufferSize
+				if accumulatedData > maxAccumulatedData:
+					print("Accumulated data limit reached. Closing.")
+					self.connection.close()
+					self.queue = False
+					return
+				data = self.queue.get()
+				self.connection.sendall(data)
+		except:
+			self.connection.close()
+			self.queue = False
+			raise
+
 def listToCommand(lst):
 	cmd = ""
 	for arg in lst:
@@ -65,17 +88,27 @@ def main():
 	connection.sendall(makePayload(sys.argv[2:]))
 	
 	if sys.argv[2] == "watch":
-		stdoutThr = stdoutThread()
-		stdoutThr.start()
-		while True:
-			data = connection.recv(bufferSize)
-			if data == b"": return
-			stdoutThr.queue.put(data)
+		try:
+			stdoutThr = stdoutThread()
+			stdoutThr.start()
+			while True:
+				data = connection.recv(bufferSize)
+				if data == b"": return
+				stdoutThr.queue.put(data)
+		except:
+			connection.close()
+			raise
 	
 	if sys.argv[2] == "broadcast":
-		while True:
-			data = sys.stdin.buffer.read(bufferSize)
-			connection.sendall(data)
+		try:
+			stdinThr = stdinThread(connection)
+			stdinThr.start()
+			while True:
+				data = sys.stdin.buffer.read(bufferSize)
+				stdinThr.queue.put(data)
+		except:
+			connection.close()
+			raise
 
 if __name__ == '__main__':
 	main()
\ No newline at end of file