From dcccdaddafd7a6a848627685f9ee109316ca3728 Mon Sep 17 00:00:00 2001
From: gVisor bot <gvisor-bot@google.com>
Date: Sun, 12 Mar 2023 19:03:03 +0800
Subject: [PATCH] fix: tuic relay tuic

---
 transport/tuic/conn.go     | 30 ++++++++++++++++++++++++++----
 transport/tuic/protocol.go | 20 +++++---------------
 transport/tuic/server.go   | 24 +-----------------------
 3 files changed, 32 insertions(+), 42 deletions(-)

diff --git a/transport/tuic/conn.go b/transport/tuic/conn.go
index 7ecc3f0d..ee687425 100644
--- a/transport/tuic/conn.go
+++ b/transport/tuic/conn.go
@@ -200,7 +200,7 @@ func (q *quicStreamPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err err
 }
 
 func (q *quicStreamPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
-	if len(p) > q.maxUdpRelayPacketSize {
+	if q.udpRelayMode != "quic" && len(p) > q.maxUdpRelayPacketSize {
 		return 0, fmt.Errorf("udp packet too large(%d > %d)", len(p), q.maxUdpRelayPacketSize)
 	}
 	if q.closed {
@@ -215,7 +215,6 @@ func (q *quicStreamPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err erro
 			q.deferQuicConnFn(q.quicConn, err)
 		}()
 	}
-	addr.String()
 	buf := pool.GetBuffer()
 	defer pool.PutBuffer(buf)
 	addrPort, err := netip.ParseAddrPort(addr.String())
@@ -239,7 +238,8 @@ func (q *quicStreamPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err erro
 			return
 		}
 	default: // native
-		err = q.quicConn.SendMessage(buf.Bytes())
+		data := buf.Bytes()
+		err = q.quicConn.SendMessage(data)
 		if err != nil {
 			return
 		}
@@ -250,7 +250,29 @@ func (q *quicStreamPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err erro
 }
 
 func (q *quicStreamPacketConn) LocalAddr() net.Addr {
-	return q.quicConn.LocalAddr()
+	addr := q.quicConn.LocalAddr()
+	if q.inputConn != nil { // client
+		return &packetAddr{addrStr: q.quicConn.LocalAddr().String(), connId: q.connId, rawAddr: addr}
+	}
+	return addr // server
 }
 
 var _ net.PacketConn = &quicStreamPacketConn{}
+
+type packetAddr struct {
+	addrStr string
+	connId  uint32
+	rawAddr net.Addr
+}
+
+func (a packetAddr) Network() string {
+	return "tuic"
+}
+
+func (a packetAddr) String() string {
+	return fmt.Sprintf("%s-%d", a.addrStr, a.connId)
+}
+
+func (a packetAddr) RawAddr() net.Addr {
+	return a.rawAddr
+}
diff --git a/transport/tuic/protocol.go b/transport/tuic/protocol.go
index a54c8e54..b14074a7 100644
--- a/transport/tuic/protocol.go
+++ b/transport/tuic/protocol.go
@@ -114,9 +114,6 @@ func NewAuthenticate(TKN [32]byte) Authenticate {
 
 func ReadAuthenticateWithHead(head CommandHead, reader BufferedReader) (c Authenticate, err error) {
 	c.CommandHead = head
-	if err != nil {
-		return
-	}
 	if c.CommandHead.TYPE != AuthenticateType {
 		err = fmt.Errorf("error command type: %s", c.CommandHead.TYPE)
 		return
@@ -170,9 +167,6 @@ func NewConnect(ADDR Address) Connect {
 
 func ReadConnectWithHead(head CommandHead, reader BufferedReader) (c Connect, err error) {
 	c.CommandHead = head
-	if err != nil {
-		return
-	}
 	if c.CommandHead.TYPE != ConnectType {
 		err = fmt.Errorf("error command type: %s", c.CommandHead.TYPE)
 		return
@@ -228,9 +222,6 @@ func NewPacket(ASSOC_ID uint32, LEN uint16, ADDR Address, DATA []byte) Packet {
 
 func ReadPacketWithHead(head CommandHead, reader BufferedReader) (c Packet, err error) {
 	c.CommandHead = head
-	if err != nil {
-		return
-	}
 	if c.CommandHead.TYPE != PacketType {
 		err = fmt.Errorf("error command type: %s", c.CommandHead.TYPE)
 		return
@@ -305,9 +296,6 @@ func NewDissociate(ASSOC_ID uint32) Dissociate {
 
 func ReadDissociateWithHead(head CommandHead, reader BufferedReader) (c Dissociate, err error) {
 	c.CommandHead = head
-	if err != nil {
-		return
-	}
 	if c.CommandHead.TYPE != DissociateType {
 		err = fmt.Errorf("error command type: %s", c.CommandHead.TYPE)
 		return
@@ -476,15 +464,17 @@ func NewAddress(metadata *C.Metadata) Address {
 
 func NewAddressAddrPort(addrPort netip.AddrPort) Address {
 	var addrType byte
-	if addrPort.Addr().Is4() {
+	port := addrPort.Port()
+	addr := addrPort.Addr().Unmap()
+	if addr.Is4() {
 		addrType = AtypIPv4
 	} else {
 		addrType = AtypIPv6
 	}
 	return Address{
 		TYPE: addrType,
-		ADDR: addrPort.Addr().AsSlice(),
-		PORT: addrPort.Port(),
+		ADDR: addr.AsSlice(),
+		PORT: port,
 	}
 }
 
diff --git a/transport/tuic/server.go b/transport/tuic/server.go
index 19398fde..770e1869 100644
--- a/transport/tuic/server.go
+++ b/transport/tuic/server.go
@@ -5,7 +5,6 @@ import (
 	"bytes"
 	"context"
 	"crypto/tls"
-	"fmt"
 	"net"
 	"sync"
 	"sync/atomic"
@@ -154,14 +153,10 @@ func (s *serverHandler) parsePacket(packet Packet, udpRelayMode string) (err err
 	return s.HandleUdpFn(packet.ADDR.SocksAddr(), &serverUDPPacket{
 		pc:     pc,
 		packet: &packet,
-		rAddr:  s.genServerAssocIdAddr(assocId, s.quicConn.RemoteAddr()),
+		rAddr:  &packetAddr{addrStr: "tuic-" + s.uuid.String(), connId: assocId, rawAddr: s.quicConn.RemoteAddr()},
 	})
 }
 
-func (s *serverHandler) genServerAssocIdAddr(assocId uint32, addr net.Addr) net.Addr {
-	return &ServerAssocIdAddr{assocId: fmt.Sprintf("tuic-%s-%d", s.uuid.String(), assocId), addr: addr}
-}
-
 func (s *serverHandler) handleStream() (err error) {
 	for {
 		var quicStream quic.Stream
@@ -276,23 +271,6 @@ func (s *serverHandler) handleUniStream() (err error) {
 	}
 }
 
-type ServerAssocIdAddr struct {
-	assocId string
-	addr    net.Addr
-}
-
-func (a ServerAssocIdAddr) Network() string {
-	return "ServerAssocIdAddr"
-}
-
-func (a ServerAssocIdAddr) String() string {
-	return a.assocId
-}
-
-func (a ServerAssocIdAddr) RawAddr() net.Addr {
-	return a.addr
-}
-
 type serverUDPPacket struct {
 	pc     *quicStreamPacketConn
 	packet *Packet