note

This article was last updated on June 7, 2023, 8 months ago. The content may be out of date.

Continuing from previous post, this post is about how to handle ttyd using golang.

Overview

The golang handler consists of 3 parts: sending frontend, sending auth token and handling websocket connection.

Sending Frontend

Golang can embed static assets in the compiled binary so there is no need to specify a path to the file. The frontend can be extracted using devtools.

//go:embed ttyd.html
var ttydHtml string

func index(writer http.ResponseWriter, request *http.Request) {
	_, _ = io.WriteString(writer, ttydHtml)
}

Sending Auth Token

We just write a static json string with an empty token since we don’t check the token in our implementation.

func token(writer http.ResponseWriter, request *http.Request) {
	_, _ = io.WriteString(writer, "{\"token\": \"\"}")
}

Handling Websocket Connection

This the main part of the ttyd protocol. We create 2 types: one to handle reading from and writing to the websocket connection, another to handle the protocol itself.

classDiagram class wsConn { brw *bufio.ReadWriter conn net.Conn lr io.LimitedReader rb bytes.Buffer wb bytes.Buffer e wsflate.Parameters accepted bool eof bool sw bytes.Buffer br bytes.Reader fr io.ReadCloser r flate.Resetter fw *flate.Writer hdr ws.Header lock sync.Mutex } class daemon { conn *wsConn cmd *exec.Cmd file *os.File paused atomic.Bool resume chan struct#123;#125; ioErr atomic.Bool }

wsConn

wsConn is used to handle the underlying websocket connection. It also handles websocket’s permessage-deflate extension. It uses github.com/gobwas/ws.

Close

Send a websocket close frame then close the connection. Don’t wait for client’s response.

func (w *wsConn) Close() {
	w.lock.Lock()
	_, _ = w.conn.Write(ws.CompiledCloseNormalClosure)
	w.lock.Unlock()
	_ = w.conn.Close()
}

handleControl

Handle control frames the client may send. There are 3 types of control frame, ping, pong and close and we’ll only need to act when receiving ping and pong frames.

func (w *wsConn) handleControl(hdr ws.Header) error {
	data, err := w.brw.Peek(int(hdr.Length))
	if err != nil {
		return err
	}

	switch hdr.OpCode {
	case ws.OpPing:
		ws.Cipher(data, hdr.Mask, 0)
		w.lock.Lock()
		_ = ws.WriteFrame(w.brw, ws.NewPongFrame(data))
		err = w.brw.Flush()
		w.lock.Unlock()
	case ws.OpClose:
		err = io.EOF
	}
	_, _ = w.brw.Discard(int(hdr.Length))
	return err
}

nextFrame

Read the next frame header and if it’s a control frame, handle it appropriately.

func (w *wsConn) nextFrame() error {
	for {
		hdr, err := ws.ReadHeader(w.brw)
		if err != nil {
			return err
		}

		if !hdr.Masked {
			return ws.ErrProtocolMaskRequired
		}

		if hdr.OpCode.IsControl() {
			err = w.handleControl(hdr)
			if err != nil {
				return err
			}
			continue
		}

		w.hdr = hdr
		return nil
	}
}

readFrame

Read the binary data contained within each text or binary frame into rb field. It handles deflate automatically.

info

Every data frame is read fully into memory, but because usually ttyd frames are 4k large at most, this doesn’t cause any memory problems in practice.

It also always assumes the maximum sliding window size when using context takeover.

Because permessage-deflate trims the trailing bytes, we add these bytes and trailing eof back to deflate successfully. There is no way to get the sliding windows from a deflate reader, so we have to keep our own. It’s extra memory usage, but it is what it is.

Also, the deflate reader can only be reset to be able to be read from again. And we’ll need to specify the sliding window when resetting if using context takeover.

func (w *wsConn) readFrame() error {
	r1 := w.hdr.Rsv1()
	for {
		idx := w.rb.Len()
		w.lr.N = w.hdr.Length
		_, err := w.rb.ReadFrom(&w.lr)
		if err != nil {
			return err
		}

		ws.Cipher(w.rb.Bytes()[idx:], w.hdr.Mask, 0)
		if w.hdr.Fin {
			break
		}
		err = w.nextFrame()
		if err != nil {
			return err
		}
	}

	if !w.accepted || !r1 {
		return nil
	}

	w.rb.Write(compressionReadTail)
	w.br.Reset(w.rb.Bytes())
	if w.e.ClientNoContextTakeover || w.eof {
		w.eof = false
		w.sw.Reset()
		_ = w.r.Reset(&w.br, nil)
	} else {
		_ = w.r.Reset(&w.br, w.sw.Bytes())
	}

	_, err := w.rb.ReadFrom(w.fr)
	if err != nil {
		return err
	}
	w.rb.Next(int(w.br.Size()))
	w.sw.Write(w.rb.Bytes())
	next := w.sw.Len() - 32768
	if next > 0 {
		w.sw.Next(next)
	}
	if w.br.Len() != 0 {
		w.eof = true
	}
	return nil
}

Write

Encode the data as a binary websocket frame and send it. The frame will not be fragmented. Compressed if permessage-deflate is negotiated.

func (w *wsConn) Write(p []byte) (n int, err error) {
	w.wb.Reset()
	w.lock.Lock()
	var frame ws.Frame
	if !w.accepted {
		frame = ws.NewBinaryFrame(p)
	} else {
		if w.e.ServerNoContextTakeover {
			w.fw.Reset(&w.wb)
		}
		_, _ = w.fw.Write(p)
		if w.e.ServerNoContextTakeover {
			_ = w.fw.Close()
		} else {
			_ = w.fw.Flush()
		}
		w.wb.Truncate(w.wb.Len() - 4)

		frame = ws.NewBinaryFrame(w.wb.Bytes())
		frame.Header.Rsv = ws.Rsv(true, false, false)
	}
	_ = ws.WriteFrame(w.brw, frame)
	err = w.brw.Flush()
	w.lock.Unlock()
	if err == nil {
		n = len(p)
	}
	return
}

daemon

daemon is used to handle ttyd protocol and communicate with the backing pty process.

cleanup

Close the network connection and then release the pty process’s resources.

func (d *daemon) cleanup() {
	d.conn.Close()
	if d.file != nil {
		_ = d.file.Close()
		_ = d.cmd.Wait()

		select {
		case d.resume <- struct{}{}:
		default:
		}
	}
}

initWrite

Send the title the ttyd tab should use and ttyd client options.

func (d *daemon) initWrite() error {
	d.conn.rb.WriteByte(setWindowTitle)
	d.conn.rb.WriteString("title of ttyd tab")
	_, err := d.conn.rb.WriteTo(d.conn)
	if err != nil {
		return err
	}

	d.conn.rb.WriteByte(setPreference)
	d.conn.rb.WriteString("{ }")
	_, err = d.conn.rb.WriteTo(d.conn)
	return err
}

readLoop

Send the initial ttyd messages and read from the connection until there is an error.

func (d *daemon) readLoop() {
	err := d.initWrite()
	if err != nil {
		return
	}
	d.conn.lr.R = d.conn.brw
	for !d.ioErr.Load() {
		d.conn.rb.Reset()
		for d.conn.rb.Len() == 0 {
			err = d.conn.nextFrame()
			if err != nil {
				return
			}

			err = d.conn.readFrame()
			if err != nil {
				return
			}
		}

		cmd, _ := d.conn.rb.ReadByte()
		if (cmd == jsonData && d.file != nil) || (cmd != jsonData && d.file == nil) {
			continue
		}

		switch cmd {
		case input:
			_, err = d.conn.rb.WriteTo(d.file)
			if err != nil {
				return
			}
		case resizeTerminal:
			var rr resizeRequest
			err = json.NewDecoder(&d.conn.rb).Decode(&rr)
			if err != nil {
				return
			}

			err = pty.Setsize(d.file, &pty.Winsize{
				Rows: rr.Rows,
				Cols: rr.Columns,
			})
			if err != nil {
				return
			}
		case pause:
			d.paused.Store(true)
		case resume:
			d.paused.Store(false)
			select {
			case d.resume <- struct{}{}:
			default:
			}
		case jsonData:
			_ = d.conn.rb.UnreadByte()
			var rr resizeRequest
			err = json.NewDecoder(&d.conn.rb).Decode(&rr)
			if err != nil {
				return
			}

			d.file, err = pty.StartWithSize(d.cmd, &pty.Winsize{
				Rows: rr.Rows,
				Cols: rr.Columns,
			})
			if err != nil {
				return
			}
			go d.writeLoop()
		}
	}
}

writeLoop

Read from the backing pty process and write to the connection util there is an error. Pause reading when the client requests to pause.

func (d *daemon) writeLoop() {
	d.conn.wb.Grow(4097)
	buf := d.conn.wb.Bytes()[:d.conn.wb.Cap()]
	for !d.ioErr.Load() {
		buf[0] = output
		n, err := d.file.Read(buf[1:])
		if err != nil {
			break
		}

		_, err = d.conn.Write(buf[:1+n])
		if err != nil {
			break
		}

		if d.paused.Load() {
			<-d.resume
		}
	}
	d.ioErr.Store(true)
	d.cleanup()
}

Websocket Handler

This part handles websocket connection establishment with permessage-deflate extension support.

func websocket(writer http.ResponseWriter, request *http.Request) {
	var (
		extension = &wsflate.Extension{}
		upgrader  = &ws.HTTPUpgrader{
			Protocol:  protocol,
			Negotiate: extension.Negotiate,
		}
	)
	conn, bw, _, err := upgrader.Upgrade(request, writer)
	if err != nil {
		return
	}

	e, accepted := extension.Accepted()
	d := &daemon{
		conn: &wsConn{
			brw:      bw,
			conn:     conn,
			e:        e,
			accepted: accepted,
		},
		cmd:    exec.Command("command to run"),
		resume: make(chan struct{}, 1),
	}
	if accepted {
		d.conn.fr = flate.NewReader(&d.conn.br)
		d.conn.r = d.conn.fr.(flate.Resetter)
		d.conn.fw, _ = flate.NewWriter(&d.conn.wb, 6)
	}

	d.readLoop()
	d.ioErr.Store(true)
	d.cleanup()
}

tip

We can specify what command to run based on the request.

info

We are using permessage-deflate with context takeover here. It’s easy to switch to no context takeover to save some memory.