Refactor: optimize UDP module

Symmetric NAT support for now.

core/stack/tcp.go

@@ -2,12 +2,10 @@ package stack
 
 import (
 	"fmt"
-	"net"
 	"time"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
-	"gvisor.dev/gvisor/pkg/tcpip/stack"
 	"gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
 	"gvisor.dev/gvisor/pkg/waiter"
 )
@@ -36,10 +34,9 @@ func withTCPHandler() Option {
 	return func(s *Stack) error {
 		tcpForwarder := tcp.NewForwarder(s.Stack, defaultWndSize, maxConnAttempts, func(r *tcp.ForwarderRequest) {
 			var wq waiter.Queue
-			id := r.ID()
 			ep, err := r.CreateEndpoint(&wq)
 			if err != nil {
-				// prevent potential half-open TCP connection leak.
+				// RST: prevent potential half-open TCP connection leak.
 				r.Complete(true)
 				return
 			}
@@ -47,11 +44,7 @@ func withTCPHandler() Option {
 
 			setKeepalive(ep)
 
-			conn := &tcpConn{
-				Conn: gonet.NewTCPConn(&wq, ep),
-				id:   &id,
-			}
-			s.handler.Add(conn)
+			s.handler.HandleTCPConn(gonet.NewTCPConn(&wq, ep))
 		})
 		s.SetTransportProtocolHandler(tcp.ProtocolNumber, tcpForwarder.HandlePacket)
 		return nil
@@ -72,12 +65,3 @@ func setKeepalive(ep tcpip.Endpoint) error {
 	}
 	return nil
 }
-
-type tcpConn struct {
-	net.Conn
-	id *stack.TransportEndpointID
-}
-
-func (c *tcpConn) ID() *stack.TransportEndpointID {
-	return c.id
-}