From 2cf809021bba4637243cc8f432d4a35ca105587d Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Thu, 13 Apr 2017 20:56:32 +0200
Subject: [PATCH] force close session manage when main connection is finished

---
 app/proxyman/mux/mux.go     |  3 +++
 app/proxyman/mux/session.go | 18 ++++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/app/proxyman/mux/mux.go b/app/proxyman/mux/mux.go
index 5e876311..16cb6ae9 100644
--- a/app/proxyman/mux/mux.go
+++ b/app/proxyman/mux/mux.go
@@ -121,6 +121,7 @@ func (m *Client) monitor() {
 	for {
 		select {
 		case <-m.ctx.Done():
+			m.sessionManager.Close()
 			m.inboundRay.InboundInput().Close()
 			m.inboundRay.InboundOutput().CloseError()
 			return
@@ -302,6 +303,8 @@ func handle(ctx context.Context, s *Session, output buf.Writer) {
 func (w *ServerWorker) run(ctx context.Context) {
 	input := w.outboundRay.OutboundInput()
 	reader := NewReader(input)
+
+	defer w.sessionManager.Close()
 L:
 	for {
 		select {
diff --git a/app/proxyman/mux/session.go b/app/proxyman/mux/session.go
index f66065ae..963ae5fd 100644
--- a/app/proxyman/mux/session.go
+++ b/app/proxyman/mux/session.go
@@ -74,10 +74,26 @@ func (m *SessionManager) CloseIfNoSession() bool {
 	m.RLock()
 	defer m.RUnlock()
 
+	if m.closed {
+		return true
+	}
+
 	if len(m.sessions) == 0 {
 		return false
 	}
 
+	m.closed = true
+	return true
+}
+
+func (m *SessionManager) Close() {
+	m.RLock()
+	defer m.RUnlock()
+
+	if m.closed {
+		return
+	}
+
 	m.closed = true
 
 	for _, s := range m.sessions {
@@ -86,8 +102,6 @@ func (m *SessionManager) CloseIfNoSession() bool {
 	}
 
 	m.sessions = make(map[uint16]*Session)
-
-	return true
 }
 
 type Session struct {