From 596f67b718d2cfcdc510a6ff249596d94d3ee809 Mon Sep 17 00:00:00 2001
From: Paul Lorenz <paul.lorenz@netfoundry.io>
Date: Wed, 8 Jun 2022 10:29:18 -0400
Subject: [PATCH] Fix missing memory pool buffer returns. Fixes #1045

---
 tunnel/intercept/tproxy/tproxy_linux.go |  3 +++
 tunnel/udp_vconn/api.go                 |  2 +-
 tunnel/udp_vconn/conn.go                | 13 ++++++++++---
 tunnel/udp_vconn/manager.go             |  2 +-
 4 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/tunnel/intercept/tproxy/tproxy_linux.go b/tunnel/intercept/tproxy/tproxy_linux.go
index 506b622fd..75823bcc0 100644
--- a/tunnel/intercept/tproxy/tproxy_linux.go
+++ b/tunnel/intercept/tproxy/tproxy_linux.go
@@ -310,16 +310,19 @@ func (event *udpReadEvent) Handle(manager udp_vconn.Manager) error {
 		log := pfxlog.Logger()
 		origDest, err := getOriginalDest(event.oob)
 		if err != nil {
+			event.buf.Release()
 			return fmt.Errorf("error while getting original destination packet: %v", err)
 		}
 		log.Infof("received datagram from %v (original dest %v). Creating udp listen socket on original dest", event.srcAddr, origDest)
 		packetConn, err := listenConfig.ListenPacket(context.Background(), "udp", origDest.String())
 		if err != nil {
+			event.buf.Release()
 			return err
 		}
 		writeConn := packetConn.(*net.UDPConn)
 		writeQueue, err = manager.CreateWriteQueue(origDest, event.srcAddr, event.interceptor.service, writeConn)
 		if err != nil {
+			event.buf.Release()
 			return err
 		}
 	}
diff --git a/tunnel/udp_vconn/api.go b/tunnel/udp_vconn/api.go
index 8ac651577..7bf2a2be1 100644
--- a/tunnel/udp_vconn/api.go
+++ b/tunnel/udp_vconn/api.go
@@ -77,7 +77,7 @@ func (u UnpooledBuffer) Release() {
 
 func NewManager(provider tunnel.FabricProvider, newConnPolicy NewConnPolicy, expirationPolicy ConnExpirationPolicy) Manager {
 	manager := &manager{
-		eventC:           make(chan Event),
+		eventC:           make(chan Event, 4),
 		provider:         provider,
 		connMap:          make(map[string]*udpConn),
 		newConnPolicy:    newConnPolicy,
diff --git a/tunnel/udp_vconn/conn.go b/tunnel/udp_vconn/conn.go
index 59de80090..129627c88 100644
--- a/tunnel/udp_vconn/conn.go
+++ b/tunnel/udp_vconn/conn.go
@@ -75,14 +75,17 @@ func (conn *udpConn) WriteTo(w io.Writer) (n int64, err error) {
 	var bytesWritten int64
 	for {
 		var buf mempool.PooledBuffer
-		var ok bool
 
 		select {
-		case buf, ok = <-conn.readC:
+		case buf = <-conn.readC:
 		case <-conn.closeNotify:
+			select {
+			case buf = <-conn.readC:
+			default:
+			}
 		}
 
-		if !ok {
+		if buf == nil {
 			return bytesWritten, io.EOF
 		}
 
@@ -123,6 +126,10 @@ func (conn *udpConn) Read(b []byte) (n int, err error) {
 	select {
 	case buf = <-conn.readC:
 	case <-conn.closeNotify:
+		select {
+		case buf = <-conn.readC:
+		default:
+		}
 	}
 
 	if buf == nil {
diff --git a/tunnel/udp_vconn/manager.go b/tunnel/udp_vconn/manager.go
index 4095ecc21..c9679bbbb 100644
--- a/tunnel/udp_vconn/manager.go
+++ b/tunnel/udp_vconn/manager.go
@@ -89,7 +89,7 @@ func (manager *manager) CreateWriteQueue(targetAddr *net.UDPAddr, srcAddr net.Ad
 		return nil, errors.New("max connections exceeded")
 	}
 	conn := &udpConn{
-		readC:       make(chan mempool.PooledBuffer),
+		readC:       make(chan mempool.PooledBuffer, 4),
 		closeNotify: make(chan struct{}),
 		service:     service.Name,
 		srcAddr:     srcAddr,