From f4a15a06b0836272a188e3626d01e22c2b59c6fc Mon Sep 17 00:00:00 2001
From: Zhi Wang <zwang@cs.fsu.edu>
Date: Sun, 9 Jan 2022 08:56:43 -0500
Subject: [PATCH] one runner works

---
 assets/index.html |  20 ++--
 assets/main.js    |   4 +-
 main.go           | 255 ++++---------------------------------------
 relay.go          | 272 ++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 308 insertions(+), 243 deletions(-)
 create mode 100644 relay.go
diff --git a/assets/index.html b/assets/index.html
index 2930b44..4af0a57 100644
--- a/assets/index.html
+++ b/assets/index.html
@@ -8,26 +8,26 @@
   <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
   <link href="https://fonts.googleapis.com/css2?family=Fira+Code&family=Fira+Sans&display=swap" rel="stylesheet">
 
-  <script src="xterm.js"></script>
-  <script src="xterm-addon-attach.js"></script>
-  <script src="xterm-addon-fit.js"></script>
-  <script src="xterm-addon-web-links.js"></script>
+  <script src="/assets/external/xterm.js"></script>
+  <script src="/assets/external/xterm-addon-attach.js"></script>
+  <script src="/assets/external/xterm-addon-fit.js"></script>
+  <script src="/assets/external/xterm-addon-web-links.js"></script>
 
-  <script src="main.js"></script>
+  <script src="/assets/main.js"></script>
 
-  <link rel="stylesheet" href="xterm.css" />
-  <link rel="stylesheet" href="main.css" />
+  <link rel="stylesheet" href="/assets/external/xterm.css" />
+  <link rel="stylesheet" href="/assets/main.css" />
 
-  <title>Websocket Terminal</title>
+  <title>{{.title}}</title>
 </head>
 
 <body>
-  <h2>Web Terminal</h2>
+  <h2>{{.title}}</h2>
   <div id="terminal">
     <div id="terminal_view"></div>
   </div>
   <script>
-    term = createTerminal();
+    term = createTerminal({{.path}});
     // print something to test output and scroll
     var str = [
       ' ┌────────────────────────────────────────────────────────────────────────────┐',
diff --git a/assets/main.js b/assets/main.js
index 3df8894..c3ca192 100644
--- a/assets/main.js
+++ b/assets/main.js
@@ -1,4 +1,4 @@
-function createTerminal() {
+function createTerminal(path) {
   // vscode-snazzy https://github.com/Tyriar/vscode-snazzy
   // copied from xterm.js website
   var baseTheme = {
@@ -43,7 +43,7 @@ function createTerminal() {
   fitAddon.fit();
 
   // create the websocket and connect to the server
-  const ws_uri = "wss://" + window.location.host + "/ws";
+  const ws_uri = "wss://" + window.location.host + path;
   const socket = new WebSocket(ws_uri);
   const attachAddon = new AttachAddon.AttachAddon(socket);
   term.loadAddon(attachAddon);
diff --git a/main.go b/main.go
index 81f7f78..f44861b 100644
--- a/main.go
+++ b/main.go
@@ -1,237 +1,16 @@
 package main
 
 import (
-	"fmt"
 	"log"
 	"net/http"
-	"net/url"
 	"os"
-	"os/exec"
-	"strings"
-	"time"
 
-	"github.com/creack/pty"
 	"github.com/gin-gonic/gin"
-	"github.com/gorilla/websocket"
 )
 
-const (
-	// Time allowed to write a message to the peer.
-	writeWait = 5 * time.Second
-
-	// Time allowed to read the next pong message from the peer.
-	pongWait = 30 * time.Second
-
-	// Maximum message size allowed from peer.
-	maxMessageSize = 8192
-
-	// Send pings to peer with this period. Must be less than pongWait.
-	pingPeriod = (pongWait * 9) / 10
-
-	// Time to wait before force close on connection.
-	closeGracePeriod = 10 * time.Second
-)
-
-func createPty(cmdline []string) (*os.File, *exec.Cmd, error) {
-	// Create a shell command.
-	cmd := exec.Command(cmdline[0], cmdline[1:]...)
-
-	// Start the command with a pty.
-	ptmx, err := pty.Start(cmd)
-
-	if err != nil {
-		return nil, nil, err
-	}
-
-	// Use fixed size, the xterm is initalized as 122x37,
-	// But we set pty to 120x36. Using fullsize will lead
-	// some program to misbehaive.
-	pty.Setsize(ptmx, &pty.Winsize{
-		Cols: 120,
-		Rows: 36,
-	})
-
-	log.Printf("Create shell process %v (%v)", cmdline, cmd.Process.Pid)
-	return ptmx, cmd, nil
-}
-
-var host *string = nil
-
-var upgrader = websocket.Upgrader{
-	ReadBufferSize:  4096,
-	WriteBufferSize: 4096,
-	CheckOrigin: func(r *http.Request) bool {
-		org := r.Header.Get("Origin")
-		h, err := url.Parse(org)
-
-		if err != nil {
-			return false
-		}
-
-		if (host == nil) || (*host != h.Host) {
-			log.Println("Failed origin check of ", org)
-		}
-
-		return (host != nil) && (*host == h.Host)
-	},
-}
-
-// Periodically send ping message to detect the status of the ws
-func ping(ws *websocket.Conn, done chan struct{}) {
-	ticker := time.NewTicker(pingPeriod)
-	defer ticker.Stop()
-	for {
-		select {
-		case <-ticker.C:
-			err := ws.WriteControl(websocket.PingMessage,
-				[]byte{}, time.Now().Add(writeWait))
-
-			if err != nil {
-				log.Println("Failed to write ping message:", err)
-			}
-
-		case <-done:
-			log.Println("Exit ping routine as stdout is going away")
-			return
-		}
-	}
-}
-
-// shovel data from websocket to pty stdin
-func toPtyStdin(ws *websocket.Conn, ptmx *os.File) {
-	ws.SetReadLimit(maxMessageSize)
-
-	// set the readdeadline. The idea here is simple,
-	// as long as we keep receiving pong message,
-	// the readdeadline will keep updating. Otherwise
-	// read will timeout.
-	ws.SetReadDeadline(time.Now().Add(pongWait))
-	ws.SetPongHandler(func(string) error {
-		ws.SetReadDeadline(time.Now().Add(pongWait))
-		return nil
-	})
-
-	for {
-		_, buf, err := ws.ReadMessage()
-
-		if err != nil {
-			log.Println("Failed to receive data from ws:", err)
-			break
-		}
-
-		_, err = ptmx.Write(buf)
-
-		if err != nil {
-			log.Println("Failed to send data to pty stdin: ", err)
-			break
-		}
-	}
-}
-
-// shovel data from pty Stdout to WS
-func fromPtyStdout(ws *websocket.Conn, ptmx *os.File, done chan struct{}) {
-	readBuf := make([]byte, 4096)
-
-	for {
-		n, err := ptmx.Read(readBuf)
-
-		if err != nil {
-			log.Println("Failed to read from pty stdout: ", err)
-			break
-		}
-
-		ws.SetWriteDeadline(time.Now().Add(writeWait))
-		if err = ws.WriteMessage(websocket.BinaryMessage, readBuf[:n]); err != nil {
-			log.Println("Failed to write message: ", err)
-			break
-		}
-	}
-
-	close(done)
-
-	ws.SetWriteDeadline(time.Now().Add(writeWait))
-	ws.WriteMessage(websocket.CloseMessage,
-		websocket.FormatCloseMessage(websocket.CloseNormalClosure, "Pty closed"))
-	time.Sleep(closeGracePeriod)
-}
-
+// command line options
 var cmdToExec = []string{"bash"}
 
-// handle websockets
-func wsHandler(w http.ResponseWriter, r *http.Request) {
-	ws, err := upgrader.Upgrade(w, r, nil)
-
-	if err != nil {
-		log.Println("Failed to create websocket: ", err)
-		return
-	}
-
-	defer ws.Close()
-
-	log.Println("\n\nCreated the websocket")
-
-	ptmx, cmd, err := createPty(cmdToExec)
-
-	if err != nil {
-		log.Println("Failed to create PTY: ", err)
-		return
-	}
-
-	done := make(chan struct{})
-
-	go fromPtyStdout(ws, ptmx, done)
-	go ping(ws, done)
-
-	toPtyStdin(ws, ptmx)
-
-	// cleanup the pty and its related process
-	ptmx.Close()
-	proc := cmd.Process
-
-	// send an interrupt, this will cause the shell process to
-	// return from syscalls if any is pending
-	if err := proc.Signal(os.Interrupt); err != nil {
-		log.Printf("Failed to send Interrupt to shell process(%v): %v ", proc.Pid, err)
-	}
-
-	// Wait for a second for shell process to interrupt before kill it
-	time.Sleep(time.Second)
-
-	log.Printf("Try to kill the shell process(%v)", proc.Pid)
-
-	if err := proc.Signal(os.Kill); err != nil {
-		log.Printf("Failed to send KILL to shell process(%v): %v", proc.Pid, err)
-	}
-
-	if _, err := proc.Wait(); err != nil {
-		log.Printf("Failed to wait for shell process(%v): %v", proc.Pid, err)
-	}
-}
-
-// return files
-func fileHandler(c *gin.Context, fname string) {
-	// if the URL has no fname, c.Param returns "/"
-	if fname == "/" {
-		fname = "/index.html"
-		host = &c.Request.Host
-	}
-
-	fname = fname[1:] //fname always starts with /
-	log.Println("Sending ", fname)
-
-	if strings.HasSuffix(fname, "html") {
-		c.HTML(200, fname, nil)
-	} else {
-		//c.HTML interprets the file as HTML file
-		//we do not need that for regular files
-		if strings.HasPrefix(fname, "xterm") {
-			c.File(fmt.Sprint("./assets/external/", fname))
-		} else {
-			c.File(fmt.Sprint("./assets/", fname))
-		}
-	}
-}
-
 func main() {
 	fp, err := os.OpenFile("web_term.log", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
 
@@ -257,16 +36,30 @@ func main() {
 	rt.SetTrustedProxies(nil)
 	rt.LoadHTMLGlob("./assets/*.html")
 
-	rt.GET("/*fname", func(c *gin.Context) {
-		fname := c.Param("fname")
+	rt.GET("/watch/*sname", func(c *gin.Context) {
+		c.HTML(http.StatusOK, "index.html", gin.H{
+			"title": "Watcher terminal",
+			"path":  "/ws_watch",
+		})
+	})
 
-		// ws is a special case to create a new websocket
-		switch fname {
-		case "/ws":
-			wsHandler(c.Writer, c.Request)
-		default:
-			fileHandler(c, fname)
-		}
+	rt.GET("/ws_run", func(c *gin.Context) {
+		wsHandler(c.Writer, c.Request, false)
+	})
+
+	rt.GET("/ws_watch", func(c *gin.Context) {
+		wsHandler(c.Writer, c.Request, true)
+	})
+
+	// handle static files
+	rt.Static("/assets", "./assets")
+
+	rt.GET("/", func(c *gin.Context) {
+		c.HTML(http.StatusOK, "index.html", gin.H{
+			"title": "Master terminal",
+			"path":  "/ws_run ",
+		})
+		host = &c.Request.Host
 	})
 
 	rt.RunTLS(":8080", "./tls/cert.pem", "./tls/private-key.pem")
diff --git a/relay.go b/relay.go
new file mode 100644
index 0000000..52ed1f2
--- /dev/null
+++ b/relay.go
@@ -0,0 +1,272 @@
+//This file contains code to relay traffic between websocket and pty
+package main
+
+import (
+	"log"
+	"net/http"
+	"net/url"
+	"os"
+	"os/exec"
+	"time"
+
+	"github.com/creack/pty"
+	"github.com/gorilla/websocket"
+)
+
+const (
+	// Time allowed to write a message to the peer.
+	writeWait = 5 * time.Second
+
+	// Time allowed to read the next pong message from the peer.
+	pongWait = 30 * time.Second
+
+	// Maximum message size allowed from peer.
+	maxMessageSize = 8192
+
+	// Send pings to peer with this period. Must be less than pongWait.
+	pingPeriod = (pongWait * 9) / 10
+
+	// Time to wait before force close on connection.
+	closeGracePeriod = 10 * time.Second
+)
+
+func createPty(cmdline []string) (*os.File, *exec.Cmd, error) {
+	// Create a shell command.
+	cmd := exec.Command(cmdline[0], cmdline[1:]...)
+
+	// Start the command with a pty.
+	ptmx, err := pty.Start(cmd)
+
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// Use fixed size, the xterm is initalized as 122x37,
+	// But we set pty to 120x36. Using fullsize will lead
+	// some program to misbehaive.
+	pty.Setsize(ptmx, &pty.Winsize{
+		Cols: 120,
+		Rows: 36,
+	})
+
+	log.Printf("Create shell process %v (%v)", cmdline, cmd.Process.Pid)
+	return ptmx, cmd, nil
+}
+
+var host *string = nil
+
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  4096,
+	WriteBufferSize: 4096,
+	CheckOrigin: func(r *http.Request) bool {
+		org := r.Header.Get("Origin")
+		h, err := url.Parse(org)
+
+		if err != nil {
+			return false
+		}
+
+		if (host == nil) || (*host != h.Host) {
+			log.Println("Failed origin check of ", org)
+		}
+
+		return (host != nil) && (*host == h.Host)
+	},
+}
+
+// Periodically send ping message to detect the status of the ws
+func ping(ws *websocket.Conn, done chan struct{}) {
+	ticker := time.NewTicker(pingPeriod)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			err := ws.WriteControl(websocket.PingMessage,
+				[]byte{}, time.Now().Add(writeWait))
+
+			if err != nil {
+				log.Println("Failed to write ping message:", err)
+			}
+
+		case <-done:
+			log.Println("Exit ping routine as stdout is going away")
+			return
+		}
+	}
+}
+
+// shovel data from websocket to pty stdin
+func toPtyStdin(ws *websocket.Conn, ptmx *os.File) {
+	ws.SetReadLimit(maxMessageSize)
+
+	// set the readdeadline. The idea here is simple,
+	// as long as we keep receiving pong message,
+	// the readdeadline will keep updating. Otherwise
+	// read will timeout.
+	ws.SetReadDeadline(time.Now().Add(pongWait))
+	ws.SetPongHandler(func(string) error {
+		ws.SetReadDeadline(time.Now().Add(pongWait))
+		return nil
+	})
+
+	for {
+		_, buf, err := ws.ReadMessage()
+
+		if err != nil {
+			log.Println("Failed to receive data from ws:", err)
+			break
+		}
+
+		_, err = ptmx.Write(buf)
+
+		if err != nil {
+			log.Println("Failed to send data to pty stdin: ", err)
+			break
+		}
+	}
+}
+
+// shovel data from pty Stdout to WS
+func fromPtyStdout(ws *websocket.Conn, ptmx *os.File, done chan struct{}) {
+	var watchers []*websocket.Conn
+	readBuf := make([]byte, 4096)
+
+	for {
+		n, err := ptmx.Read(readBuf)
+
+		if err != nil {
+			log.Println("Failed to read from pty stdout: ", err)
+			break
+		}
+
+		// handle watchers, we want to use non-blocking receive
+		select {
+		case watcher := <-watcherChan:
+			log.Println("Received watcher", watcher)
+			watchers = append(watchers, watcher)
+
+		default:
+			log.Println("no watcher received")
+		}
+
+		// We could add ws to watchers as well, but we want to handle it
+		// differently if there is an error
+		ws.SetWriteDeadline(time.Now().Add(writeWait))
+		if ws.WriteMessage(websocket.BinaryMessage, readBuf[0:n]); err != nil {
+			log.Println("Failed to write message: ", err)
+			break
+		}
+
+		for i, w := range watchers {
+			if w == nil {
+				continue
+			}
+
+			if err = w.WriteMessage(websocket.BinaryMessage, readBuf[0:n]); err != nil {
+				log.Println("Failed to write message to watcher: ", err)
+
+				watchers[i] = nil
+				w.Close()
+			}
+		}
+	}
+
+	close(done)
+	close(watcherChan)
+	watcherChan = nil
+
+	// close the watcher
+	for _, w := range watchers {
+		if w != nil {
+			w.Close()
+		}
+	}
+
+	ws.SetWriteDeadline(time.Now().Add(writeWait))
+	ws.WriteMessage(websocket.CloseMessage,
+		websocket.FormatCloseMessage(websocket.CloseNormalClosure, "Pty closed"))
+	time.Sleep(closeGracePeriod)
+}
+
+var watcherChan chan *websocket.Conn
+
+// handle websockets
+func wsHandleRun(w http.ResponseWriter, r *http.Request) {
+	ws, err := upgrader.Upgrade(w, r, nil)
+
+	if err != nil {
+		log.Println("Failed to create websocket: ", err)
+		return
+	}
+
+	defer ws.Close()
+
+	log.Println("\n\nCreated the websocket")
+
+	ptmx, cmd, err := createPty(cmdToExec)
+
+	if err != nil {
+		log.Println("Failed to create PTY: ", err)
+		return
+	}
+
+	done := make(chan struct{})
+	watcherChan = make(chan *websocket.Conn)
+
+	go fromPtyStdout(ws, ptmx, done)
+	go ping(ws, done)
+
+	toPtyStdin(ws, ptmx)
+
+	// cleanup the pty and its related process
+	ptmx.Close()
+	proc := cmd.Process
+
+	// send an interrupt, this will cause the shell process to
+	// return from syscalls if any is pending
+	if err := proc.Signal(os.Interrupt); err != nil {
+		log.Printf("Failed to send Interrupt to shell process(%v): %v ", proc.Pid, err)
+	}
+
+	// Wait for a second for shell process to interrupt before kill it
+	time.Sleep(time.Second)
+
+	log.Printf("Try to kill the shell process(%v)", proc.Pid)
+
+	if err := proc.Signal(os.Kill); err != nil {
+		log.Printf("Failed to send KILL to shell process(%v): %v", proc.Pid, err)
+	}
+
+	if _, err := proc.Wait(); err != nil {
+		log.Printf("Failed to wait for shell process(%v): %v", proc.Pid, err)
+	}
+}
+
+// handle websockets
+func wsHandleWatcher(w http.ResponseWriter, r *http.Request) {
+	ws, err := upgrader.Upgrade(w, r, nil)
+
+	if err != nil {
+		log.Println("Failed to create websocket: ", err)
+		return
+	}
+
+	log.Println("\n\nCreated the websocket")
+
+	if watcherChan == nil {
+		log.Println("No active runner, create a runner first")
+		ws.Close()
+		return
+	}
+
+	// hand the websocket to runner
+	watcherChan <- ws
+}
+
+func wsHandler(w http.ResponseWriter, r *http.Request, isWatcher bool) {
+	if !isWatcher {
+		wsHandleRun(w, r)
+	} else {
+		wsHandleWatcher(w, r)
+	}
+}
