websocket/examples/command/main.go

171 lines
3.4 KiB
Go
Raw Normal View History

2015-10-18 23:32:14 +03:00
// Copyright 2015 The Gorilla WebSocket Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package main
import (
"bufio"
"flag"
"io"
"log"
"net/http"
"os"
"os/exec"
"text/template"
"time"
"github.com/gorilla/websocket"
)
var (
addr = flag.String("addr", "127.0.0.1:8080", "http service address")
cmdPath string
2015-10-18 23:32:14 +03:00
homeTempl = template.Must(template.ParseFiles("home.html"))
)
const (
// Time allowed to write a message to the peer.
writeWait = 10 * time.Second
// Maximum message size allowed from peer.
maxMessageSize = 8192
)
func pumpStdin(ws *websocket.Conn, w io.Writer) {
defer ws.Close()
ws.SetReadLimit(maxMessageSize)
2015-10-18 23:32:14 +03:00
for {
_, message, err := ws.ReadMessage()
2015-10-18 23:32:14 +03:00
if err != nil {
break
}
message = append(message, '\n')
if _, err := w.Write(message); err != nil {
2015-10-18 23:32:14 +03:00
break
}
}
log.Println("exit stdin pump")
}
func pumpStdout(ws *websocket.Conn, r io.Reader, done chan struct{}) {
defer ws.Close()
s := bufio.NewScanner(r)
2015-10-18 23:32:14 +03:00
for s.Scan() {
ws.SetWriteDeadline(time.Now().Add(writeWait))
if err := ws.WriteMessage(websocket.TextMessage, s.Bytes()); err != nil {
2015-10-18 23:32:14 +03:00
break
}
}
if s.Err() != nil {
log.Println("scan:", s.Err())
}
close(done)
2015-10-18 23:32:14 +03:00
log.Println("exit stdout pump")
}
func internalError(ws *websocket.Conn, msg string, err error) {
log.Println(msg, err)
2015-10-18 23:32:14 +03:00
ws.WriteMessage(websocket.TextMessage, []byte("Internal server error."))
}
var upgrader = websocket.Upgrader{}
func serveWs(w http.ResponseWriter, r *http.Request) {
if r.Method != "GET" {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
ws, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println("upgrade:", err)
2015-10-18 23:32:14 +03:00
return
}
defer ws.Close()
2015-10-18 23:32:14 +03:00
outr, outw, err := os.Pipe()
2015-10-18 23:32:14 +03:00
if err != nil {
internalError(ws, "stdout:", err)
2015-10-18 23:32:14 +03:00
return
}
defer outr.Close()
defer outw.Close()
2015-10-18 23:32:14 +03:00
inr, inw, err := os.Pipe()
2015-10-18 23:32:14 +03:00
if err != nil {
internalError(ws, "stdin:", err)
2015-10-18 23:32:14 +03:00
return
}
defer inr.Close()
defer inw.Close()
2015-10-18 23:32:14 +03:00
proc, err := os.StartProcess(cmdPath, flag.Args(), &os.ProcAttr{
Files: []*os.File{inr, outw, outw},
})
if err != nil {
internalError(ws, "start:", err)
2015-10-18 23:32:14 +03:00
return
}
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()
2015-10-18 23:32:14 +03:00
// 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)
}
<-done
}
if _, err := proc.Wait(); err != nil {
2015-10-18 23:32:14 +03:00
log.Println("wait:", err)
}
log.Println("exiting handler")
2015-10-18 23:32:14 +03:00
}
func serveHome(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/" {
http.Error(w, "Not found", 404)
return
}
if r.Method != "GET" {
http.Error(w, "Method not allowed", 405)
return
}
w.Header().Set("Content-Type", "text/html; charset=utf-8")
homeTempl.Execute(w, r.Host)
}
func main() {
flag.Parse()
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)
}
2015-10-18 23:32:14 +03:00
http.HandleFunc("/", serveHome)
http.HandleFunc("/ws", serveWs)
log.Fatal(http.ListenAndServe(*addr, nil))
}