From 598e01c1e9b71bdfde832205cc7df0c2733b1de3 Mon Sep 17 00:00:00 2001
From: Zhi Wang <zwang@cs.fsu.edu>
Date: Sun, 9 Jan 2022 22:23:52 -0500
Subject: [PATCH] fix a deadlock

---
 reg.go   |  5 +++--
 relay.go | 15 +++++++++------
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/reg.go b/reg.go
index a5b580e..7f0dcbf 100644
--- a/reg.go
+++ b/reg.go
@@ -24,8 +24,9 @@ func (reg *Registry) init() {
 func (d *Registry) addDoer(name string, tc *TermConn) {
 	d.mtx.Lock()
 	if val, ok := d.doers[name]; ok {
-		log.Printf(name, "already exist in the dispatcher")
-		val.release()
+		log.Println(name, "already exist in the dispatcher", val, tc)
+		delete(d.doers, name)
+		val.release(false) // do not unregister in release, otherwise it is a deadlock
 	}
 	d.doers[name] = tc
 	d.mtx.Unlock()
diff --git a/relay.go b/relay.go
index accff5f..6ca7f62 100644
--- a/relay.go
+++ b/relay.go
@@ -159,10 +159,10 @@ func (tc *TermConn) ptyStdoutToWs() {
 		// handle viewers, we want to use non-blocking receive
 		select {
 		case watcher := <-tc.vchan:
-			log.Println("Received watcher", watcher)
+			log.Println("Received viewer", watcher)
 			viewers = append(viewers, watcher)
 		default:
-			log.Println("no watcher received")
+			//log.Println("no viewer received")
 		}
 
 		// We could add ws to viewers as well (then we can use io.MultiWriter),
@@ -195,15 +195,19 @@ func (tc *TermConn) ptyStdoutToWs() {
 		}
 	}
 
+	close(tc.done)
 	tc.ws.SetWriteDeadline(time.Now().Add(writeWait))
 	tc.ws.WriteMessage(websocket.CloseMessage,
 		websocket.FormatCloseMessage(websocket.CloseNormalClosure, "Pty closed"))
 	time.Sleep(closeGracePeriod)
 }
 
-func (tc *TermConn) release() {
+func (tc *TermConn) release(unreg bool) {
 	log.Println("releasing", tc.name)
-	registry.delDoer(tc.name)
+
+	if unreg {
+		registry.delDoer(tc.name)
+	}
 
 	if tc.ptmx != nil {
 		// cleanup the pty and its related process
@@ -231,7 +235,6 @@ func (tc *TermConn) release() {
 			log.Printf("Failed to wait for shell process(%v): %v", proc.Pid, err)
 		}
 
-		close(tc.done)
 		close(tc.vchan)
 	}
 
@@ -253,7 +256,7 @@ func wsHandleDoer(w http.ResponseWriter, r *http.Request) {
 		name: "main",
 	}
 
-	defer tc.release()
+	defer tc.release(true)
 	log.Println("\n\nCreated the websocket")
 
 	if err := tc.createPty(cmdToExec); err != nil {