Fix gLazyConn race

redirect_nftables.go

@@ -64,7 +64,7 @@ func (r *autoRedirect) setupNFTables() error {
 			r.nftablesCreateRedirect(nft, table, chainOutput)
 
 			chainOutputUDP := nft.AddChain(&nftables.Chain{
-				Name:     "output_udp",
+				Name:     "output_udp_icmp",
 				Table:    table,
 				Hooknum:  nftables.ChainHookOutput,
 				Priority: nftables.ChainPriorityMangle,

redirect_nftables_rules.go

@@ -439,6 +439,20 @@ func (r *autoRedirect) nftablesCreateExcludeRules(nft *nftables.Conn, table *nft
 	if r.tunOptions.AutoRedirectMarkMode &&
 		((chain.Hooknum == nftables.ChainHookOutput && chain.Type == nftables.ChainTypeRoute) ||
 			(chain.Hooknum == nftables.ChainHookPrerouting && chain.Type == nftables.ChainTypeFilter)) {
+		ipProto := &nftables.Set{
+			Table:     table,
+			Anonymous: true,
+			Constant:  true,
+			KeyType:   nftables.TypeInetProto,
+		}
+		err := nft.AddSet(ipProto, []nftables.SetElement{
+			{Key: []byte{unix.IPPROTO_UDP}},
+			{Key: []byte{unix.IPPROTO_ICMP}},
+			{Key: []byte{unix.IPPROTO_ICMPV6}},
+		})
+		if err != nil {
+			return err
+		}
 		nft.AddRule(&nftables.Rule{
 			Table: table,
 			Chain: chain,
@@ -447,10 +461,11 @@ func (r *autoRedirect) nftablesCreateExcludeRules(nft *nftables.Conn, table *nft
 					Key:      expr.MetaKeyL4PROTO,
 					Register: 1,
 				},
-				&expr.Cmp{
-					Op:       expr.CmpOpNeq,
-					Register: 1,
-					Data:     []byte{unix.IPPROTO_UDP},
+				&expr.Lookup{
+					SourceRegister: 1,
+					SetID:          ipProto.ID,
+					SetName:        ipProto.Name,
+					Invert:         true,
 				},
 				&expr.Verdict{
 					Kind: expr.VerdictReturn,
@@ -682,6 +697,7 @@ func (r *autoRedirect) nftablesCreateDNSHijackRulesForFamily(
 			Register: 1,
 			Data:     binaryutil.BigEndian.PutUint16(53),
 		},
+		&expr.Counter{},
 		&expr.Immediate{
 			Register: 1,
 			Data:     dnsServer.AsSlice(),
@@ -691,7 +707,6 @@ func (r *autoRedirect) nftablesCreateDNSHijackRulesForFamily(
 			Family:     uint32(family),
 			RegAddrMin: 1,
 		},
-		&expr.Counter{},
 	)
 	nft.AddRule(&nftables.Rule{
 		Table: table,
@@ -727,9 +742,7 @@ func (r *autoRedirect) nftablesCreateUnreachable(
 				Data:     []byte{uint8(nfProto)},
 			},
 			&expr.Counter{},
-			&expr.Verdict{
-				Kind: expr.VerdictDrop,
-			},
+			&expr.Reject{},
 		},
 	})
 }

stack_gvisor_lazy.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"net"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/sagernet/gvisor/pkg/tcpip"
@@ -17,19 +18,25 @@ import (
 )
 
 type gLazyConn struct {
-	tcpConn       *gonet.TCPConn
-	parentCtx     context.Context
-	stack         *stack.Stack
-	request       *tcp.ForwarderRequest
-	localAddr     net.Addr
-	remoteAddr    net.Addr
-	handshakeDone bool
-	handshakeErr  error
+	tcpConn         *gonet.TCPConn
+	parentCtx       context.Context
+	stack           *stack.Stack
+	request         *tcp.ForwarderRequest
+	localAddr       net.Addr
+	remoteAddr      net.Addr
+	handshakeAccess sync.Mutex
+	handshakeDone   bool
+	handshakeErr    error
 }
 
 func (c *gLazyConn) HandshakeContext(ctx context.Context) error {
 	if c.handshakeDone {
-		return nil
+		return c.handshakeErr
+	}
+	c.handshakeAccess.Lock()
+	defer c.handshakeAccess.Unlock()
+	if c.handshakeDone {
+		return c.handshakeErr
 	}
 	defer func() {
 		c.handshakeDone = true
@@ -64,6 +71,11 @@ func (c *gLazyConn) HandshakeContext(ctx context.Context) error {
 }
 
 func (c *gLazyConn) HandshakeFailure(err error) error {
+	if c.handshakeDone {
+		return os.ErrInvalid
+	}
+	c.handshakeAccess.Lock()
+	defer c.handshakeAccess.Unlock()
 	if c.handshakeDone {
 		return os.ErrInvalid
 	}
@@ -78,25 +90,17 @@ func (c *gLazyConn) HandshakeSuccess() error {
 }
 
 func (c *gLazyConn) Read(b []byte) (n int, err error) {
-	if !c.handshakeDone {
-		err = c.HandshakeContext(context.Background())
-		if err != nil {
-			return
-		}
-	} else if c.handshakeErr != nil {
-		return 0, c.handshakeErr
+	err = c.HandshakeContext(context.Background())
+	if err != nil {
+		return
 	}
 	return c.tcpConn.Read(b)
 }
 
 func (c *gLazyConn) Write(b []byte) (n int, err error) {
-	if !c.handshakeDone {
-		err = c.HandshakeContext(context.Background())
-		if err != nil {
-			return
-		}
-	} else if c.handshakeErr != nil {
-		return 0, c.handshakeErr
+	err = c.HandshakeContext(context.Background())
+	if err != nil {
+		return
 	}
 	return c.tcpConn.Write(b)
 }
@@ -110,79 +114,80 @@ func (c *gLazyConn) RemoteAddr() net.Addr {
 }
 
 func (c *gLazyConn) SetDeadline(t time.Time) error {
-	if !c.handshakeDone {
-		err := c.HandshakeContext(context.Background())
-		if err != nil {
-			return err
-		}
-	} else if c.handshakeErr != nil {
-		return c.handshakeErr
+	err := c.HandshakeContext(context.Background())
+	if err != nil {
+		return err
 	}
 	return c.tcpConn.SetDeadline(t)
 }
 
 func (c *gLazyConn) SetReadDeadline(t time.Time) error {
-	if !c.handshakeDone {
-		err := c.HandshakeContext(context.Background())
-		if err != nil {
-			return err
-		}
-	} else if c.handshakeErr != nil {
-		return c.handshakeErr
+	err := c.HandshakeContext(context.Background())
+	if err != nil {
+		return err
 	}
 	return c.tcpConn.SetReadDeadline(t)
 }
 
 func (c *gLazyConn) SetWriteDeadline(t time.Time) error {
-	if !c.handshakeDone {
-		err := c.HandshakeContext(context.Background())
-		if err != nil {
-			return err
-		}
-	} else if c.handshakeErr != nil {
-		return c.handshakeErr
+	err := c.HandshakeContext(context.Background())
+	if err != nil {
+		return err
 	}
 	return c.tcpConn.SetWriteDeadline(t)
 }
 
 func (c *gLazyConn) Close() error {
 	if !c.handshakeDone {
-		c.request.Complete(true)
-		c.handshakeErr = net.ErrClosed
-		return nil
-	} else if c.handshakeErr != nil {
-		return nil
+		c.handshakeAccess.Lock()
+		if !c.handshakeDone {
+			c.request.Complete(true)
+			c.handshakeErr = net.ErrClosed
+			c.handshakeDone = true
+			return nil
+		}
+		c.handshakeAccess.Unlock()
 	}
 	return c.tcpConn.Close()
 }
 
 func (c *gLazyConn) CloseRead() error {
 	if !c.handshakeDone {
-		c.request.Complete(true)
-		c.handshakeErr = net.ErrClosed
-		return nil
-	} else if c.handshakeErr != nil {
-		return nil
+		c.handshakeAccess.Lock()
+		if !c.handshakeDone {
+			c.request.Complete(true)
+			c.handshakeErr = net.ErrClosed
+			c.handshakeDone = true
+			return nil
+		}
+		c.handshakeAccess.Unlock()
 	}
 	return c.tcpConn.CloseRead()
 }
 
 func (c *gLazyConn) CloseWrite() error {
 	if !c.handshakeDone {
-		c.request.Complete(true)
-		c.handshakeErr = net.ErrClosed
-		return nil
-	} else if c.handshakeErr != nil {
-		return nil
+		c.handshakeAccess.Lock()
+		if !c.handshakeDone {
+			c.request.Complete(true)
+			c.handshakeErr = net.ErrClosed
+			c.handshakeDone = true
+			return nil
+		}
+		c.handshakeAccess.Unlock()
 	}
 	return c.tcpConn.CloseRead()
 }
 
 func (c *gLazyConn) ReaderReplaceable() bool {
+	c.handshakeAccess.Lock()
+	defer c.handshakeAccess.Unlock()
 	return c.handshakeDone && c.handshakeErr == nil
 }
 
 func (c *gLazyConn) WriterReplaceable() bool {
+	c.handshakeAccess.Lock()
+	defer c.handshakeAccess.Unlock()
 	return c.handshakeDone && c.handshakeErr == nil
 }