udp listener context

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")
 		}
 	}()