From 61b1031da64f25db39fa95fb38df0245d939e9df Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Sun, 6 Dec 2020 10:50:45 +0000
Subject: [PATCH] Reduce code duplication in session closing

---
 internal/multiplex/session.go | 37 +++++++++++++++--------------------
 1 file changed, 16 insertions(+), 21 deletions(-)

diff --git a/internal/multiplex/session.go b/internal/multiplex/session.go
index d86fd4e..e1ab275 100644
--- a/internal/multiplex/session.go
+++ b/internal/multiplex/session.go
@@ -253,8 +253,7 @@ func (sesh *Session) TerminalMsg() string {
 	}
 }
 
-func (sesh *Session) passiveClose() error {
-	log.Debugf("attempting to passively close session %v", sesh.id)
+func (sesh *Session) closeSession(closeSwitchboard bool) error {
 	if atomic.SwapUint32(&sesh.closed, 1) == 1 {
 		log.Debugf("session %v has already been closed", sesh.id)
 		return errRepeatSessionClosing
@@ -273,7 +272,18 @@ func (sesh *Session) passiveClose() error {
 		return true
 	})
 
-	sesh.sb.closeAll()
+	if closeSwitchboard {
+		sesh.sb.closeAll()
+	}
+	return nil
+}
+
+func (sesh *Session) passiveClose() error {
+	log.Debugf("attempting to passively close session %v", sesh.id)
+	err := sesh.closeSession(true)
+	if err != nil {
+		return err
+	}
 	log.Debugf("session %v closed gracefully", sesh.id)
 	return nil
 }
@@ -288,25 +298,10 @@ func genRandomPadding() []byte {
 
 func (sesh *Session) Close() error {
 	log.Debugf("attempting to actively close session %v", sesh.id)
-	if atomic.SwapUint32(&sesh.closed, 1) == 1 {
-		log.Debugf("session %v has already been closed", sesh.id)
-		return errRepeatSessionClosing
+	err := sesh.closeSession(false)
+	if err != nil {
+		return err
 	}
-	sesh.acceptCh <- nil
-
-	// close all streams
-	sesh.streams.Range(func(key, streamI interface{}) bool {
-		if streamI == nil {
-			return true
-		}
-		stream := streamI.(*Stream)
-		atomic.StoreUint32(&stream.closed, 1)
-		_ = stream.recvBuf.Close() // will not block
-		sesh.streams.Delete(key)
-		sesh.streamCountDecr()
-		return true
-	})
-
 	// we send a notice frame telling remote to close the session
 	pad := genRandomPadding()
 	f := &Frame{