Revert "passive tunnel termination"

This reverts commit bf1d0d6.
xiyan130227 · Mar 16, 2018 · 03f27ec · 03f27ec
1 parent b2390a5
commit 03f27ec
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 14 deletions.
diff --git a/client/main.go b/client/main.go
@@ -9,7 +9,6 @@ import (
 	"math/rand"
 	"net"
 	"os"
-	"sync"
 	"time"
 
 	"golang.org/x/crypto/pbkdf2"
@@ -69,12 +68,18 @@ func handleClient(sess *smux.Session, p1 io.ReadWriteCloser, quiet bool) {
 	}
 	defer p2.Close()
 
-	// tunnel setup
-	var wg sync.WaitGroup
-	wg.Add(2)
-	go func() { io.Copy(p1, p2); wg.Done() }()
-	go func() { io.Copy(p2, p1); wg.Done() }()
-	wg.Wait()
+	// start tunnel
+	p1die := make(chan struct{})
+	go func() { io.Copy(p1, p2); close(p1die) }()
+
+	p2die := make(chan struct{})
+	go func() { io.Copy(p2, p1); close(p2die) }()
+
+	// wait for tunnel termination
+	select {
+	case <-p1die:
+	case <-p2die:
+	}
 }
 
 func checkError(err error) {

diff --git a/server/main.go b/server/main.go
@@ -11,7 +11,6 @@ import (
 	"net/http"
 	_ "net/http/pprof"
 	"os"
-	"sync"
 	"time"
 
 	"golang.org/x/crypto/pbkdf2"
@@ -94,12 +93,18 @@ func handleClient(p1, p2 io.ReadWriteCloser, quiet bool) {
 	defer p1.Close()
 	defer p2.Close()
 
-	// tunnel setup
-	var wg sync.WaitGroup
-	wg.Add(2)
-	go func() { io.Copy(p1, p2); wg.Done() }()
-	go func() { io.Copy(p2, p1); wg.Done() }()
-	wg.Wait()
+	// start tunnel
+	p1die := make(chan struct{})
+	go func() { io.Copy(p1, p2); close(p1die) }()
+
+	p2die := make(chan struct{})
+	go func() { io.Copy(p2, p1); close(p2die) }()
+
+	// wait for tunnel termination
+	select {
+	case <-p1die:
+	case <-p2die:
+	}
 }
 
 func checkError(err error) {