Fix proxies to set the listeners properly

RicYaben · RicYaben · commit 2fcd01383bef · 2022-09-05T23:14:22.000+02:00
diff --git a/internal/proxy/manager.go b/internal/proxy/manager.go
@@ -24,6 +24,7 @@ type ProxyManager interface {
 
 	// Methods for proxies the using ID field
 	GetProxy(id string) (Proxy, error)
+	SetProxy(pe Proxy) (Proxy, error)
 	DeleteProxy(id string) error
 
 	// Wrapper method to find a proxy using the port and protocol
@@ -81,6 +82,23 @@ func (pm *ProxyManagerItem) GetProxy(id string) (pe Proxy, err error) {
 	return
 }
 
+func (pm *ProxyManagerItem) SetProxy(px Proxy) (pe Proxy, err error) {
+	// Get all the proxies registered
+	proxies := pm.GetProxies()
+
+	for ind, proxy := range proxies {
+		if proxy.GetID() == px.GetID() {
+			// Replace the index of the proxy with the new one
+			proxies[ind] = px
+			pm.proxies = proxies
+			return
+		}
+	}
+
+	err = fmt.Errorf("proxy not found")
+	return
+}
+
 // Delete a proxy from teh registered list using the ID
 func (pm *ProxyManagerItem) DeleteProxy(id string) (err error) {
 	// Get all the proxies registered
diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
@@ -57,13 +57,13 @@ type AbstractProxy struct {
 	wg sync.WaitGroup
 
 	// Generic listener
-	listener interface{ Close() }
+	listener interface{ Close() error }
 }
 
 // Function to stop the proxy from runing
 func (pe *AbstractProxy) Stop() (err error) {
 	// Stop the proxy if it is still alive
-	if pe.GetStatus() != globals.StoppedStatus {
+	if pe.GetStatus() == globals.RunningStatus {
 		close(pe.stop)
 		pe.listener.Close()
 		// Wait for all the connections and the server to stop
diff --git a/internal/proxy/tcpproxy.go b/internal/proxy/tcpproxy.go
@@ -28,6 +28,7 @@ func (tcpProxy *TCPProxy) Start() (err error) {
 
 	// Get the listener or create a new one
 	listener, err := tcpProxy.GetListener()
+
 	if err != nil {
 		return
 	}
@@ -98,6 +99,7 @@ func (tcpProxy *TCPProxy) GetListener() (listener net.Listener, err error) {
 
 func (tcpProxy *TCPProxy) NewListener() (listener net.Listener, err error) {
 	listener, err = net.Listen(tcpProxy.GetNetwork().String(), fmt.Sprintf(":%d", tcpProxy.GetPort()))
+	tcpProxy.AbstractProxy.listener = listener
 	return
 }
 
diff --git a/internal/proxy/udpproxy.go b/internal/proxy/udpproxy.go
@@ -93,6 +93,7 @@ func (udpProxy *UDPProxy) GetListener() (listener *net.UDPConn, err error) {
 			return
 		}
 		udpProxy.listener = listener
+		udpProxy.AbstractProxy.listener = listener
 	}
 
 	return

Original file line number	Diff line number	Diff line change
`@@ -28,6 +28,7 @@ func (tcpProxy *TCPProxy) Start() (err error) {`
`28`	`28`
`29`	`29`	`// Get the listener or create a new one`
`30`	`30`	`listener, err := tcpProxy.GetListener()`
	`31`	`+`
`31`	`32`	`if err != nil {`
`32`	`33`	`return`
`33`	`34`	`}`
`@@ -98,6 +99,7 @@ func (tcpProxy *TCPProxy) GetListener() (listener net.Listener, err error) {`
`98`	`99`
`99`	`100`	`func (tcpProxy *TCPProxy) NewListener() (listener net.Listener, err error) {`
`100`	`101`	`listener, err = net.Listen(tcpProxy.GetNetwork().String(), fmt.Sprintf(":%d", tcpProxy.GetPort()))`
	`102`	`+ tcpProxy.AbstractProxy.listener = listener`
`101`	`103`	`return`
`102`	`104`	`}`
`103`	`105`
Original file line number	Diff line number	Diff line change
`@@ -93,6 +93,7 @@ func (udpProxy UDPProxy) GetListener() (listener net.UDPConn, err error) {`
`93`	`93`	`return`
`94`	`94`	`}`
`95`	`95`	`udpProxy.listener = listener`
	`96`	`+ udpProxy.AbstractProxy.listener = listener`
`96`	`97`	`}`
`97`	`98`
`98`	`99`	`return`