Added wip UDP proxy

mod/tcpprox/conn.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"strconv"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -38,8 +39,12 @@ func isReachable(target string) bool {
 	return true
 }
 
-func connCopy(conn1 net.Conn, conn2 net.Conn, wg *sync.WaitGroup, accumulator *int64) {
-	io.Copy(conn1, conn2)
+func connCopy(conn1 net.Conn, conn2 net.Conn, wg *sync.WaitGroup, accumulator *atomic.Int64) {
+	n, err := io.Copy(conn1, conn2)
+	if err != nil {
+		//Add to accumulator
+		accumulator.Add(n)
+	}
 	conn1.Close()
 	log.Println("[←]", "close the connect at local:["+conn1.LocalAddr().String()+"] and remote:["+conn1.RemoteAddr().String()+"]")
 	//conn2.Close()
@@ -47,7 +52,7 @@ func connCopy(conn1 net.Conn, conn2 net.Conn, wg *sync.WaitGroup, accumulator *i
 	wg.Done()
 }
 
-func forward(conn1 net.Conn, conn2 net.Conn, aTob *int64, bToa *int64) {
+func forward(conn1 net.Conn, conn2 net.Conn, aTob *atomic.Int64, bToa *atomic.Int64) {
 	log.Printf("[+] start transmit. [%s],[%s] <-> [%s],[%s] \n", conn1.LocalAddr().String(), conn1.RemoteAddr().String(), conn2.LocalAddr().String(), conn2.RemoteAddr().String())
 	var wg sync.WaitGroup
 	// wait tow goroutines
@@ -158,6 +163,8 @@ func (c *ProxyRelayConfig) Start() error {
 			err = c.Port2port(c.PortA, c.PortB, stopChan)
 		} else if c.Mode == ProxyMode_Starter {
 			err = c.Host2host(c.PortA, c.PortB, stopChan)
+		} else if c.Mode == ProxyMode_UDP {
+			err = c.ForwardUDP(c.PortA, c.PortB, stopChan)
 		}
 		if err != nil {
 			c.Running = false

mod/tcpprox/tcpprox.go

@@ -3,6 +3,7 @@ package tcpprox
 import (
 	"errors"
 	"net"
+	"sync/atomic"
 
 	"github.com/google/uuid"
 	"imuslab.com/zoraxy/mod/database"
@@ -20,6 +21,7 @@ const (
 	ProxyMode_Listen    = 0
 	ProxyMode_Transport = 1
 	ProxyMode_Starter   = 2
+	ProxyMode_UDP       = 3
 )
 
 type ProxyRelayOptions struct {
@@ -31,16 +33,16 @@ type ProxyRelayOptions struct {
 }
 
 type ProxyRelayConfig struct {
-	UUID                        string    //A UUIDv4 representing this config
-	Name                        string    //Name of the config
-	Running                     bool      //If the service is running
-	PortA                       string    //Ports A (config depends on mode)
-	PortB                       string    //Ports B (config depends on mode)
-	Mode                        int       //Operation Mode
-	Timeout                     int       //Timeout for connection in sec
-	stopChan                    chan bool //Stop channel to stop the listener
-	aTobAccumulatedByteTransfer int64     //Accumulated byte transfer from A to B
-	bToaAccumulatedByteTransfer int64     //Accumulated byte transfer from B to A
+	UUID                        string       //A UUIDv4 representing this config
+	Name                        string       //Name of the config
+	Running                     bool         //If the service is running
+	PortA                       string       //Ports A (config depends on mode)
+	PortB                       string       //Ports B (config depends on mode)
+	Mode                        int          //Operation Mode
+	Timeout                     int          //Timeout for connection in sec
+	stopChan                    chan bool    //Stop channel to stop the listener
+	aTobAccumulatedByteTransfer atomic.Int64 //Accumulated byte transfer from A to B
+	bToaAccumulatedByteTransfer atomic.Int64 //Accumulated byte transfer from B to A
 
 	parent *Manager `json:"-"`
 }
@@ -94,6 +96,11 @@ func NewTCProxy(options *Options) *Manager {
 }
 
 func (m *Manager) NewConfig(config *ProxyRelayOptions) string {
+	//Generate two zero value for atomic int64
+	aAcc := atomic.Int64{}
+	bAcc := atomic.Int64{}
+	aAcc.Store(0)
+	bAcc.Store(0)
 	//Generate a new config from options
 	configUUID := uuid.New().String()
 	thisConfig := ProxyRelayConfig{
@@ -105,8 +112,8 @@ func (m *Manager) NewConfig(config *ProxyRelayOptions) string {
 		Mode:                        config.Mode,
 		Timeout:                     config.Timeout,
 		stopChan:                    nil,
-		aTobAccumulatedByteTransfer: 0,
-		bToaAccumulatedByteTransfer: 0,
+		aTobAccumulatedByteTransfer: aAcc,
+		bToaAccumulatedByteTransfer: bAcc,
 
 		parent: m,
 	}

mod/tcpprox/udpprox.go

@@ -0,0 +1,89 @@
+package tcpprox
+
+import (
+	"log"
+	"net"
+)
+
+/*
+	UDP Proxy Module
+*/
+
+// Information maintained for each client/server connection
+type udpClientServerConn struct {
+	ClientAddr *net.UDPAddr // Address of the client
+	ServerConn *net.UDPConn // UDP connection to server
+}
+
+// Generate a new connection by opening a UDP connection to the server
+func createNewUDPConn(srvAddr, cliAddr *net.UDPAddr) *udpClientServerConn {
+	conn := new(udpClientServerConn)
+	conn.ClientAddr = cliAddr
+	srvudp, err := net.DialUDP("udp", nil, srvAddr)
+	if err != nil {
+		return nil
+	}
+	conn.ServerConn = srvudp
+	return conn
+}
+
+// Start listener, return inbound lisener and proxy target UDP address
+func initUDPConnections(listenAddr string, targetAddress string) (*net.UDPConn, *net.UDPAddr, error) {
+	// Set up Proxy
+	saddr, err := net.ResolveUDPAddr("udp", listenAddr)
+	if err != nil {
+		return nil, nil, err
+	}
+	inboundConn, err := net.ListenUDP("udp", saddr)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	log.Println("Proxy serving on port %s\n", listenAddr)
+
+	outboundConn, err := net.ResolveUDPAddr("udp", targetAddress)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return inboundConn, outboundConn, nil
+}
+
+func (c *ProxyRelayConfig) ForwardUDP(address1, address2 string, stopChan chan bool) error {
+	lisener, targetAddr, err := initUDPConnections(address1, address2)
+	if err != nil {
+		return err
+	}
+
+	var buffer [1500]byte
+	for {
+		n, cliaddr, err := lisener.ReadFromUDP(buffer[0:])
+		if err != nil {
+			continue
+		}
+		c.aTobAccumulatedByteTransfer.Add(int64(n))
+		saddr := cliaddr.String()
+		dlock()
+		conn, found := ClientDict[saddr]
+		if !found {
+			conn = createNewUDPConn(targetAddr, cliaddr)
+			if conn == nil {
+				dunlock()
+				continue
+			}
+			ClientDict[saddr] = conn
+			dunlock()
+			Vlogf(2, "Created new connection for client %s\n", saddr)
+			// Fire up routine to manage new connection
+			go RunConnection(conn)
+		} else {
+			Vlogf(5, "Found connection for client %s\n", saddr)
+			dunlock()
+		}
+		// Relay to server
+		_, err = conn.ServerConn.Write(buffer[0:n])
+		if checkreport(1, err) {
+			continue
+		}
+	}
+}