KubePi/thirdparty/gotty/webtty/webtty.go

package webtty

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"github.com/pkg/errors"
	"os"
	"strconv"
	"sync"
	"time"
)

// WebTTY bridges 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
	columns     int
	rows        int
	reconnect   int // in seconds
	masterPrefs []byte

	bufferSize   int
	writeMutex   sync.Mutex
	lastPingTime time.Time
}

const (
	MaxBufferSize = 1024 * 1024 * 1
)

// 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,
		columns:     0,
		rows:        0,

		bufferSize:   MaxBufferSize,
		lastPingTime: time.Now(),
	}

	for _, option := range options {
		option(wt)
	}

	return wt, nil
}

// Run starts the main process of 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.
// If the connection to one end gets closed, returns ErrSlaveClosed or ErrMasterClosed.
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, 3)

	slaveBuffer := make([]byte, wt.bufferSize)
	go func() {
		errs <- func() error {
			defer func() {
				if e := recover(); e != nil {
				}
			}()
			for {
				if slaveBuffer == nil {
					return ErrSlaveClosed
				}
				n, err := wt.slave.Read(slaveBuffer)
				if err != nil {
					return ErrSlaveClosed
				}
				err = wt.handleSlaveReadEvent(slaveBuffer[:n])
				if err != nil {
					return err
				}
			}
		}()
	}()
	masterBuffer := make([]byte, wt.bufferSize)
	go func() {
		errs <- func() error {
			defer func() {
				if e := recover(); e != nil {
				}
			}()
			for {
				if masterBuffer == nil {
					return ErrMasterClosed
				}
				n, err := wt.masterConn.Read(masterBuffer)
				if err != nil {
					return ErrMasterClosed
				}
				err = wt.handleMasterReadEvent(masterBuffer[:n])
				if err != nil {
					return err
				}
			}
		}()
	}()

	go func() {
		errs <- func() error {
			lostPingTimeout := time.Duration(180) * time.Second
			seconds, _err := strconv.Atoi(os.Getenv("LOST_PING_TIMEOUT_SECONDS"))
			if _err != nil && seconds > 30 {
				lostPingTimeout = time.Duration(seconds) * time.Second
			}
			for {
				time.Sleep(time.Duration(30) * time.Second)
				if err != nil {
					return err
				}
				if time.Now().After(wt.lastPingTime.Add(lostPingTimeout)) {
					return ErrConnectionLostPing
				}
			}
		}()
	}()

	defer func() {
		slaveBuffer = nil
		masterBuffer = nil
	}()

	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.Write(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})
		wt.lastPingTime = time.Now()
		if err != nil {
			return errors.Wrapf(err, "failed to return Pong message to master")
		}

	case ResizeTerminal:
		if wt.columns != 0 && wt.rows != 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.rows
		if rows == 0 {
			rows = int(args.Rows)
		}

		columns := wt.columns
		if columns == 0 {
			columns = int(args.Columns)
		}

		return wt.slave.ResizeTerminal(columns, rows)
	default:
		return errors.Errorf("unknown message type `%c`", data[0])
	}

	return nil
}

type argResizeTerminal struct {
	Columns float64
	Rows    float64
}