From 04e376991e3ba82b181c6ad3f86d7a465f7cd679 Mon Sep 17 00:00:00 2001
From: Vasyl Gello <vasek.gello@gmail.com>
Date: Tue, 6 Aug 2024 15:35:26 +0300
Subject: [PATCH] Fix UDP port forward not working with Yggdrasil client

Signed-off-by: Vasyl Gello <vasek.gello@gmail.com>
---
 cmd/yggstack/main.go | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/cmd/yggstack/main.go b/cmd/yggstack/main.go
index 19fbad9..771efe9 100644
--- a/cmd/yggstack/main.go
+++ b/cmd/yggstack/main.go
@@ -465,13 +465,16 @@ func main() {
 				for {
 					bytesRead, remoteUdpAddr, err := udpListenConn.ReadFrom(udpBuffer)
 					if err != nil {
-						if bytesRead == 0 {
-							continue
-						}
+						logger.Debugf("udp readFrom error: %v", err)
+					}
+					if bytesRead == 0 {
+						continue
 					}
 
 					remoteUdpAddrStr := remoteUdpAddr.String()
 
+					var udpSession *UDPSession = nil
+
 					connVal, ok := remoteUdpConnections.Load(remoteUdpAddrStr)
 
 					if !ok {
@@ -481,23 +484,24 @@ func main() {
 							logger.Errorf("Failed to connect to %s: %s", mapping.Mapped, err)
 							continue
 						}
-						udpSession := &UDPSession{
+						udpSession = &UDPSession{
 							conn:       udpFwdConn,
 							remoteAddr: remoteUdpAddr,
 						}
 						remoteUdpConnections.Store(remoteUdpAddrStr, udpSession)
 						go types.ReverseProxyUDP(mtu, udpListenConn, remoteUdpAddr, udpFwdConn)
-					}
+					} else {
+						udpSession, ok = connVal.(*UDPSession)
 
-					udpSession, ok := connVal.(*UDPSession)
-					if !ok {
-						continue
+						if !ok {
+							continue
+						}
 					}
 
 					udpFwdConnPtr := udpSession.conn.(*net.UDPConn)
 					udpFwdConn := *udpFwdConnPtr
 
-					_, err = udpFwdConn.Write(udpBuffer[:bytesRead])
+					bytesWritten, err := udpFwdConn.Write(udpBuffer[:bytesRead])
 					if err != nil {
 						logger.Debugf("Cannot write from yggdrasil to udp listener: %q", err)
 						udpFwdConn.Close()