Browse Source

Added stop function to UDP proxsy

Toby Chui 9 months ago
parent
commit
a9738b8ad2
2 changed files with 44 additions and 8 deletions
  1. 2 0
      mod/tcpprox/conn.go
  2. 42 8
      mod/tcpprox/udpprox.go

+ 2 - 0
mod/tcpprox/conn.go

@@ -191,12 +191,14 @@ func (c *ProxyRelayConfig) IsRunning() bool {
 
 // Stop a running proxy if running
 func (c *ProxyRelayConfig) Stop() {
+	log.Println("[PROXY] Stopping Stream Proxy " + c.Name)
 	if c.Running || c.stopChan != nil {
 		c.stopChan <- true
 		time.Sleep(300 * time.Millisecond)
 		c.stopChan = nil
 		c.Running = false
 	}
+	log.Println("[PROXY] Stopped Stream Proxy " + c.Name)
 }
 
 /*

+ 42 - 8
mod/tcpprox/udpprox.go

@@ -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
 		}
+
 	}
 }