From 55328921b508293657b2d2959fd87ab6df538539 Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Wed, 5 Apr 2017 13:46:11 +0200
Subject: [PATCH] wait for done

---
 app/proxyman/mux/mux.go | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/app/proxyman/mux/mux.go b/app/proxyman/mux/mux.go
index 23ac4707..02ede296 100644
--- a/app/proxyman/mux/mux.go
+++ b/app/proxyman/mux/mux.go
@@ -182,6 +182,13 @@ func fetchInput(ctx context.Context, s *session, output buf.Writer) {
 	}
 }
 
+func waitForDone(ctx context.Context, s *session) {
+	<-ctx.Done()
+	s.closeUplink()
+	s.closeDownlink()
+	s.output.Close()
+}
+
 func (m *Client) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) bool {
 	m.access.Lock()
 	defer m.access.Unlock()
@@ -210,6 +217,7 @@ func (m *Client) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) bool
 	}
 	m.sessions[id] = s
 	go fetchInput(ctx, s, m.inboundRay.InboundInput())
+	go waitForDone(ctx, s)
 	return true
 }