From 7a26ffdcf7c2f3fdd7fe7d7a659c7c35eb888d98 Mon Sep 17 00:00:00 2001
From: CPunch <sethtstubbs@gmail.com>
Date: Sat, 11 Mar 2023 15:02:01 -0600
Subject: [PATCH] server/peer: wrap 'defer pool.Put()' in a closure

---
 server/peer.go | 43 ++++++++++++++++++++++---------------------
 1 file changed, 22 insertions(+), 21 deletions(-)

diff --git a/server/peer.go b/server/peer.go
index bf28cbe..4b17382 100644
--- a/server/peer.go
+++ b/server/peer.go
@@ -119,32 +119,33 @@ func (peer *Peer) Handler() {
 		}
 
 		// grab buffer && read packet body
-		buf := pool.Get()
-		defer pool.Put(buf)
-		if _, err := buf.ReadFrom(io.LimitReader(peer.conn, int64(sz))); err != nil {
-			log.Printf("[FATAL] failed to read packet body! %v", err)
-			return
-		}
+		if err := func() error { // we wrap this in a closure so we can easily defer the buffer return to pool
+			buf := pool.Get()
+			defer pool.Put(buf)
+			if _, err := buf.ReadFrom(io.LimitReader(peer.conn, int64(sz))); err != nil {
+				return fmt.Errorf("failed to read packet body! %v", err)
+			}
 
-		// decrypt
-		protocol.DecryptData(buf.Bytes(), peer.E_key)
-		pkt := protocol.NewPacket(buf)
+			// decrypt
+			protocol.DecryptData(buf.Bytes(), peer.E_key)
+			pkt := protocol.NewPacket(buf)
 
-		// create packet && read typeID
-		var typeID uint32
-		if err := pkt.Decode(&typeID); err != nil {
-			log.Printf("[FATAL] failed to read packet type! %v", err)
-			return
-		}
+			// create packet && read typeID
+			var typeID uint32
+			if err := pkt.Decode(&typeID); err != nil {
+				return fmt.Errorf("failed to read packet type! %v", err)
+			}
 
-		// dispatch packet
-		log.Printf("Got packet ID: %x, with a sizeof: %d\n", typeID, sz)
-		if err := peer.handler.HandlePacket(peer, typeID, pkt); err != nil {
+			// dispatch packet
+			log.Printf("Got packet ID: %x, with a sizeof: %d\n", typeID, sz)
+			if err := peer.handler.HandlePacket(peer, typeID, pkt); err != nil {
+				return err
+			}
+
+			return nil
+		}(); err != nil {
 			log.Printf("[FATAL] %v", err)
 			return
 		}
-
-		// restore buffer to pool
-		pool.Put(buf)
 	}
 }