Fixup buffer freeing

3 files changed, 21 insertions, 9 deletions

diff --git a/device.go b/device.go
index b1f3569..bbcd0fc 100644
--- a/device.go
+++ b/device.go
@@ -19,8 +19,7 @@ const (
 	DeviceRoutineNumberAdditional = 2
 )
 
-
-const preallocatedBuffers = 0
+var preallocatedBuffers = 0
 
 type Device struct {
 	isUp     AtomicBool // device is (going) up
diff --git a/receive.go b/receive.go
index 6b6543c..9bf3af3 100644
--- a/receive.go
+++ b/receive.go
@@ -279,6 +279,7 @@ func (device *Device) RoutineDecryption() {
 			if err != nil {
 				elem.Drop()
 				device.PutMessageBuffer(elem.buffer)
+				elem.buffer = nil
 				elem.mutex.Unlock()
 			}
 			elem.mutex.Unlock()
@@ -294,18 +295,25 @@ func (device *Device) RoutineHandshake() {
 	logError := device.log.Error
 	logDebug := device.log.Debug
 
+	var elem QueueHandshakeElement
+	var ok bool
+
 	defer func() {
 		logDebug.Println("Routine: handshake worker - stopped")
 		device.state.stopping.Done()
+		if elem.buffer != nil {
+			device.PutMessageBuffer(elem.buffer)
+		}
 	}()
 
 	logDebug.Println("Routine: handshake worker - started")
 	device.state.starting.Done()
 
-	var elem QueueHandshakeElement
-	var ok bool
-
 	for {
+		if elem.buffer != nil {
+			device.PutMessageBuffer(elem.buffer)
+		}
+
 		select {
 		case elem, ok = <-device.queue.handshake:
 		case <-device.signals.stop:
@@ -478,9 +486,14 @@ func (peer *Peer) RoutineSequentialReceiver() {
 	logError := device.log.Error
 	logDebug := device.log.Debug
 
+	var elem *QueueInboundElement
+
 	defer func() {
 		logDebug.Println(peer, "- Routine: sequential receiver - stopped")
 		peer.routines.stopping.Done()
+		if elem != nil && elem.buffer != nil {
+			device.PutMessageBuffer(elem.buffer)
+		}
 	}()
 
 	logDebug.Println(peer, "- Routine: sequential receiver - started")
@@ -488,6 +501,9 @@ func (peer *Peer) RoutineSequentialReceiver() {
 	peer.routines.starting.Done()
 
 	for {
+		if elem != nil && elem.buffer != nil {
+			device.PutMessageBuffer(elem.buffer)
+		}
 
 		select {
 
@@ -608,7 +624,6 @@ func (peer *Peer) RoutineSequentialReceiver() {
 			_, err := device.tun.device.Write(
 				elem.buffer[:offset+len(elem.packet)],
 				offset)
-			device.PutMessageBuffer(elem.buffer)
 			if err != nil {
 				logError.Println("Failed to write packet to TUN device:", err)
 			}
diff --git a/send.go b/send.go
index bc0f5e4..24e2f39 100644
--- a/send.go
+++ b/send.go
@@ -74,8 +74,7 @@ func addToNonceQueue(queue chan *QueueOutboundElement, element *QueueOutboundEle
 		default:
 			select {
 			case old := <-queue:
-				old.Drop()
-				device.PutMessageBuffer(element.buffer)
+				device.PutMessageBuffer(old.buffer)
 			default:
 			}
 		}
@@ -98,7 +97,6 @@ func addToOutboundAndEncryptionQueues(outboundQueue chan *QueueOutboundElement,
 	}
 }
 
-
 /* Queues a keepalive if no packets are queued for peer
  */
 func (peer *Peer) SendKeepalive() bool {