Refactor

webtty/doc.go

@@ -0,0 +1,3 @@
+// Package webtty provides a protocl and an implementation to
+// controll terminals thorough networks.
+package webtty

webtty/errors.go

@@ -0,0 +1,10 @@
+package webtty
+
+import (
+	"errors"
+)
+
+var (
+	ErrSlaveClosed  = errors.New("slave closed")
+	ErrMasterClosed = errors.New("master closed")
+)

webtty/master.go

@@ -0,0 +1,55 @@
+/*
+Copyright (c) 2013 The Gorilla WebSocket Authors. All rights reserved.
+
+	Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+Redistributions of source code must retain the above copyright notice, this
+list of conditions and the following disclaimer.
+
+	Redistributions in binary form must reproduce the above copyright notice,
+this list of conditions and the following disclaimer in the documentation
+and/or other materials provided with the distribution.
+
+	THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+	SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+*/
+
+package webtty
+
+// Master represents a PTY master, usually it's a websocket connection.
+type Master interface {
+	WriteMessage(messageType int, data []byte) error
+	ReadMessage() (messageType int, p []byte, err error)
+}
+
+// The message types are defined in RFC 6455, section 11.8.
+const (
+	// TextMessage denotes a text data message. The text message payload is
+	// interpreted as UTF-8 encoded text data.
+	WSTextMessage = 1
+
+	// BinaryMessage denotes a binary data message.
+	WSBinaryMessage = 2
+
+	// CloseMessage denotes a close control message. The optional message
+	// payload contains a numeric code and text. Use the FormatCloseMessage
+	// function to format a close message payload.
+	WSCloseMessage = 8
+
+	// PingMessage denotes a ping control message. The optional message payload
+	// is UTF-8 encoded text.
+	WSPingMessage = 9
+
+	// PongMessage denotes a ping control message. The optional message payload
+	// is UTF-8 encoded text.
+	WSPongMessage = 10
+)

webtty/message_types.go

@@ -0,0 +1,29 @@
+package webtty
+
+var Protocols = []string{"webtty"}
+
+const (
+	// Unknown message type, maybe sent by a bug
+	UnknownInput = '0'
+	// User input typically from a keyboard
+	Input = '1'
+	// Ping to the server
+	Ping = '2'
+	// Notify that the browser size has been changed
+	ResizeTerminal = '3'
+)
+
+const (
+	// Unknown message type, maybe set by a bug
+	UnknownOutput = '0'
+	// Normal output to the terminal
+	Output = '1'
+	// Pong to the browser
+	Pong = '2'
+	// Set window title of the terminal
+	SetWindowTitle = '3'
+	// Set terminal preference
+	SetPreferences = '4'
+	// Make terminal to reconnect
+	SetReconnect = '5'
+)

webtty/option.go

@@ -0,0 +1,55 @@
+package webtty
+
+import (
+	"encoding/json"
+
+	"github.com/pkg/errors"
+)
+
+// Option is an option for WebTTY.
+type Option func(*WebTTY) error
+
+// WithPermitWrite sets a WebTTY to accept input from slaves.
+func WithPermitWrite() Option {
+	return func(wt *WebTTY) error {
+		wt.permitWrite = true
+		return nil
+	}
+}
+
+// WithFixedSize sets a fixed size to TTY master.
+func WithFixedSize(width int, height int) Option {
+	return func(wt *WebTTY) error {
+		wt.width = width
+		wt.height = height
+		return nil
+	}
+}
+
+// WithWindowTitle sets the default window title of the session
+func WithWindowTitle(windowTitle []byte) Option {
+	return func(wt *WebTTY) error {
+		wt.windowTitle = windowTitle
+		return nil
+	}
+}
+
+// WithReconnect enables reconnection on the master side.
+func WithReconnect(timeInSeconds int) Option {
+	return func(wt *WebTTY) error {
+		wt.reconnect = timeInSeconds
+		return nil
+	}
+}
+
+// WithMasterPreferences sets an optional configuration of master.
+func WithMasterPreferences(preferences interface{}) Option {
+	return func(wt *WebTTY) error {
+		prefs, err := json.Marshal(preferences)
+		if err != nil {
+			return errors.Wrapf(err, "failed to marshal preferences as JSON")
+		}
+		wt.masterPrefs = prefs
+		return nil
+	}
+}

webtty/slave.go

@@ -0,0 +1,13 @@
+package webtty
+
+import (
+	"io"
+)
+
+// Slave represents a PTY slave, typically it's a local command.
+type Slave interface {
+	io.ReadWriteCloser
+
+	WindowTitleVariables() map[string]interface{}
+	ResizeTerminal(columns int, rows int) error
+}

webtty/webtty.go

@@ -0,0 +1,220 @@
+package webtty
+
+import (
+	"context"
+	"encoding/base64"
+	"encoding/json"
+	"sync"
+
+	"github.com/pkg/errors"
+)
+
+// WebTTY bridges sets of a PTY slave and its PTY master.
+// To support text-based streams and side channel commands such as
+// terminal resizing, WebTTY uses an original protocol.
+type WebTTY struct {
+	// PTY Master, which probably a connection to browser
+	masterConn Master
+	// PTY Slave
+	slave Slave
+
+	windowTitle []byte
+	permitWrite bool
+	width       int
+	height      int
+	reconnect   int // in milliseconds
+	masterPrefs []byte
+
+	bufferSize int
+	writeMutex sync.Mutex
+}
+
+// New creates a new instance of WebTTY.
+// masterConn is a connection to the PTY master,
+// typically it's a websocket connection to a client.
+// slave is a PTY slave such as a local command with a PTY.
+func New(masterConn Master, slave Slave, options ...Option) (*WebTTY, error) {
+	wt := &WebTTY{
+		masterConn: masterConn,
+		slave:      slave,
+
+		permitWrite: false,
+		width:       0,
+		height:      0,
+
+		bufferSize: 1024,
+	}
+
+	for _, option := range options {
+		option(wt)
+	}
+
+	return wt, nil
+}
+
+// Run starts the WebTTY.
+// This method blocks until the context is canceled.
+// Note that the master and slave are left intact even
+// after the context is canceled. Closing them is caller's
+// responsibility.
+func (wt *WebTTY) Run(ctx context.Context) error {
+	err := wt.sendInitializeMessage()
+	if err != nil {
+		return errors.Wrapf(err, "failed to send initializing message")
+	}
+
+	errs := make(chan error, 2)
+
+	go func() {
+		errs <- func() error {
+			buffer := make([]byte, wt.bufferSize)
+			for {
+				n, err := wt.slave.Read(buffer)
+				if err != nil {
+					return ErrSlaveClosed
+				}
+
+				err = wt.handleSlaveReadEvent(buffer[:n])
+				if err != nil {
+					return err
+				}
+			}
+		}()
+	}()
+
+	go func() {
+		errs <- func() error {
+			for {
+				typ, data, err := wt.masterConn.ReadMessage()
+				if err != nil {
+					return ErrMasterClosed
+				}
+				if typ != WSTextMessage {
+					continue
+				}
+
+				err = wt.handleMasterReadEvent(data)
+				if err != nil {
+					return err
+				}
+			}
+		}()
+	}()
+
+	select {
+	case <-ctx.Done():
+		err = ctx.Err()
+	case err = <-errs:
+	}
+
+	return err
+}
+
+func (wt *WebTTY) sendInitializeMessage() error {
+	err := wt.masterWrite(append([]byte{SetWindowTitle}, wt.windowTitle...))
+	if err != nil {
+		return errors.Wrapf(err, "failed to send window title")
+	}
+
+	if wt.reconnect > 0 {
+		reconnect, _ := json.Marshal(wt.reconnect)
+		err := wt.masterWrite(append([]byte{SetReconnect}, reconnect...))
+		if err != nil {
+			return errors.Wrapf(err, "failed to set reconnect")
+		}
+	}
+
+	if wt.masterPrefs != nil {
+		err := wt.masterWrite(append([]byte{SetPreferences}, wt.masterPrefs...))
+		if err != nil {
+			return errors.Wrapf(err, "failed to set preferences")
+		}
+	}
+
+	return nil
+}
+
+func (wt *WebTTY) handleSlaveReadEvent(data []byte) error {
+	safeMessage := base64.StdEncoding.EncodeToString(data)
+	err := wt.masterWrite(append([]byte{Output}, []byte(safeMessage)...))
+	if err != nil {
+		return errors.Wrapf(err, "failed to send message to master")
+	}
+
+	return nil
+}
+
+func (wt *WebTTY) masterWrite(data []byte) error {
+	wt.writeMutex.Lock()
+	defer wt.writeMutex.Unlock()
+
+	err := wt.masterConn.WriteMessage(WSTextMessage, data)
+	if err != nil {
+		return errors.Wrapf(err, "failed to write to master")
+	}
+
+	return nil
+}
+
+func (wt *WebTTY) handleMasterReadEvent(data []byte) error {
+	if len(data) == 0 {
+		return errors.New("unexpected zero length read from master")
+	}
+
+	switch data[0] {
+	case Input:
+		if !wt.permitWrite {
+			return nil
+		}
+
+		if len(data) <= 1 {
+			return nil
+		}
+
+		_, err := wt.slave.Write(data[1:])
+		if err != nil {
+			return errors.Wrapf(err, "failed to write received data to slave")
+		}
+
+	case Ping:
+		err := wt.masterWrite([]byte{Pong})
+		if err != nil {
+			return errors.Wrapf(err, "failed to return Pong message to master")
+		}
+
+	case ResizeTerminal:
+		if wt.width != 0 && wt.height != 0 {
+			break
+		}
+
+		if len(data) <= 1 {
+			return errors.New("received malformed remote command for terminal resize: empty payload")
+		}
+
+		var args argResizeTerminal
+		err := json.Unmarshal(data[1:], &args)
+		if err != nil {
+			return errors.Wrapf(err, "received malformed data for terminal resize")
+		}
+		rows := wt.height
+		if rows == 0 {
+			rows = int(args.Rows)
+		}
+
+		columns := wt.width
+		if columns == 0 {
+			columns = int(args.Columns)
+		}
+
+		wt.slave.ResizeTerminal(columns, rows)
+	default:
+		return errors.Errorf("unknown message type `%c`", data[0])
+	}
+
+	return nil
+}
+
+type argResizeTerminal struct {
+	Columns float64
+	Rows    float64
+}

webtty/webtty_test.go

@@ -0,0 +1,139 @@
+package webtty
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"io"
+	"sync"
+	"testing"
+)
+
+type pipePair struct {
+	*io.PipeReader
+	*io.PipeWriter
+}
+
+func TestWriteFromPTY(t *testing.T) {
+	connInPipeReader, connInPipeWriter := io.Pipe() // in to conn
+	connOutPipeReader, _ := io.Pipe()               // out from conn
+
+	conn := pipePair{
+		connOutPipeReader,
+		connInPipeWriter,
+	}
+	dt, err := New(conn)
+	if err != nil {
+		t.Fatalf("Unexpected error from New(): %s", err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		wg.Done()
+		err := dt.Run(ctx)
+		if err != nil {
+			t.Fatalf("Unexpected error from Run(): %s", err)
+		}
+	}()
+
+	message := []byte("foobar")
+	n, err := dt.TTY().Write(message)
+	if err != nil {
+		t.Fatalf("Unexpected error from Write(): %s", err)
+	}
+	if n != len(message) {
+		t.Fatalf("Write() accepted `%d` for message `%s`", n, message)
+	}
+
+	buf := make([]byte, 1024)
+	n, err = connInPipeReader.Read(buf)
+	if err != nil {
+		t.Fatalf("Unexpected error from Read(): %s", err)
+	}
+	if buf[0] != Output {
+		t.Fatalf("Unexpected message type `%c`", buf[0])
+	}
+	decoded := make([]byte, 1024)
+	n, err = base64.StdEncoding.Decode(decoded, buf[1:n])
+	if err != nil {
+		t.Fatalf("Unexpected error from Decode(): %s", err)
+	}
+	if !bytes.Equal(decoded[:n], message) {
+		t.Fatalf("Unexpected message received: `%s`", decoded[:n])
+	}
+
+	cancel()
+	wg.Wait()
+}
+
+func TestWriteFromConn(t *testing.T) {
+	connInPipeReader, connInPipeWriter := io.Pipe()   // in to conn
+	connOutPipeReader, connOutPipeWriter := io.Pipe() // out from conn
+
+	conn := pipePair{
+		connOutPipeReader,
+		connInPipeWriter,
+	}
+
+	dt, err := New(conn)
+	if err != nil {
+		t.Fatalf("Unexpected error from New(): %s", err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		wg.Done()
+		err := dt.Run(ctx)
+		if err != nil {
+			t.Fatalf("Unexpected error from Run(): %s", err)
+		}
+	}()
+
+	var (
+		message []byte
+		n       int
+	)
+	readBuf := make([]byte, 1024)
+
+	// input
+	message = []byte("0hello\n") // line buffered canonical mode
+	n, err = connOutPipeWriter.Write(message)
+	if err != nil {
+		t.Fatalf("Unexpected error from Write(): %s", err)
+	}
+	if n != len(message) {
+		t.Fatalf("Write() accepted `%d` for message `%s`", n, message)
+	}
+
+	n, err = dt.TTY().Read(readBuf)
+	if err != nil {
+		t.Fatalf("Unexpected error from Write(): %s", err)
+	}
+	if !bytes.Equal(readBuf[:n], message[1:]) {
+		t.Fatalf("Unexpected message received: `%s`", readBuf[:n])
+	}
+
+	// ping
+	message = []byte("1\n") // line buffered canonical mode
+	n, err = connOutPipeWriter.Write(message)
+	if n != len(message) {
+		t.Fatalf("Write() accepted `%d` for message `%s`", n, message)
+	}
+
+	n, err = connInPipeReader.Read(readBuf)
+	if err != nil {
+		t.Fatalf("Unexpected error from Read(): %s", err)
+	}
+	if !bytes.Equal(readBuf[:n], []byte{'1'}) {
+		t.Fatalf("Unexpected message received: `%s`", readBuf[:n])
+	}
+
+	// TODO: resize
+
+	cancel()
+	wg.Wait()
+}