Refactor: improve metadata structure

metadata/metadata.go

@@ -3,38 +3,15 @@ package metadata
 import (
 	"bytes"
 	"encoding/binary"
-	"fmt"
 	"net"
 	"strconv"
 
 	"github.com/xjasonlyu/tun2socks/v2/transport/socks5"
 )
 
-const (
+// Metadata contains metadata of transport protocol sessions.
-	TCP Network = iota
-	UDP
-)
-
-type Network uint8
-
-func (n Network) String() string {
-	switch n {
-	case TCP:
-		return "tcp"
-	case UDP:
-		return "udp"
-	default:
-		return fmt.Sprintf("network(%d)", n)
-	}
-}
-
-func (n Network) MarshalText() ([]byte, error) {
-	return []byte(n.String()), nil
-}
-
-// Metadata implements the net.Addr interface.
 type Metadata struct {
-	Net     Network `json:"network"`
+	Network Network `json:"network"`
 	SrcIP   net.IP  `json:"sourceIP"`
 	MidIP   net.IP  `json:"dialerIP"`
 	DstIP   net.IP  `json:"destinationIP"`
@@ -51,8 +28,15 @@ func (m *Metadata) SourceAddress() string {
 	return net.JoinHostPort(m.SrcIP.String(), strconv.FormatUint(uint64(m.SrcPort), 10))
 }
 
+func (m *Metadata) Addr() net.Addr {
+	if udpAddr := m.UDPAddr(); udpAddr != nil {
+		return udpAddr
+	}
+	return &Addr{metadata: m}
+}
+
 func (m *Metadata) UDPAddr() *net.UDPAddr {
-	if m.Net != UDP || m.DstIP == nil {
+	if m.Network != UDP || m.DstIP == nil {
 		return nil
 	}
 	return &net.UDPAddr{
@@ -78,12 +62,19 @@ func (m *Metadata) SerializeSocksAddr() socks5.Addr {
 	return bytes.Join(buf, nil)
 }
 
-func (m *Metadata) Network() string {
+// Addr implements the net.Addr interface.
-	return m.Net.String()
+type Addr struct {
+	metadata *Metadata
 }
 
-// String returns destination address of this metadata.
+func (a *Addr) Metadata() *Metadata {
-// Also, it implements net.Addr interface.
+	return a.metadata
-func (m *Metadata) String() string {
+}
-	return m.DestinationAddress()
+
+func (a *Addr) Network() string {
+	return a.metadata.Network.String()
+}
+
+func (a *Addr) String() string {
+	return a.metadata.DestinationAddress()
 }

metadata/network.go

@@ -0,0 +1,27 @@
+package metadata
+
+import (
+	"fmt"
+)
+
+const (
+	TCP Network = iota
+	UDP
+)
+
+type Network uint8
+
+func (n Network) String() string {
+	switch n {
+	case TCP:
+		return "tcp"
+	case UDP:
+		return "udp"
+	default:
+		return fmt.Sprintf("network(%d)", n)
+	}
+}
+
+func (n Network) MarshalText() ([]byte, error) {
+	return []byte(n.String()), nil
+}

proxy/direct.go

@@ -45,8 +45,8 @@ type directPacketConn struct {
 }
 
 func (pc *directPacketConn) WriteTo(b []byte, addr net.Addr) (int, error) {
-	if m, ok := addr.(*M.Metadata); ok && m.DstIP != nil {
+	if ma, ok := addr.(*M.Addr); ok && ma.Metadata().DstIP != nil {
-		return pc.PacketConn.WriteTo(b, m.UDPAddr())
+		return pc.PacketConn.WriteTo(b, ma.Metadata().UDPAddr())
 	}
 
 	udpAddr, err := net.ResolveUDPAddr("udp", addr.String())

proxy/shadowsocks.go

@@ -88,8 +88,8 @@ type ssPacketConn struct {
 
 func (pc *ssPacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
 	var packet []byte
-	if m, ok := addr.(*M.Metadata); ok {
+	if ma, ok := addr.(*M.Addr); ok {
-		packet, err = socks5.EncodeUDPPacket(m.SerializeSocksAddr(), b)
+		packet, err = socks5.EncodeUDPPacket(ma.Metadata().SerializeSocksAddr(), b)
 	} else {
 		packet, err = socks5.EncodeUDPPacket(socks5.ParseAddrToSocksAddr(addr), b)
 	}

proxy/socks5.go

@@ -141,8 +141,8 @@ type socksPacketConn struct {
 
 func (pc *socksPacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
 	var packet []byte
-	if m, ok := addr.(*M.Metadata); ok {
+	if ma, ok := addr.(*M.Addr); ok {
-		packet, err = socks5.EncodeUDPPacket(m.SerializeSocksAddr(), b)
+		packet, err = socks5.EncodeUDPPacket(ma.Metadata().SerializeSocksAddr(), b)
 	} else {
 		packet, err = socks5.EncodeUDPPacket(socks5.ParseAddrToSocksAddr(addr), b)
 	}

tunnel/tcp.go

@@ -30,7 +30,7 @@ func handleTCPConn(localConn core.TCPConn) {
 		dstIP, dstPort = parseAddr(localConn.LocalAddr())
 	)
 	metadata := &M.Metadata{
-		Net:     M.TCP,
+		Network: M.TCP,
 		SrcIP:   srcIP,
 		SrcPort: srcPort,
 		DstIP:   dstIP,

tunnel/udp.go

@@ -33,7 +33,7 @@ func handleUDPConn(uc core.UDPConn) {
 		dstIP, dstPort = parseAddr(uc.LocalAddr())
 	)
 	metadata := &M.Metadata{
-		Net:     M.UDP,
+		Network: M.UDP,
 		SrcIP:   srcIP,
 		SrcPort: srcPort,
 		DstIP:   dstIP,
@@ -50,8 +50,9 @@ func handleUDPConn(uc core.UDPConn) {
 	pc = newUDPTracker(pc, metadata)
 	defer pc.Close()
 
-	go handleUDPToRemote(uc, pc, metadata)
+	remote := metadata.Addr()
-	handleUDPToLocal(uc, pc, metadata)
+	go handleUDPToRemote(uc, pc, remote)
+	handleUDPToLocal(uc, pc, remote)
 }
 
 func handleUDPToRemote(uc core.UDPConn, pc net.PacketConn, remote net.Addr) {