[skip CI] Refactor UDP port forwarding and document usage

yggdrasil-network · Jul 19, 2024 · 475863e · 475863e
1 parent 0783b42
commit 475863e
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 43 deletions.
diff --git a/README.md b/README.md
@@ -72,7 +72,8 @@ other configuration such as listen addresses or multicast addresses, etc.
 
 ### Run Yggstack
 
-To run SOCKS proxy server listening on local port 1080 using generated configuration:
+To run SOCKS proxy server listening on local port 1080 using generated
+configuration (like `ssh -D`):
 
 ```
 ./yggstack -useconffile /path/to/yggdrasil.conf -socks 127.0.0.1:1080
@@ -84,11 +85,35 @@ To run SOCKS proxy server listening on UNIX socket file `/tmp/yggstack.sock`:
 ./yggstack -useconffile /path/to/yggdrasil.conf -socks /tmp/yggstack.sock
 ```
 
-To expose network services (like a Web server) listening on local port 8080 to Yggdrasil
-network address at port 80:
+To expose network services (like a Web server) listening on local port 8080
+to Yggdrasil network address at port 80 (like `ssh -R`):
+
+TCP:
+
+```
+./yggstack -useconffile /path/to/yggdrasil.conf -remote-tcp 80:127.0.0.1:8080
+```
+
+UDP:
+
+```
+./yggstack -useconffile /path/to/yggdrasil.conf -remote-udp 53:127.0.0.1:53
+```
+
+To forward remote port on some other Yggdrasil node to local machine (like `ssh -L`):
+
+TCP:
+
+```
+./yggstack -useconffile /path/to/yggdrasil.conf -local-tcp 127.0.0.1:8080:<remote-yggdrasil-ipv6>:8080
+./yggstack -useconffile /path/to/yggdrasil.conf -local-tcp [::1]:8080:<remote-yggdrasil-ipv6>:8080
+```
+
+UDP:
 
 ```
-./yggstack -useconffile /path/to/yggdrasil.conf -exposetcp 80:127.0.0.1:8080
+./yggstack -useconffile /path/to/yggdrasil.conf -local-udp 127.0.0.1:5353:<remote-yggdrasil-ipv6>:53
+./yggstack -useconffile /path/to/yggdrasil.conf -local-udp [::1]:5353:<remote-yggdrasil-ipv6>:53
 ```
 
 To run as a standalone node without SOCKS server or TCP port forwarding:

diff --git a/cmd/yggstack/main.go b/cmd/yggstack/main.go
@@ -31,6 +31,8 @@ import (
 
 	"github.com/yggdrasil-network/yggstack/src/netstack"
 	"github.com/yggdrasil-network/yggstack/src/types"
+
+	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 )
 
 type node struct {
@@ -40,13 +42,8 @@ type node struct {
 	socks5Listener net.Listener
 }
 
-type UDPConnSession struct {
-	conn       *net.UDPConn
-	remoteAddr net.Addr
-}
-
-type UDPPacketConnSession struct {
-	conn       *net.PacketConn
+type UDPSession struct {
+	conn       interface{}
 	remoteAddr net.Addr
 }
 
@@ -396,22 +393,23 @@ func main() {
 							logger.Errorf("Failed to connect to %s: %s", mapping.Mapped, err)
 							continue
 						}
-						udpSession := &UDPPacketConnSession{
-							conn:       &udpFwdConn,
+						udpSession := &UDPSession{
+							conn:       udpFwdConn,
 							remoteAddr: remoteUdpAddr,
 						}
 						localUdpConnections.Store(remoteUdpAddrStr, udpSession)
-						go types.ReverseProxyUDPPacketConn(mtu, udpListenConn, remoteUdpAddr, udpFwdConn)
+						go types.ReverseProxyUDP(mtu, udpListenConn, remoteUdpAddr, udpFwdConn)
 					}
 
-					udpSession, ok := connVal.(*UDPPacketConnSession)
+					udpSession, ok := connVal.(*UDPSession)
 					if !ok {
 						continue
 					}
 
-					udpFwdConn := *udpSession.conn
+					udpFwdConnPtr := udpSession.conn.(*gonet.UDPConn)
+					udpFwdConn := *udpFwdConnPtr
 
-					_, err = udpFwdConn.WriteTo(udpBuffer[:bytesRead], mapping.Mapped)
+					_, err = udpFwdConn.Write(udpBuffer[:bytesRead])
 					if err != nil {
 						logger.Debugf("Cannot write from yggdrasil to udp listener: %q", err)
 						udpFwdConn.Close()
@@ -482,23 +480,26 @@ func main() {
 							logger.Errorf("Failed to connect to %s: %s", mapping.Mapped, err)
 							continue
 						}
-						udpSession := &UDPConnSession{
+						udpSession := &UDPSession{
 							conn:       udpFwdConn,
 							remoteAddr: remoteUdpAddr,
 						}
 						remoteUdpConnections.Store(remoteUdpAddrStr, udpSession)
-						go types.ReverseProxyUDPConn(mtu, udpListenConn, remoteUdpAddr, *udpFwdConn)
+						go types.ReverseProxyUDP(mtu, udpListenConn, remoteUdpAddr, udpFwdConn)
 					}
 
-					udpSession, ok := connVal.(*UDPConnSession)
+					udpSession, ok := connVal.(*UDPSession)
 					if !ok {
 						continue
 					}
 
-					_, err = udpSession.conn.Write(udpBuffer[:bytesRead])
+					udpFwdConnPtr := udpSession.conn.(*net.UDPConn)
+					udpFwdConn := *udpFwdConnPtr
+
+					_, err = udpFwdConn.Write(udpBuffer[:bytesRead])
 					if err != nil {
 						logger.Debugf("Cannot write from yggdrasil to udp listener: %q", err)
-						udpSession.conn.Close()
+						udpFwdConn.Close()
 						remoteUdpConnections.Delete(remoteUdpAddrStr)
 						continue
 					}

diff --git a/src/netstack/netstack.go b/src/netstack/netstack.go
@@ -86,12 +86,12 @@ func (s *YggdrasilNetstack) DialContext(ctx context.Context, network, address st
 	}
 }
 
-func (s *YggdrasilNetstack) DialTCP(addr *net.TCPAddr) (net.Conn, error) {
+func (s *YggdrasilNetstack) DialTCP(addr *net.TCPAddr) (*gonet.TCPConn, error) {
 	fa, pn, _ := convertToFullAddr(addr.IP, addr.Port)
 	return gonet.DialTCP(s.stack, fa, pn)
 }
 
-func (s *YggdrasilNetstack) DialUDP(addr *net.UDPAddr) (net.PacketConn, error) {
+func (s *YggdrasilNetstack) DialUDP(addr *net.UDPAddr) (*gonet.UDPConn, error) {
 	fa, pn, _ := convertToFullAddr(addr.IP, addr.Port)
 	return gonet.DialUDP(s.stack, nil, &fa, pn)
 }
@@ -101,7 +101,7 @@ func (s *YggdrasilNetstack) ListenTCP(addr *net.TCPAddr) (net.Listener, error) {
 	return gonet.ListenTCP(s.stack, fa, pn)
 }
 
-func (s *YggdrasilNetstack) ListenUDP(addr *net.UDPAddr) (net.PacketConn, error) {
+func (s *YggdrasilNetstack) ListenUDP(addr *net.UDPAddr) (*gonet.UDPConn, error) {
 	fa, pn, _ := convertToFullAddr(addr.IP, addr.Port)
 	return gonet.DialUDP(s.stack, &fa, nil, pn)
 }
diff --git a/src/types/udpproxy.go b/src/types/udpproxy.go
@@ -4,7 +4,7 @@ import (
 	"net"
 )
 
-func ReverseProxyUDPConn(mtu uint64, dst net.PacketConn, dstAddr net.Addr, src net.UDPConn) error {
+func ReverseProxyUDP(mtu uint64, dst net.PacketConn, dstAddr net.Addr, src net.Conn) error {
 	buf := make([]byte, mtu)
 	for {
 		n, err := src.Read(buf[:])
@@ -20,20 +20,3 @@ func ReverseProxyUDPConn(mtu uint64, dst net.PacketConn, dstAddr net.Addr, src n
 	}
 	return nil
 }
-
-func ReverseProxyUDPPacketConn(mtu uint64, dst net.PacketConn, dstAddr net.Addr, src net.PacketConn) error {
-	buf := make([]byte, mtu)
-	for {
-		n, _, err := src.ReadFrom(buf[:])
-		if err != nil {
-			return err
-		}
-		if n > 0 {
-			n, err = dst.WriteTo(buf[:n], dstAddr)
-			if err != nil {
-				return err
-			}
-		}
-	}
-	return nil
-}