From 9985e0996932b954268a93532290f159b94d6399 Mon Sep 17 00:00:00 2001
From: Jake Hillion <jake@hillion.co.uk>
Date: Wed, 31 Mar 2021 18:45:41 +0100
Subject: [PATCH] udp listener context

---
 udp/listener.go | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/udp/listener.go b/udp/listener.go
index f2b9e0d..cf3e04f 100644
--- a/udp/listener.go
+++ b/udp/listener.go
@@ -5,6 +5,7 @@ import (
 	"log"
 	"mpbl3p/proxy"
 	"net"
+	"time"
 )
 
 type ComparableUdpAddress struct {
@@ -45,15 +46,19 @@ func NewListener(ctx context.Context, p *proxy.Proxy, local string, v func() pro
 	receivedConnections := make(map[ComparableUdpAddress]*Flow)
 
 	go func() {
-		for {
+		for ctx.Err() == nil {
 			buf := make([]byte, 6000)
 
-			log.Println("listening...")
-			n, addr, err := pconn.ReadFromUDP(buf)
-			if err != nil {
+			if err := pconn.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
+				panic(err)
+			}
+			n, addr, err := pconn.ReadFromUDP(buf)
+			if err != nil {
+				if e, ok := err.(net.Error); ok && e.Timeout() {
+					continue
+				}
 				panic(err)
 			}
-			log.Println("listened")
 
 			raddr := fromUdpAddress(*addr)
 			if f, exists := receivedConnections[raddr]; exists {
@@ -84,7 +89,9 @@ func NewListener(ctx context.Context, p *proxy.Proxy, local string, v func() pro
 			p.AddProducer(ctx, &f, v)
 
 			log.Println("handling...")
-			f.queueDatagram(ctx, buf[:n])
+			if err := f.queueDatagram(ctx, buf[:n]); err != nil {
+				return
+			}
 			log.Println("handled")
 		}
 	}()