gotty/webtty/webtty.go

221 lines
4.5 KiB
Go
Raw Normal View History

2017-02-25 22:37:07 +00:00
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
}