tstest/natlab: fix conn.Close race with conn.ReadFromUDPAddrPort (#16710)

If a conn.Close call raced conn.ReadFromUDPAddrPort before it could
"register" itself as an active read, the conn.ReadFromUDPAddrPort would
never return.

This commit replaces all the activeRead and breakActiveReads machinery
with a channel. These constructs were only depended upon by
SetReadDeadline, and SetReadDeadline was unused.

Updates #16707

Signed-off-by: Jordan Whited <jordan@tailscale.com>

tstest/natlab/natlab.go

@@ -687,6 +687,7 @@ func (m *Machine) ListenPacket(ctx context.Context, network, address string) (ne
 		m:        m,
 		fam:      fam,
 		ipp:      ipp,
+		closedCh: make(chan struct{}),
 		in:       make(chan *Packet, 100), // arbitrary
 	}
 	switch c.fam {
@@ -716,50 +717,14 @@ type conn struct {
 	fam uint8 // 0, 4, or 6
 	ipp netip.AddrPort
 
-	mu           sync.Mutex
+	closeOnce sync.Once
-	closed       bool
+	closedCh  chan struct{} // closed by Close
-	readDeadline time.Time
-	activeReads  map[*activeRead]bool
-	in           chan *Packet
-}
-
-type activeRead struct {
-	cancel context.CancelFunc
-}
-
-// canRead reports whether we can do a read.
-func (c *conn) canRead() error {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	if c.closed {
-		return net.ErrClosed
-	}
-	if !c.readDeadline.IsZero() && c.readDeadline.Before(time.Now()) {
-		return errors.New("read deadline exceeded")
-	}
-	return nil
-}
 
-func (c *conn) registerActiveRead(ar *activeRead, active bool) {
+	in chan *Packet
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	if c.activeReads == nil {
-		c.activeReads = make(map[*activeRead]bool)
-	}
-	if active {
-		c.activeReads[ar] = true
-	} else {
-		delete(c.activeReads, ar)
-	}
 }
 
 func (c *conn) Close() error {
-	c.mu.Lock()
+	c.closeOnce.Do(func() {
-	defer c.mu.Unlock()
-	if c.closed {
-		return nil
-	}
-	c.closed = true
 		switch c.fam {
 		case 0:
 			c.m.unregisterConn4(c)
@@ -769,17 +734,11 @@ func (c *conn) Close() error {
 		case 6:
 			c.m.unregisterConn6(c)
 		}
-	c.breakActiveReadsLocked()
+		close(c.closedCh)
+	})
 	return nil
 }
 
-func (c *conn) breakActiveReadsLocked() {
-	for ar := range c.activeReads {
-		ar.cancel()
-	}
-	c.activeReads = nil
-}
-
 func (c *conn) LocalAddr() net.Addr {
 	return &net.UDPAddr{
 		IP:   c.ipp.Addr().AsSlice(),
@@ -809,25 +768,13 @@ func (c *conn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
 }
 
 func (c *conn) ReadFromUDPAddrPort(p []byte) (n int, addr netip.AddrPort, err error) {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	ar := &activeRead{cancel: cancel}
-
-	if err := c.canRead(); err != nil {
-		return 0, netip.AddrPort{}, err
-	}
-
-	c.registerActiveRead(ar, true)
-	defer c.registerActiveRead(ar, false)
-
 	select {
+	case <-c.closedCh:
+		return 0, netip.AddrPort{}, net.ErrClosed
 	case pkt := <-c.in:
 		n = copy(p, pkt.Payload)
 		pkt.Trace("PacketConn.ReadFrom")
 		return n, pkt.Src, nil
-	case <-ctx.Done():
-		return 0, netip.AddrPort{}, context.DeadlineExceeded
 	}
 }
 
@@ -857,18 +804,5 @@ func (c *conn) SetWriteDeadline(t time.Time) error {
 	panic("SetWriteDeadline unsupported; TODO when needed")
 }
 func (c *conn) SetReadDeadline(t time.Time) error {
-	c.mu.Lock()
+	panic("SetReadDeadline unsupported; TODO when needed")
-	defer c.mu.Unlock()
-
-	now := time.Now()
-	if t.After(now) {
-		panic("SetReadDeadline in the future not yet supported; TODO?")
-	}
-
-	if !t.IsZero() && t.Before(now) {
-		c.breakActiveReadsLocked()
-	}
-	c.readDeadline = t
-
-	return nil
 }