Browse Source

Replaced the logic of when the tcp and sctp connections are closed to the original one.

Eliezer Croitoru 6 months ago
parent
commit
2c15aa34ab
2 changed files with 26 additions and 31 deletions
  1. 13 16
      sctp-to-tcp/main.go
  2. 13 15
      tcp-to-sctp/main.go

+ 13 - 16
sctp-to-tcp/main.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"os/signal"
 	"strconv"
+	"sync"
 	"syscall"
 
 	"github.com/asaskevich/govalidator"
@@ -24,6 +25,7 @@ var (
 	listenIP      string
 	listenPort    string
 	debug         int
+	// closeOnFirstDone int
 )
 
 func parseIP(s string) (net.IP, uint16, error) {
@@ -82,24 +84,17 @@ func copyConn(src net.Conn) {
 		}
 	}
 
-	done := make(chan struct{})
-
-	go func() {
-		defer src.Close()
-		defer dst.Close()
+	close := sync.Once{}
+	cp := func(dst io.WriteCloser, src io.ReadCloser) {
 		io.Copy(dst, src)
-		done <- struct{}{}
-	}()
-
-	go func() {
-		defer src.Close()
-		defer dst.Close()
-		io.Copy(src, dst)
-		done <- struct{}{}
-	}()
+		close.Do(func() {
+			dst.Close()
+			src.Close()
+		})
+	}
+	go cp(src, dst)
+	cp(dst, src)
 
-	<-done
-	<-done
 	if debug > 0 {
 		log.Println("Ended connection to:", remoteTCPIP+":"+remoteTCPPort, "From:", remoteAddress, "At:", localAddresss)
 	}
@@ -137,6 +132,8 @@ func init() {
 
 	flag.BoolVar(&dialWithPROXY, "proxy-protocol-connect", false, "Connect to the TCP host with a PROXY protocol header")
 
+	// flag.IntVar(&closeOnFirstDone, "close-on-first", 1, "The proxy will close the connection if one of the sides of the cnnection is closed instead of both")
+
 	flag.Parse()
 }
 

+ 13 - 15
tcp-to-sctp/main.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"os"
 	"os/signal"
+	"sync"
 	"syscall"
 
 	"github.com/ishidawataru/sctp"
@@ -19,6 +20,7 @@ var (
 	listenIP       string
 	listenPort     string
 	debug          int
+	// closeOnFirstDone int
 )
 
 func copyConn(src net.Conn) {
@@ -37,24 +39,18 @@ func copyConn(src net.Conn) {
 	if debug > 0 {
 		log.Printf("Dialed SCTP to %s", dst.RemoteAddr())
 	}
-	done := make(chan struct{})
 
-	go func() {
-		defer src.Close()
-		defer dst.Close()
+	close := sync.Once{}
+	cp := func(dst io.WriteCloser, src io.ReadCloser) {
 		io.Copy(dst, src)
-		done <- struct{}{}
-	}()
-
-	go func() {
-		defer src.Close()
-		defer dst.Close()
-		io.Copy(src, dst)
-		done <- struct{}{}
-	}()
+		close.Do(func() {
+			dst.Close()
+			src.Close()
+		})
+	}
+	go cp(src, dst)
+	cp(dst, src)
 
-	<-done
-	<-done
 	if debug > 0 {
 		log.Println("Ended connection to:", remoteSCTPIP+":"+remoteSCTPPort, "From:", remoteAddress, "At:", localAddresss)
 	}
@@ -91,6 +87,8 @@ func init() {
 	// flag.IntVar(&retries, "retries", 4, "The number of http and connect retries")
 	flag.IntVar(&debug, "debug", 0, "The Debug level of the service")
 
+	// flag.IntVar(&closeOnFirstDone, "close-on-first", 1, "The proxy will close the connection if one of the sides of the cnnection is closed instead of both")
+
 	// flag.BoolVar(&dialWithPROXY, "proxy-protocol-connect", false, "Connect to the SCTP host with a PROXY protocol header")
 
 	flag.Parse()