From 0e33ab35f9dfe69b56eb100e613794ec744b13ef Mon Sep 17 00:00:00 2001
From: Gary Burd <gary@beagledreams.com>
Date: Sun, 18 Oct 2015 15:32:45 -0700
Subject: [PATCH] Fix race in command example

Close files in handler goroutine to avoid race.
Remove use of exec.Cmd. It's not adding anything.
---
 examples/command/home.html |   6 +-
 examples/command/main.go   | 118 +++++++++++++++++++++----------------
 2 files changed, 71 insertions(+), 53 deletions(-)

diff --git a/examples/command/home.html b/examples/command/home.html
index 9ddcb6e..72fd02b 100644
--- a/examples/command/home.html
+++ b/examples/command/home.html
@@ -37,7 +37,7 @@
             appendLog($("<div><b>Connection closed.</b></div>"))
         }
         conn.onmessage = function(evt) {
-            appendLog($("<div/>").text(evt.data))
+            appendLog($("<pre/>").text(evt.data))
         }
     } else {
         appendLog($("<div><b>Your browser does not support WebSockets.</b></div>"))
@@ -70,6 +70,10 @@ body {
     overflow: auto;
 }
 
+#log pre {
+  margin: 0;
+}
+
 #form {
     padding: 0 0.5em 0 0.5em;
     margin: 0;
diff --git a/examples/command/main.go b/examples/command/main.go
index b971fe8..ccd042d 100644
--- a/examples/command/main.go
+++ b/examples/command/main.go
@@ -20,6 +20,7 @@ import (
 
 var (
 	addr      = flag.String("addr", "127.0.0.1:8080", "http service address")
+	cmdPath   string
 	homeTempl = template.Must(template.ParseFiles("home.html"))
 )
 
@@ -31,49 +32,40 @@ const (
 	maxMessageSize = 8192
 )
 
-// connection is an middleman between the websocket connection and the command.
-type connection struct {
-	ws     *websocket.Conn
-	stdout io.ReadCloser
-	stdin  io.WriteCloser
-	cmd    *exec.Cmd
-}
-
-func (c *connection) pumpStdin() {
-	defer c.ws.Close()
-	c.ws.SetReadLimit(maxMessageSize)
+func pumpStdin(ws *websocket.Conn, w io.Writer) {
+	defer ws.Close()
+	ws.SetReadLimit(maxMessageSize)
 	for {
-		_, message, err := c.ws.ReadMessage()
+		_, message, err := ws.ReadMessage()
 		if err != nil {
 			break
 		}
 		message = append(message, '\n')
-		if _, err := c.stdin.Write(message); err != nil {
+		if _, err := w.Write(message); err != nil {
 			break
 		}
 	}
-	c.stdin.Close()
 	log.Println("exit stdin pump")
 }
 
-func (c *connection) pumpStdout() {
-	defer c.ws.Close()
-	s := bufio.NewScanner(c.stdout)
+func pumpStdout(ws *websocket.Conn, r io.Reader, done chan struct{}) {
+	defer ws.Close()
+	s := bufio.NewScanner(r)
 	for s.Scan() {
-		c.ws.SetWriteDeadline(time.Now().Add(writeWait))
-		if err := c.ws.WriteMessage(websocket.TextMessage, s.Bytes()); err != nil {
+		ws.SetWriteDeadline(time.Now().Add(writeWait))
+		if err := ws.WriteMessage(websocket.TextMessage, s.Bytes()); err != nil {
 			break
 		}
 	}
 	if s.Err() != nil {
 		log.Println("scan:", s.Err())
 	}
-	c.stdout.Close()
+	close(done)
 	log.Println("exit stdout pump")
 }
 
-func internalError(ws *websocket.Conn, fmt string, err error) {
-	log.Println(fmt, err)
+func internalError(ws *websocket.Conn, msg string, err error) {
+	log.Println(msg, err)
 	ws.WriteMessage(websocket.TextMessage, []byte("Internal server error."))
 }
 
@@ -87,49 +79,66 @@ func serveWs(w http.ResponseWriter, r *http.Request) {
 
 	ws, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
-		log.Println(err)
+		log.Println("upgrade:", err)
 		return
 	}
 
-	c := &connection{
-		cmd: exec.Command(flag.Args()[0], flag.Args()[1:]...),
-		ws:  ws,
-	}
+	defer ws.Close()
 
-	c.stdout, err = c.cmd.StdoutPipe()
+	outr, outw, err := os.Pipe()
 	if err != nil {
-		internalError(ws, "stdout: %v", err)
-		ws.Close()
+		internalError(ws, "stdout:", err)
+		return
+	}
+	defer outr.Close()
+	defer outw.Close()
+
+	inr, inw, err := os.Pipe()
+	if err != nil {
+		internalError(ws, "stdin:", err)
+		return
+	}
+	defer inr.Close()
+	defer inw.Close()
+
+	proc, err := os.StartProcess(cmdPath, flag.Args(), &os.ProcAttr{
+		Files: []*os.File{inr, outw, outw},
+	})
+	if err != nil {
+		internalError(ws, "start:", err)
 		return
 	}
 
-	c.stdin, err = c.cmd.StdinPipe()
-	if err != nil {
-		internalError(ws, "stdin: %v", err)
-		c.stdout.Close()
-		if closer, ok := c.cmd.Stdout.(io.Closer); ok {
-			closer.Close()
+	inr.Close()
+	outw.Close()
+
+	done := make(chan struct{})
+	go pumpStdout(ws, outr, done)
+
+	pumpStdin(ws, inw)
+
+	// Some commands will exit when stdin is closed.
+	inw.Close()
+
+	// Other comamnds need a bonk on the head.
+	if err := proc.Signal(os.Interrupt); err != nil {
+		log.Println("inter:", err)
+	}
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		// A bigger bonk on the head.
+		if err := proc.Signal(os.Kill); err != nil {
+			log.Println("term:", err)
 		}
-		ws.Close()
-		return
+		<-done
 	}
 
-	if err := c.cmd.Start(); err != nil {
-		internalError(ws, "start: %v", err)
-		c.stdout.Close()
-		c.stdin.Close()
-		ws.Close()
-		return
-	}
-
-	go c.pumpStdout()
-	c.pumpStdin()
-
-	c.cmd.Process.Signal(os.Interrupt)
-	if err := c.cmd.Wait(); err != nil {
+	if _, err := proc.Wait(); err != nil {
 		log.Println("wait:", err)
 	}
-	log.Println("exit serveWs")
+	log.Println("exiting handler")
 }
 
 func serveHome(w http.ResponseWriter, r *http.Request) {
@@ -150,6 +159,11 @@ func main() {
 	if len(flag.Args()) < 1 {
 		log.Fatal("must specify at least one argument")
 	}
+	var err error
+	cmdPath, err = exec.LookPath(flag.Args()[0])
+	if err != nil {
+		log.Fatal(err)
+	}
 	http.HandleFunc("/", serveHome)
 	http.HandleFunc("/ws", serveWs)
 	log.Fatal(http.ListenAndServe(*addr, nil))