From 9b552f55a4d0e00be44a147eb60c8be18556dce0 Mon Sep 17 00:00:00 2001
From: Qian Wang <cbeuw.andy@gmail.com>
Date: Wed, 24 Jul 2019 14:25:57 +0100
Subject: [PATCH] Allow an terminal message to be set as a session terminates

---
 cmd/ck-server/ck-server.go    |  2 +-
 internal/multiplex/session.go | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/cmd/ck-server/ck-server.go b/cmd/ck-server/ck-server.go
index 4d594d2..d97deef 100644
--- a/cmd/ck-server/ck-server.go
+++ b/cmd/ck-server/ck-server.go
@@ -186,7 +186,7 @@ func dispatchConnection(conn net.Conn, sta *server.State) {
 			newStream, err := sesh.Accept()
 			if err != nil {
 				if err == mux.ErrBrokenSession {
-					log.Printf("Session closed for UID:%v, sessionID:%v\n", b64.EncodeToString(UID), sessionID)
+					log.Printf("Session closed for UID:%v, sessionID:%v, reason:%v\n", b64.EncodeToString(UID), sessionID, sesh.TerminalMsg())
 					user.DelSession(sessionID)
 					return
 				} else {
diff --git a/internal/multiplex/session.go b/internal/multiplex/session.go
index 5b78d74..d8bc76c 100644
--- a/internal/multiplex/session.go
+++ b/internal/multiplex/session.go
@@ -41,6 +41,8 @@ type Session struct {
 	broken  uint32
 	die     chan struct{}
 	suicide sync.Once
+
+	terminalMsg atomic.Value
 }
 
 func MakeSession(id uint32, valve *Valve, obfs Obfser, deobfs Deobfser, obfsedRead func(net.Conn, []byte) (int, error)) *Session {
@@ -124,6 +126,19 @@ func (sesh *Session) getStream(id uint32, closingFrame bool) *Stream {
 	}
 }
 
+func (sesh *Session) SetTerminalMsg(msg string) {
+	sesh.terminalMsg.Store(msg)
+}
+
+func (sesh *Session) TerminalMsg() string {
+	msg := sesh.terminalMsg.Load()
+	if msg != nil {
+		return msg.(string)
+	} else {
+		return ""
+	}
+}
+
 func (sesh *Session) Close() error {
 	// Because closing a closed channel causes panic
 	sesh.suicide.Do(func() { close(sesh.die) })