gistfile1.txt

package main

/*
 * websocket/pty proxy server:
 * This program wires a websocket to a pty master.
 *
 * Usage:
 * go build -o ws-pty-proxy server.go
 * ./websocket-terminal -cmd /bin/bash -addr :9000 -static $HOME/src/websocket-terminal
 * ./websocket-terminal -cmd /bin/bash -- -i
 *
 * TODO:
 *  * make more things configurable
 *  * switch back to binary encoding after fixing term.js (see index.html)
 *  * make errors return proper codes to the web client
 *
 * Copyright 2014 Al Tobey [email protected]
 * MIT License, see the LICENSE file
 */

import (
	"flag"
	"github.com/eclipse/che-lib/websocket"
	"github.com/eclipse/che-lib/pty"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"encoding/json"
	"bufio"
	"bytes"
	"unicode/utf8"
	"errors"
	"fmt"
	"sync"
)

var addrFlag, cmdFlag, staticFlag string

type WebSocketMessage struct {
	Type string          `json:"type"`
	Data json.RawMessage `json:"data"`
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1,
	WriteBufferSize: 1,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type wsPty struct {
	Cmd *exec.Cmd // pty builds on os.exec
	Pty *os.File  // a pty is simply an os.File
}

func (wp *wsPty) Start() error {
	var err error
	args := flag.Args()
	wp.Cmd = exec.Command(cmdFlag, args...)

	env := os.Environ()
	env = append(env, "TERM=xterm")
	wp.Cmd.Env = env

	wp.Pty, err = pty.Start(wp.Cmd)
	if err != nil {
		return err
	}

	//Set the size of the pty
	pty.Setsize(wp.Pty, 60, 200)

	return nil
}

func (wp *wsPty) Stop() {
	wp.Pty.Close()
	fmt.Println("Pty file closed")
	wp.Cmd.Process.Kill()
	fmt.Println("Process killed")
}

// read from the web socket, copying to the pty master
// messages are expected to be text and base64 encoded
func sendConnectionInputToPty(ptyFile *os.File, conn *websocket.Conn, wg *sync.WaitGroup) {
	defer func() {
		ptyFile.Close()
		fmt.Println("ptyFile.Close()")
	}();
	defer fmt.Println("sendConnectionInputToPty finished")
	for {
		mt, payload, err := conn.ReadMessage()
		if err != nil {
			if err != io.EOF {
				log.Printf("conn.ReadMessage failed: %s\n", err)
				wg.Done()
				return
			}
		}
		var msg WebSocketMessage;
		switch mt {
		case websocket.BinaryMessage:
			log.Printf("Ignoring binary message: %q\n", payload)
		case websocket.TextMessage:
			if err := json.Unmarshal(payload, &msg); err != nil {
				log.Printf("Invalid message %s\n", err);
				continue
			}
			if errMsg := handleMessage(msg, ptyFile); errMsg != nil {
				log.Printf(errMsg.Error())
				wg.Done()
				return
			}

		default:
			log.Printf("Invalid websocket message type %d\n", mt)
			wg.Done()
			return
		}
	}
}

func handleMessage(msg WebSocketMessage, ptyFile *os.File) error {
	switch msg.Type {
	case "resize" :
		var size []float64;
		if err := json.Unmarshal(msg.Data, &size); err != nil {
			log.Printf("Invalid resize message: %s\n", err);
		} else {
			pty.Setsize(ptyFile, uint16(size[1]), uint16(size[0]));
		}

	case "data" :
		var dat string;
		if err := json.Unmarshal(msg.Data, &dat); err != nil {
			log.Printf("Invalid data message %s\n", err);
		} else {
			ptyFile.Write([]byte(dat));
		}

	default:
		return errors.New("Invalid field message type: " + msg.Type + "\n")
	}
	return nil
}

// copy everything from the pty master to the websocket
// using base64 encoding for now due to limitations in term.js
func sendPtyOutputToConnection(ptyFile *os.File, conn *websocket.Conn, wg *sync.WaitGroup) {
	defer fmt.Println("sendPtyOutputToConnection finished")
	buf := make([]byte, 8192)
	reader := bufio.NewReader(ptyFile)
	var buffer bytes.Buffer;
	// TODO: more graceful exit on socket close / process exit
	for {
		fmt.Println("sendPtyOutputToConnection before Read")
		n, err := reader.Read(buf);
		fmt.Println("sendPtyOutputToConnection after Read")
		if err != nil {
			log.Printf("Failed to read from pty master: %s", err)
			wg.Done()
			return
		}
		//read byte array as Unicode code points (rune in go)

		bufferBytes := buffer.Bytes()
		runeReader := bufio.NewReader(bytes.NewReader(append(bufferBytes[:], buf[:n]...)))
		buffer.Reset()
		i := 0;
		for i < n {
			char, charLen, e := runeReader.ReadRune()
			if e != nil {
				log.Printf("Failed to read from pty master: %s", e)
				wg.Done()
				return
			}

			if char == utf8.RuneError {
				runeReader.UnreadRune()
				break
			}
			i += charLen;
			buffer.WriteRune(char)
		}
		err = conn.WriteMessage(websocket.TextMessage, buffer.Bytes())
		if err != nil {
			log.Printf("Failed to send UTF-8 char: %s", err)
			wg.Done()
			return
		}
		buffer.Reset();
		if i < n {
			buffer.Write(buf[i:n])
		}
	}
}

var i int  = 0;

func ptyHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Fatalf("Websocket upgrade failed: %s\n", err)
	}
	defer conn.Close()

	wp := wsPty{}
	if err := wp.Start(); err != nil {
		log.Printf("Failed to start command: %s\n", err)
		http.Error(w, err.Error(), 500)
		return
	}

	wg := sync.WaitGroup{}
	wg.Add(2)

	go sendPtyOutputToConnection(wp.Pty, conn, &wg)
	go sendConnectionInputToPty(wp.Pty, conn, &wg)

	wg.Wait()
	fmt.Println("DONE")
}

func init() {
	cwd, _ := os.Getwd()
	flag.StringVar(&addrFlag, "addr", ":9000", "IP:PORT or :PORT address to listen on")
	flag.StringVar(&cmdFlag, "cmd", "/bin/bash", "command to execute on slave side of the pty")
	flag.StringVar(&staticFlag, "static", cwd, "path to static content")
	// TODO: make sure paths exist and have correct permissions
}

func main() {
	flag.Parse()

	http.HandleFunc("/pty", ptyHandler)

	// serve html & javascript
	http.Handle("/", http.FileServer(http.Dir(staticFlag)))

	err := http.ListenAndServe(addrFlag, nil)
	if err != nil {
		log.Fatalf("net.http could not listen on address '%s': %s\n", addrFlag, err)
	}
}