refactor: implement UDP frame read/write functions for improved data handling

Author: mkyla

internal/common.go

@@ -326,6 +326,40 @@ func (c *Common) decode(data []byte) ([]byte, error) {
 	return c.xor(decoded), nil
 }
 
+func writeUDPFrame(w net.Conn, data []byte) error {
+	length := len(data)
+	if length > 65535 {
+		return fmt.Errorf("writeUDPFrame: datagram too large: %d", length)
+	}
+	header := [2]byte{byte(length >> 8), byte(length)}
+	if _, err := w.Write(header[:]); err != nil {
+		return err
+	}
+	_, err := w.Write(data)
+	return err
+}
+
+func readUDPFrame(conn net.Conn, buf []byte, timeout time.Duration) (int, error) {
+	if timeout > 0 {
+		conn.SetReadDeadline(time.Now().Add(timeout))
+	}
+	var header [2]byte
+	if _, err := io.ReadFull(conn, header[:]); err != nil {
+		return 0, err
+	}
+	length := int(header[0])<<8 | int(header[1])
+	if length == 0 {
+		return 0, nil
+	}
+	if length > len(buf) {
+		return 0, fmt.Errorf("readUDPFrame: datagram too large: %d > buffer %d", length, len(buf))
+	}
+	if _, err := io.ReadFull(conn, buf[:length]); err != nil {
+		return 0, err
+	}
+	return length, nil
+}
+
 func (c *Common) resolve(network, address string) (any, error) {
 	now := time.Now()
 
@@ -1389,10 +1423,9 @@ func (c *Common) commonUDPLoop() {
 
 				buffer := c.getUDPBuffer()
 				defer c.putUDPBuffer(buffer)
-				reader := &conn.TimeoutReader{Conn: remoteConn, Timeout: udpReadTimeout}
 
 				for c.ctx.Err() == nil {
-					x, err := reader.Read(buffer)
+					x, err := readUDPFrame(remoteConn, buffer, udpReadTimeout)
 					if err != nil {
 						if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
 							c.logger.Debug("UDP session abort: %v", err)
@@ -1426,10 +1459,9 @@ func (c *Common) commonUDPLoop() {
 			c.logger.Debug("Starting transfer: %v <-> %v", remoteConn.LocalAddr(), c.targetUDPConn.LocalAddr())
 		}
 
-		_, err = remoteConn.Write(buffer[:x])
-		if err != nil {
+		if err = writeUDPFrame(remoteConn, buffer[:x]); err != nil {
 			if err != io.EOF {
-				c.logger.Error("commonUDPLoop: write to tunnel failed: %v", err)
+				c.logger.Error("tunnelUDPLoop: write to tunnel failed: %v", err)
 			}
 			c.targetUDPSession.Delete(sessionKey)
 			remoteConn.Close()
@@ -1704,10 +1736,9 @@ func (c *Common) commonUDPOnce(signal Signal) {
 
 		buffer := c.getUDPBuffer()
 		defer c.putUDPBuffer(buffer)
-		reader := &conn.TimeoutReader{Conn: remoteConn, Timeout: udpReadTimeout}
 
 		for c.ctx.Err() == nil {
-			x, err := reader.Read(buffer)
+			x, err := readUDPFrame(remoteConn, buffer, udpReadTimeout)
 			if err != nil {
 				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
 					c.logger.Debug("UDP session abort: %v", err)
@@ -1747,10 +1778,9 @@ func (c *Common) commonUDPOnce(signal Signal) {
 				return
 			}
 
-			_, err = remoteConn.Write(buffer[:x])
-			if err != nil {
+			if err = writeUDPFrame(remoteConn, buffer[:x]); err != nil {
 				if err != io.EOF {
-					c.logger.Error("commonUDPOnce: write to tunnel failed: %v", err)
+					c.logger.Error("tunnelUDPOnce: write to tunnel failed: %v", err)
 				}
 				return
 			}