note
This article was last updated on June 7, 2023, 2 years 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.
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.