Make UDP conn duplex

Author: selvakn

client/client.go

@@ -6,8 +6,8 @@ import (
 	"io"
 	"net"
 	"p2p-port-forward/constants"
-	"p2p-port-forward/utils"
 	"p2p-port-forward/logger"
+	"p2p-port-forward/utils"
 )
 
 var log = logger.Logger
@@ -56,6 +56,8 @@ func (c *Client) listenUDP() func() (net.Conn, error) {
 	return func() (net.Conn, error) {
 		addr, _ := net.ResolveUDPAddr(c.networkProto.GetName(), fmt.Sprintf(":%s", c.port))
 		fmt.Println("Listening....")
-		return net.ListenUDP(c.networkProto.GetName(), addr)
+		conn, err := net.ListenUDP(c.networkProto.GetName(), addr)
+		return &utils.DuplexUDPConnection{UDPConn: conn}, err
+
 	}
 }

main.go

@@ -6,8 +6,8 @@ import (
 	"io"
 	"p2p-port-forward/client"
 	"p2p-port-forward/logger"
-	"p2p-port-forward/utils"
 	"p2p-port-forward/server"
+	"p2p-port-forward/utils"
 )
 
 var (

server/server.go

@@ -1,13 +1,13 @@
 package server
 
 import (
-	"io"
+	"fmt"
 	"github.com/selvakn/libzt"
+	"io"
 	"net"
 	"p2p-port-forward/constants"
-	"p2p-port-forward/utils"
-	"fmt"
 	"p2p-port-forward/logger"
+	"p2p-port-forward/utils"
 )
 
 var log = logger.Logger

utils/duplex_udp_connection.go

@@ -0,0 +1,28 @@
+package utils
+
+import (
+	"net"
+	"time"
+)
+
+type DuplexUDPConnection struct {
+	UDPConn  *net.UDPConn
+	rUDPAddr *net.UDPAddr
+}
+
+func (c *DuplexUDPConnection) Read(b []byte) (n int, err error) {
+	len, rUDPAddr, err := c.UDPConn.ReadFromUDP(b)
+	c.rUDPAddr = rUDPAddr
+	return len, err
+}
+
+func (c *DuplexUDPConnection) Write(b []byte) (n int, err error) {
+	return c.UDPConn.WriteToUDP(b, c.rUDPAddr)
+}
+
+func (c *DuplexUDPConnection) Close() error                       { return c.UDPConn.Close() }
+func (c *DuplexUDPConnection) LocalAddr() net.Addr                { return c.UDPConn.LocalAddr() }
+func (c *DuplexUDPConnection) RemoteAddr() net.Addr               { return c.UDPConn.RemoteAddr() }
+func (c *DuplexUDPConnection) SetDeadline(t time.Time) error      { return c.UDPConn.SetDeadline(t) }
+func (c *DuplexUDPConnection) SetReadDeadline(t time.Time) error  { return c.UDPConn.SetReadDeadline(t) }
+func (c *DuplexUDPConnection) SetWriteDeadline(t time.Time) error { return c.UDPConn.SetReadDeadline(t) }

utils/io.go

@@ -29,10 +29,22 @@ func Sync(stream1 func() (net.Conn, error), stream2 func() (net.Conn, error), li
 
 func sync(source1 io.ReadWriteCloser, source2 io.ReadWriteCloser) {
 	go func() {
-		defer source2.Close()
-		defer source1.Close()
+		defer closeAll(source2, source1)
 
-		io.Copy(source2, source1)
+		_, err := io.Copy(source2, source1)
+		if err != nil {
+			log.Error(err)
+		}
 	}()
-	io.Copy(source1, source2)
+	_, err := io.Copy(source1, source2)
+	if err != nil {
+		log.Error(err)
+	}
+}
+
+func closeAll(sources ...io.Closer) {
+	log.Info("Closing all connections")
+	for _, source := range sources {
+		source.Close()
+	}
 }