tun: fix race during WriteNotify and Close

Signed-off-by: Spike Curtis <spike@coder.com>

Author: spikecurtis

tun/netstack/tun.go

@@ -44,6 +44,7 @@ type netTun struct {
 	stack          *stack.Stack
 	events         chan tun.Event
 	incomingPacket chan *buffer.View
+	closed         chan struct{}
 	mtu            int
 	dnsServers     []netip.Addr
 	hasV4, hasV6   bool
@@ -62,6 +63,7 @@ func CreateNetTUN(localAddresses, dnsServers []netip.Addr, mtu int) (tun.Device,
 		stack:          stack.New(opts),
 		events:         make(chan tun.Event, 10),
 		incomingPacket: make(chan *buffer.View),
+		closed:         make(chan struct{}),
 		dnsServers:     dnsServers,
 		mtu:            mtu,
 	}
@@ -120,17 +122,17 @@ func (tun *netTun) Events() <-chan tun.Event {
 }
 
 func (tun *netTun) Read(buf [][]byte, sizes []int, offset int) (int, error) {
-	view, ok := <-tun.incomingPacket
-	if !ok {
+	select {
+	case view := <-tun.incomingPacket:
+		n, err := view.Read(buf[0][offset:])
+		if err != nil {
+			return 0, err
+		}
+		sizes[0] = n
+		return 1, nil
+	case <-tun.closed:
 		return 0, os.ErrClosed
 	}
-
-	n, err := view.Read(buf[0][offset:])
-	if err != nil {
-		return 0, err
-	}
-	sizes[0] = n
-	return 1, nil
 }
 
 func (tun *netTun) Write(buf [][]byte, offset int) (int, error) {
@@ -162,7 +164,10 @@ func (tun *netTun) WriteNotify() {
 	view := pkt.ToView()
 	pkt.DecRef()
 
-	tun.incomingPacket <- view
+	select {
+	case tun.incomingPacket <- view:
+	case <-tun.closed:
+	}
 }
 
 func (tun *netTun) Close() error {
@@ -174,8 +179,8 @@ func (tun *netTun) Close() error {
 
 	tun.ep.Close()
 
-	if tun.incomingPacket != nil {
-		close(tun.incomingPacket)
+	if tun.closed != nil {
+		close(tun.closed)
 	}
 
 	return nil