From 4faba64dc091c706f7a4c5fe256faa8e9caa5474 Mon Sep 17 00:00:00 2001
From: Manan Garg <manantechy@gmail.com>
Date: Mon, 27 Nov 2023 17:33:13 +0530
Subject: [PATCH] add try catch in message broker

---
 amqp/broker.py | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/amqp/broker.py b/amqp/broker.py
index 5730c75..92d1a28 100644
--- a/amqp/broker.py
+++ b/amqp/broker.py
@@ -47,9 +47,14 @@ def start(self):
     def _handle_client(self, client: socket.socket):
         with client:
             # Receive key
-            key = client.recv(1024).decode()
-            print(key)
-            print(type(key))
+            try:
+                key = client.recv(1024).decode()
+                print(key)
+                print(type(key))
+            except Exception as e:
+                print("Error receiving key: {}".format(e))
+                client.sendall("Error receiving key: {}\nClosing your connection".format(e).encode())
+                return
             # Bind client to exchange
             try:
                 self._exchange.bind(key, client)
@@ -66,8 +71,16 @@ def _handle_client(self, client: socket.socket):
                     print("Closing connection to client: {}".format(client))
                     break
                 # Parse message as json
-                message = base64.b64decode(message)
-                msg_json = json.loads(message)
+                try:
+                    message = base64.b64decode(message)
+                except:
+                    client.sendall(b"Message is not properly base64 encoded")
+                    continue
+                try:
+                    msg_json = json.loads(message)
+                except:
+                    client.sendall(b"Message is not valid json")
+                    continue
                 # Check if message has routing key
                 if 'key' not in msg_json:
                     client.sendall(b"Message has no routing key")