gterm/gterm.go

package gterm

import (
	"bytes"
	_ "embed"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/creack/pty"
	"github.com/gorilla/websocket"
	"kumoly.io/lib/klog"
	"kumoly.io/lib/ksrv"
	"kumoly.io/lib/ksrv/engine"
	"kumoly.io/tools/gterm/public"
)

//go:embed public/index.html
var index string

var tmpl *engine.Engine
var servePublic = http.FileServer(http.FS(public.FS))

var pool map[string]chan struct{}

func init() {
	tmpl = engine.Must(engine.New("").Parse(index))
	pool = map[string]chan struct{}{}
}

type GTerm struct {
	AppName string

	// Arguments is a list of strings to pass as arguments to the specified Command
	Args []string
	// Command is the path to the binary we should create a TTY for
	Cmd string
	// Envs env pairs to pass to the command
	Envs []string
	// ErrorLimit defines the number of consecutive errors that can happen
	// before a connection is considered unusable
	ErrorLimit int
	// Timeout defines the maximum duration between which a ping and pong
	// cycle should be tolerated, beyond this the connection should be deemed dead
	Timeout    time.Duration
	BufferSize int
}

func New() *GTerm {
	// if g.Cmd == "" {
	// 	g.Cmd = "bash"
	// }
	// if len(g.Envs) == 0 {
	// 	g.Envs = os.Environ()
	// }
	// if g.Timeout < time.Second {
	// 	g.Timeout = 20 * time.Second
	// }
	// if g.ErrorLimit == 0 {
	// 	g.ErrorLimit = 10
	// }
	// if g.BufferSize == 0 {
	// 	g.BufferSize = 512
	// }
	go func() {
		//print pool
		log := klog.Sub("gterm")
		for {
			cons := []string{}
			for k := range pool {
				cons = append(cons, k)
			}
			log.Info("current connections: ", cons)
			time.Sleep(5 * time.Minute)
		}
	}()
	return &GTerm{
		AppName:    "GTERM",
		Cmd:        "bash",
		Envs:       append(os.Environ(), "TERM=xterm-256color"),
		Timeout:    20 * time.Second,
		ErrorLimit: 10,
		BufferSize: 512,
	}
}

func (g *GTerm) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if strings.HasPrefix(r.URL.String(), "/ws") {
		g.WS(w, r)
	} else {
		g.App(w, r)
	}
}

type App struct {
	AppName string
}

func (g *GTerm) App(w http.ResponseWriter, r *http.Request) {
	if r.URL.String() == "/" {
		tmpl.Execute(w, App{g.AppName})
		return
	}
	file, err := public.FS.Open(strings.TrimPrefix(r.URL.String(), "/"))
	if err != nil {
		// tmpl.Execute(w, App{g.AppName})
		http.NotFound(w, r)
		return
	}
	stat, err := file.Stat()
	if err != nil || stat.IsDir() {
		// tmpl.Execute(w, App{g.AppName})
		http.NotFound(w, r)
		return
	}
	servePublic.ServeHTTP(w, r)
}

func (g *GTerm) WS(w http.ResponseWriter, r *http.Request) {
	id := fmt.Sprintf("[%02d] %s", ctr(), ksrv.GetIP(r))
	pool[id] = make(chan struct{}, 1)
	defer func() {
		close(pool[id])
		delete(pool, id)
	}()
	l := klog.Sub(id)
	l.Info("connection established.")
	upgrader := websocket.Upgrader{
		HandshakeTimeout: 0,
		ReadBufferSize:   g.BufferSize,
		WriteBufferSize:  g.BufferSize,
	}
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		l.Error(err)
		return
	}

	cmd := exec.Command(g.Cmd, g.Args...)
	cmd.Env = g.Envs
	tty, err := pty.Start(cmd)
	if err != nil {
		l.Error(err)
		conn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
	}
	defer func() {
		if err := cmd.Process.Kill(); err != nil {
			l.Error(err)
		}
		if _, err := cmd.Process.Wait(); err != nil {
			l.Error(err)
		}
		if err := tty.Close(); err != nil {
			l.Error(err)
		}
		if err := conn.Close(); err != nil {
			l.Error(err)
		}
	}()

	var connectionClosed bool
	waiter := make(chan struct{}, 1)

	// this is a keep-alive loop that ensures connection does not hang-up itself
	lastPongTime := time.Now()
	conn.SetPongHandler(func(msg string) error {
		lastPongTime = time.Now()
		return nil
	})
	go func() {
		for {
			if err := conn.WriteMessage(websocket.PingMessage, []byte("keepalive")); err != nil {
				l.Warn("failed to write ping message")
				return
			}
			time.Sleep(g.Timeout / 2)
			if time.Since(lastPongTime) > g.Timeout {
				l.Warn("failed to get response from ping, triggering disconnect now...")
				waiter <- struct{}{}
				return
			}
			l.Debug("received response from ping successfully")
		}
	}()

	// tty >> xterm.js
	go func() {
		errorCounter := 0
		for {
			// consider the connection closed/errored out so that the socket handler
			// can be terminated - this frees up memory so the service doesn't get
			// overloaded
			if errorCounter > g.ErrorLimit {
				waiter <- struct{}{}
				break
			}
			buffer := make([]byte, g.BufferSize)
			readLength, err := tty.Read(buffer)
			if err != nil {
				l.Warn("failed to read from tty: ", err)
				if err := conn.WriteMessage(websocket.TextMessage, []byte("bye!")); err != nil {
					l.Warn("failed to send termination message from tty to xterm.js: ", err)
				}
				waiter <- struct{}{}
				return
			}
			if err := conn.WriteMessage(websocket.BinaryMessage, buffer[:readLength]); err != nil {
				l.Warn("failed to send ", readLength, " bytes from tty to xterm.js")
				errorCounter++
				continue
			}
			l.Debug("sent message of size ", readLength, " bytes from tty to xterm.js")
			errorCounter = 0
		}
	}()

	// tty << xterm.js
	go func() {
		for {
			// data processing
			messageType, data, err := conn.ReadMessage()
			if err != nil {
				if !connectionClosed {
					l.Warn("failed to get next reader: ", err)
				}
				return
			}
			dataLength := len(data)
			dataBuffer := bytes.Trim(data, "\x00")
			dataType, ok := WebsocketMessageType[messageType]
			if !ok {
				dataType = "uunknown"
			}
			l.Debug(fmt.Sprintf("received %s (type: %v) message of size %v byte(s) from xterm.js with key sequence: %v", dataType, messageType, dataLength, dataBuffer))

			// process
			if dataLength == -1 { // invalid
				l.Warn("failed to get the correct number of bytes read, ignoring message")
				continue
			}

			// handle resizing
			if messageType == websocket.BinaryMessage {
				if dataBuffer[0] == 1 {
					ttySize := &TTYSize{}
					resizeMessage := bytes.Trim(dataBuffer[1:], " \n\r\t\x00\x01")
					if err := json.Unmarshal(resizeMessage, ttySize); err != nil {
						l.Warn(fmt.Sprintf("failed to unmarshal received resize message '%s': %s", string(resizeMessage), err))
						continue
					}
					l.DebugF(klog.H{"rows": ttySize.Rows, "columns": ttySize.Cols}, "resizing tty...")
					if err := pty.Setsize(tty, &pty.Winsize{
						Rows: ttySize.Rows,
						Cols: ttySize.Cols,
					}); err != nil {
						l.Warn("failed to resize tty, error: ", err)
					}
					continue
				}
			}

			// write to tty
			bytesWritten, err := tty.Write(dataBuffer)
			if err != nil {
				l.Warn(fmt.Sprintf("failed to write %v bytes to tty: %s", len(dataBuffer), err))
				continue
			}
			l.Debug(bytesWritten, " bytes written to tty...")
		}
	}()

	<-waiter
	close(waiter)
	l.Info("closing connection...")
	connectionClosed = true
}

var WebsocketMessageType = map[int]string{
	websocket.BinaryMessage: "binary",
	websocket.TextMessage:   "text",
	websocket.CloseMessage:  "close",
	websocket.PingMessage:   "ping",
	websocket.PongMessage:   "pong",
}

// TTYSize represents a JSON structure to be sent by the frontend
// xterm.js implementation to the xterm.js websocket handler
type TTYSize struct {
	Cols uint16 `json:"cols"`
	Rows uint16 `json:"rows"`
	X    uint16 `json:"x"`
	Y    uint16 `json:"y"`
}