From 2fd486a6e4e40fa1c535a68ad1a82b965920962c Mon Sep 17 00:00:00 2001
From: Cyberes <cyberes@evulid.cc>
Date: Thu, 27 Jun 2024 17:38:04 -0600
Subject: [PATCH] handle more failures

---
 server/server.go | 16 ++++++++++------
 test/test.py     |  8 +++++++-
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/server/server.go b/server/server.go
index 22cec5b..c15f341 100644
--- a/server/server.go
+++ b/server/server.go
@@ -82,7 +82,7 @@ func main() {
 
 			forwardConn, err := net.DialUDP("udp", nil, forwardAddr)
 			if err != nil {
-				logger.Errorln("Error dialing to forward address:", err)
+				logger.Fatalln("Error dialing to forward address:", err)
 				return
 			}
 			defer forwardConn.Close()
@@ -93,27 +93,31 @@ func main() {
 				buf := make([]byte, 1500) // 1500 is the standard internet MTU
 				n, addr, err := conn.ReadFromUDP(buf)
 				if err != nil {
-					logger.Errorln("Error reading from UDP:", err)
-					return
+					logger.Fatalf(formatLogMsg(addr.IP.String(), dest.Address, dest.Port, fmt.Sprintf("Error reading from UDP: %s", err)))
+					continue
 				}
 
 				// Forward the decrypted message
 				go func(addr *net.UDPAddr, buf []byte, n int) {
 					plaintext, err := encryption.Decrypt(key, buf[:n]) // Use only the part of the buffer that has data
 					if err != nil {
-						fmt.Println("Error decrypting message:", err)
+						logger.Errorf(formatLogMsg(addr.IP.String(), dest.Address, dest.Port, fmt.Sprintf("Error decrypting message: %s", err)))
 						return
 					}
 
 					_, err = forwardConn.Write(plaintext)
 					if err != nil {
-						fmt.Println("Error forwarding message:", err)
+						logger.Errorf(formatLogMsg(addr.IP.String(), dest.Address, dest.Port, fmt.Sprintf("Error forwarding message: %s", err)))
 						return
 					}
-					logger.Infof("%s -> %s:%d -- %s\n", addr.IP, dest.Address, dest.Port, string(plaintext))
+					logger.Infof(formatLogMsg(addr.IP.String(), dest.Address, dest.Port, string(plaintext)))
 				}(addr, buf, n)
 			}
 		}(port, dest)
 	}
 	select {}
 }
+
+func formatLogMsg(srcIp string, destIp string, destPort int, msg string) string {
+	return fmt.Sprintf("%s -> %s:%d -- %s", srcIp, destIp, destPort, msg)
+}
diff --git a/test/test.py b/test/test.py
index 90d5003..46d6e4c 100644
--- a/test/test.py
+++ b/test/test.py
@@ -21,8 +21,14 @@ ciphertext, tag = cipher.encrypt_and_digest(message)
 
 # Send the encrypted message to the server
 sock.sendto(cipher.nonce + ciphertext + tag, server_address)
-sock.close()
 
 cipher = ChaCha20_Poly1305.new(key=key, nonce=cipher.nonce)
 plaintext = cipher.decrypt_and_verify(ciphertext, tag)
 print(plaintext)
+
+time.sleep(1)
+
+# Send a bad message to the server
+sock.sendto(b"this should fail", server_address)
+
+sock.close()