|
@@ -1,9 +1,10 @@
|
|
|
package tcpprox
|
|
|
|
|
|
import (
|
|
|
- "fmt"
|
|
|
+ "errors"
|
|
|
"log"
|
|
|
"net"
|
|
|
+ "time"
|
|
|
)
|
|
|
|
|
|
/*
|
|
@@ -40,7 +41,7 @@ func initUDPConnections(listenAddr string, targetAddress string) (*net.UDPConn,
|
|
|
return nil, nil, err
|
|
|
}
|
|
|
|
|
|
- log.Println("Proxy serving on port %s\n", listenAddr)
|
|
|
+ log.Println("[UDP] Proxy listening on " + listenAddr)
|
|
|
|
|
|
outboundConn, err := net.ResolveUDPAddr("udp", targetAddress)
|
|
|
if err != nil {
|
|
@@ -51,12 +52,15 @@ func initUDPConnections(listenAddr string, targetAddress string) (*net.UDPConn,
|
|
|
}
|
|
|
|
|
|
// Go routine which manages connection from server to single client
|
|
|
-func (c *ProxyRelayConfig) RunUDPConnection(conn *udpClientServerConn, lisenter *net.UDPConn) {
|
|
|
+func (c *ProxyRelayConfig) RunUDPConnectionRelay(conn *udpClientServerConn, lisenter *net.UDPConn) {
|
|
|
var buffer [1500]byte
|
|
|
for {
|
|
|
// Read from server
|
|
|
n, err := conn.ServerConn.Read(buffer[0:])
|
|
|
if err != nil {
|
|
|
+ if errors.Is(err, net.ErrClosed) {
|
|
|
+ return
|
|
|
+ }
|
|
|
continue
|
|
|
}
|
|
|
// Relay it to client
|
|
@@ -64,25 +68,55 @@ func (c *ProxyRelayConfig) RunUDPConnection(conn *udpClientServerConn, lisenter
|
|
|
if err != nil {
|
|
|
continue
|
|
|
}
|
|
|
- log.Println("[UDP] Relayed '%s' from server to %s.\n",
|
|
|
- string(buffer[0:n]), conn.ClientAddr.String())
|
|
|
+
|
|
|
}
|
|
|
}
|
|
|
|
|
|
+// Close all connections that waiting for read from server
|
|
|
+func (c *ProxyRelayConfig) CloseAllUDPConnections() {
|
|
|
+ c.udpClientMap.Range(func(clientAddr, clientServerConn interface{}) bool {
|
|
|
+ conn := clientServerConn.(*udpClientServerConn)
|
|
|
+ conn.ServerConn.Close()
|
|
|
+ return true
|
|
|
+ })
|
|
|
+}
|
|
|
+
|
|
|
func (c *ProxyRelayConfig) ForwardUDP(address1, address2 string, stopChan chan bool) error {
|
|
|
//By default the incoming listen Address is int
|
|
|
//We need to add the loopback address into it
|
|
|
address1 = "127.0.0.1:" + address1
|
|
|
- fmt.Println(address1, address2)
|
|
|
lisener, targetAddr, err := initUDPConnections(address1, address2)
|
|
|
if err != nil {
|
|
|
return err
|
|
|
}
|
|
|
|
|
|
+ go func() {
|
|
|
+ //Stop channel receiver
|
|
|
+ for {
|
|
|
+ select {
|
|
|
+ case <-stopChan:
|
|
|
+ //Stop signal received
|
|
|
+ //Stop server -> client forwarder
|
|
|
+ c.CloseAllUDPConnections()
|
|
|
+ //Stop client -> server forwarder
|
|
|
+ //Force close, will terminate ReadFromUDP for inbound listener
|
|
|
+ lisener.Close()
|
|
|
+ return
|
|
|
+ default:
|
|
|
+ time.Sleep(100 * time.Millisecond)
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ }()
|
|
|
+
|
|
|
var buffer [1500]byte
|
|
|
for {
|
|
|
n, cliaddr, err := lisener.ReadFromUDP(buffer[0:])
|
|
|
if err != nil {
|
|
|
+ if errors.Is(err, net.ErrClosed) {
|
|
|
+ //Proxy stopped
|
|
|
+ return nil
|
|
|
+ }
|
|
|
continue
|
|
|
}
|
|
|
c.aTobAccumulatedByteTransfer.Add(int64(n))
|
|
@@ -95,10 +129,9 @@ func (c *ProxyRelayConfig) ForwardUDP(address1, address2 string, stopChan chan b
|
|
|
continue
|
|
|
}
|
|
|
c.udpClientMap.Store(saddr, conn)
|
|
|
- fmt.Println(saddr)
|
|
|
log.Println("[UDP] Created new connection for client " + saddr)
|
|
|
// Fire up routine to manage new connection
|
|
|
- go c.RunUDPConnection(conn, lisener)
|
|
|
+ go c.RunUDPConnectionRelay(conn, lisener)
|
|
|
|
|
|
} else {
|
|
|
log.Println("[UDP] Found connection for client " + saddr)
|
|
@@ -110,5 +143,6 @@ func (c *ProxyRelayConfig) ForwardUDP(address1, address2 string, stopChan chan b
|
|
|
if err != nil {
|
|
|
continue
|
|
|
}
|
|
|
+
|
|
|
}
|
|
|
}
|