vendor/github.com/gorilla/websocket/examples/command/main.go

   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 != "GET" {
 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 }