main.go

Documentation: github.com/gorilla/websocket/examples/command

     1  // Copyright 2015 The Gorilla WebSocket Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"bufio"
     9  	"flag"
    10  	"io"
    11  	"log"
    12  	"net/http"
    13  	"os"
    14  	"os/exec"
    15  	"time"
    16  
    17  	"github.com/gorilla/websocket"
    18  )
    19  
    20  var (
    21  	addr    = flag.String("addr", "127.0.0.1:8080", "http service address")
    22  	cmdPath string
    23  )
    24  
    25  const (
    26  	// Time allowed to write a message to the peer.
    27  	writeWait = 10 * time.Second
    28  
    29  	// Maximum message size allowed from peer.
    30  	maxMessageSize = 8192
    31  
    32  	// Time allowed to read the next pong message from the peer.
    33  	pongWait = 60 * time.Second
    34  
    35  	// Send pings to peer with this period. Must be less than pongWait.
    36  	pingPeriod = (pongWait * 9) / 10
    37  
    38  	// Time to wait before force close on connection.
    39  	closeGracePeriod = 10 * time.Second
    40  )
    41  
    42  func pumpStdin(ws *websocket.Conn, w io.Writer) {
    43  	defer ws.Close()
    44  	ws.SetReadLimit(maxMessageSize)
    45  	ws.SetReadDeadline(time.Now().Add(pongWait))
    46  	ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
    47  	for {
    48  		_, message, err := ws.ReadMessage()
    49  		if err != nil {
    50  			break
    51  		}
    52  		message = append(message, '\n')
    53  		if _, err := w.Write(message); err != nil {
    54  			break
    55  		}
    56  	}
    57  }
    58  
    59  func pumpStdout(ws *websocket.Conn, r io.Reader, done chan struct{}) {
    60  	defer func() {
    61  	}()
    62  	s := bufio.NewScanner(r)
    63  	for s.Scan() {
    64  		ws.SetWriteDeadline(time.Now().Add(writeWait))
    65  		if err := ws.WriteMessage(websocket.TextMessage, s.Bytes()); err != nil {
    66  			ws.Close()
    67  			break
    68  		}
    69  	}
    70  	if s.Err() != nil {
    71  		log.Println("scan:", s.Err())
    72  	}
    73  	close(done)
    74  
    75  	ws.SetWriteDeadline(time.Now().Add(writeWait))
    76  	ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
    77  	time.Sleep(closeGracePeriod)
    78  	ws.Close()
    79  }
    80  
    81  func ping(ws *websocket.Conn, done chan struct{}) {
    82  	ticker := time.NewTicker(pingPeriod)
    83  	defer ticker.Stop()
    84  	for {
    85  		select {
    86  		case <-ticker.C:
    87  			if err := ws.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait)); err != nil {
    88  				log.Println("ping:", err)
    89  			}
    90  		case <-done:
    91  			return
    92  		}
    93  	}
    94  }
    95  
    96  func internalError(ws *websocket.Conn, msg string, err error) {
    97  	log.Println(msg, err)
    98  	ws.WriteMessage(websocket.TextMessage, []byte("Internal server error."))
    99  }
   100  
   101  var upgrader = websocket.Upgrader{}
   102  
   103  func serveWs(w http.ResponseWriter, r *http.Request) {
   104  	ws, err := upgrader.Upgrade(w, r, nil)
   105  	if err != nil {
   106  		log.Println("upgrade:", err)
   107  		return
   108  	}
   109  
   110  	defer ws.Close()
   111  
   112  	outr, outw, err := os.Pipe()
   113  	if err != nil {
   114  		internalError(ws, "stdout:", err)
   115  		return
   116  	}
   117  	defer outr.Close()
   118  	defer outw.Close()
   119  
   120  	inr, inw, err := os.Pipe()
   121  	if err != nil {
   122  		internalError(ws, "stdin:", err)
   123  		return
   124  	}
   125  	defer inr.Close()
   126  	defer inw.Close()
   127  
   128  	proc, err := os.StartProcess(cmdPath, flag.Args(), &os.ProcAttr{
   129  		Files: []*os.File{inr, outw, outw},
   130  	})
   131  	if err != nil {
   132  		internalError(ws, "start:", err)
   133  		return
   134  	}
   135  
   136  	inr.Close()
   137  	outw.Close()
   138  
   139  	stdoutDone := make(chan struct{})
   140  	go pumpStdout(ws, outr, stdoutDone)
   141  	go ping(ws, stdoutDone)
   142  
   143  	pumpStdin(ws, inw)
   144  
   145  	// Some commands will exit when stdin is closed.
   146  	inw.Close()
   147  
   148  	// Other commands need a bonk on the head.
   149  	if err := proc.Signal(os.Interrupt); err != nil {
   150  		log.Println("inter:", err)
   151  	}
   152  
   153  	select {
   154  	case <-stdoutDone:
   155  	case <-time.After(time.Second):
   156  		// A bigger bonk on the head.
   157  		if err := proc.Signal(os.Kill); err != nil {
   158  			log.Println("term:", err)
   159  		}
   160  		<-stdoutDone
   161  	}
   162  
   163  	if _, err := proc.Wait(); err != nil {
   164  		log.Println("wait:", err)
   165  	}
   166  }
   167  
   168  func serveHome(w http.ResponseWriter, r *http.Request) {
   169  	if r.URL.Path != "/" {
   170  		http.Error(w, "Not found", http.StatusNotFound)
   171  		return
   172  	}
   173  	if r.Method != http.MethodGet {
   174  		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
   175  		return
   176  	}
   177  	http.ServeFile(w, r, "home.html")
   178  }
   179  
   180  func main() {
   181  	flag.Parse()
   182  	if len(flag.Args()) < 1 {
   183  		log.Fatal("must specify at least one argument")
   184  	}
   185  	var err error
   186  	cmdPath, err = exec.LookPath(flag.Args()[0])
   187  	if err != nil {
   188  		log.Fatal(err)
   189  	}
   190  	http.HandleFunc("/", serveHome)
   191  	http.HandleFunc("/ws", serveWs)
   192  	log.Fatal(http.ListenAndServe(*addr, nil))
   193  }
   194
View as plain text