Refactor

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
+}