tunnel del with active streams

2 files changed, 59 insertions, 39 deletions

diff --git a/pkg/server/socket/socket.go b/pkg/server/socket/socket.go
index cad1ad3..bf754cf 100644
--- a/pkg/server/socket/socket.go
+++ b/pkg/server/socket/socket.go
@@ -61,7 +61,7 @@ func (cc *connChannel) Recv(rq queue.Q) (err error) {
 
 func (cc *connChannel) String() string {
 	local, remote := cc.conn.LocalAddr(), cc.conn.RemoteAddr()
-	return fmt.Sprintf("%s/%s->%s", local.Network(), local, remote)
+	return fmt.Sprintf("%s/%s->%s", local.Network(), remote, local)
 }
 
 func (cc *connChannel) shutdown(err *error) {
diff --git a/pkg/server/tunnel.go b/pkg/server/tunnel.go
index 04f838b..e4a324c 100644
--- a/pkg/server/tunnel.go
+++ b/pkg/server/tunnel.go
@@ -5,7 +5,6 @@ import (
 	"tunnel/pkg/server/queue"
 	"tunnel/pkg/server/socket"
 	"tunnel/pkg/config"
-	"sync/atomic"
 	"strings"
 	"time"
 	"sort"
@@ -16,9 +15,9 @@ import (
 
 type stream struct {
 	id int
-	n int32
 	t *tunnel
 	since time.Time
+	wg sync.WaitGroup
 	in, out socket.Channel
 }
 
@@ -29,6 +28,7 @@ type tunnel struct {
 	streams map[int]*stream
 
 	mu sync.Mutex
+	wg sync.WaitGroup
 
 	nextSid int
 
@@ -47,12 +47,26 @@ func (t *tunnel) String() string {
 	return fmt.Sprintf("tunnel(%s)", t.id)
 }
 
-/* FIXME close streams also */
-func (t *tunnel) Close() {
+func (t *tunnel) stopServe() {
 	close(t.quit)
 	t.in.Close()
 	t.out.Close()
 	<-t.done
+}
+
+func (t *tunnel) stopStreams() {
+	t.mu.Lock()
+	for _, s := range t.streams {
+		s.stop()
+	}
+	t.mu.Unlock()
+
+	t.wg.Wait()
+}
+
+func (t *tunnel) Close() {
+	t.stopServe()
+	t.stopStreams()
 
 	log.Println(t, "delete")
 }
@@ -94,6 +108,21 @@ func (t *tunnel) serve() {
 	close(t.done)
 }
 
+func (t *tunnel) handle(in socket.Channel) {
+	out, err := t.out.Open()
+	if err != nil {
+		log.Println(t, err)
+		in.Close()
+		return
+	}
+
+	log.Println(t, "open", out)
+
+	s := t.newStream(in, out)
+
+	log.Println(t, s, "create", in, out)
+}
+
 func (t *tunnel) newStream(in, out socket.Channel) *stream {
 	s := &stream{
 		t: t,
@@ -108,33 +137,34 @@ func (t *tunnel) newStream(in, out socket.Channel) *stream {
 	t.streams[s.id] = s
 	t.mu.Unlock()
 
-	return s
-}
+	s.run()
 
-func (s *stream) ref() {
-	atomic.AddInt32(&s.n, 1)
-}
-
-func (s *stream) unref() {
-	if atomic.AddInt32(&s.n, -1) == 0 {
-		log.Println(s.t, s, "close")
+	go func () {
+		s.wg.Wait()
 
 		s.t.mu.Lock()
 		delete(s.t.streams, s.id)
 		s.t.mu.Unlock()
-	}
+
+		s.t.wg.Done()
+
+		log.Println(s.t, s, "close")
+	}()
+
+	return s
 }
 
 func (s *stream) watchChannel(rq, wq queue.Q, c socket.Channel) {
 	watch := func (q queue.Q, f func (q queue.Q) error) {
-		s.ref()
-		defer s.unref()
+		defer s.wg.Done()
 
 		if err := f(q); err != nil {
 			log.Println(s.t, s, err)
 		}
 	}
 
+	s.wg.Add(2)
+
 	go func () {
 		watch(wq, c.Send)
 		close(wq)
@@ -144,9 +174,10 @@ func (s *stream) watchChannel(rq, wq queue.Q, c socket.Channel) {
 }
 
 func (s *stream) watchPipe(rq, wq queue.Q, f func (rq, wq queue.Q) error) {
+	s.wg.Add(1)
+
 	go func () {
-		s.ref()
-		defer s.unref()
+		defer s.wg.Done()
 
 		if err := f(rq, wq); err != nil {
 			log.Println(s.t, s, err)
@@ -156,28 +187,14 @@ func (s *stream) watchPipe(rq, wq queue.Q, f func (rq, wq queue.Q) error) {
 	}()
 }
 
-func (t *tunnel) handle(in socket.Channel) {
-	log.Println(t, "handle")
-
-	out, err := t.out.Open()
-	if err != nil {
-		log.Println(t, err)
-		in.Close()
-		return
-	}
-
-	log.Println(t, "open", out)
-
-	s := t.newStream(in, out)
-
-	s.ref()
-	defer s.unref()
+func (s *stream) run() {
+	s.t.wg.Add(1)
 
 	rq, wq := queue.New(), queue.New()
 
-	s.watchChannel(rq, wq, in)
+	s.watchChannel(rq, wq, s.in)
 
-	for _, m := range t.m {
+	for _, m := range s.t.m {
 		send, recv := m.Open()
 		if send != nil {
 			q := queue.New()
@@ -191,9 +208,12 @@ func (t *tunnel) handle(in socket.Channel) {
 		}
 	}
 
-	s.watchChannel(wq, rq, out)
+	s.watchChannel(wq, rq, s.out)
+}
 
-	log.Println(t, s, "create", in, out)
+func (s *stream) stop() {
+	s.in.Close()
+	s.out.Close()
 }
 
 func newTunnel(args []string) (*tunnel, error) {