Record asciicast v2 input and marker events

cmd/sshpiperd/asciicast.go

@@ -83,7 +83,21 @@ func (l *asciicastLogger) uphook(msg []byte) error {
 }
 
 func (l *asciicastLogger) downhook(msg []byte) error {
-	if msg[0] == msgChannelRequest {
+	switch msg[0] {
+	case msgChannelData:
+		serverChannelID := binary.BigEndian.Uint32(msg[1:5])
+		clientChannelID := l.channelIDMap[serverChannelID]
+		f, ok := l.channels[clientChannelID]
+		if ok {
+			buf := msg[9:]
+			t := time.Since(l.starttime).Seconds()
+
+			_, err := fmt.Fprintf(f, "[%v,\"i\",\"%s\"]\n", t, jsonEscape(string(buf)))
+			if err != nil {
+				return err
+			}
+		}
+	case msgChannelRequest:
 		t := time.Since(l.starttime).Seconds()
 		serverChannelID := binary.BigEndian.Uint32(msg[1:5])
 		clientChannelID := l.channelIDMap[serverChannelID]
@@ -92,30 +106,51 @@ func (l *asciicastLogger) downhook(msg []byte) error {
 
 		switch reqType {
 		case "pty-req":
-			_, _ = buf.ReadByte()
+			if _, err := buf.ReadByte(); err != nil {
+				return err
+			}
 			term := readString(buf)
-			_ = binary.Read(buf, binary.BigEndian, &l.initWidth)
-			_ = binary.Read(buf, binary.BigEndian, &l.initHeight)
+			if err := binary.Read(buf, binary.BigEndian, &l.initWidth); err != nil {
+				return err
+			}
+			if err := binary.Read(buf, binary.BigEndian, &l.initHeight); err != nil {
+				return err
+			}
 			l.envs["TERM"] = term
 		case "env":
-			_, _ = buf.ReadByte()
+			if _, err := buf.ReadByte(); err != nil {
+				return err
+			}
 			varName := readString(buf)
 			varValue := readString(buf)
 			l.envs[varName] = varValue
 		case "window-change":
 			f, ok := l.channels[clientChannelID]
-			if !ok {
-				_, _ = buf.ReadByte()
+			if ok {
+				if _, err := buf.ReadByte(); err != nil {
+					return err
+				}
 				var width, height uint32
-				_ = binary.Read(buf, binary.BigEndian, &width)
-				_ = binary.Read(buf, binary.BigEndian, &height)
+				if err := binary.Read(buf, binary.BigEndian, &width); err != nil {
+					return err
+				}
+				if err := binary.Read(buf, binary.BigEndian, &height); err != nil {
+					return err
+				}
 
 				_, err := fmt.Fprintf(f, "[%v,\"r\", \"%vx%v\"]\n", t, width, height)
 				if err != nil {
 					return err
 				}
 			}
 		case "shell", "exec":
+			if _, err := buf.ReadByte(); err != nil {
+				return err
+			}
+			var marker string
+			if reqType == "exec" {
+				marker = readString(buf)
+			}
 			jsonEnvs, err := json.Marshal(l.envs)
 			if err != nil {
 				return err
@@ -145,6 +180,14 @@ func (l *asciicastLogger) downhook(msg []byte) error {
 			if err != nil {
 				return err
 			}
+
+			if reqType == "exec" {
+				t := time.Since(l.starttime).Seconds()
+				_, err = fmt.Fprintf(f, "[%v,\"m\",\"%s\"]\n", t, jsonEscape(marker))
+				if err != nil {
+					return err
+				}
+			}
 		}
 	}
 	return nil

cmd/sshpiperd/asciicast_test.go

@@ -0,0 +1,113 @@
+package main
+
+import (
+	"bytes"
+	"encoding/binary"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+func TestAsciicastLogger_InputEvent(t *testing.T) {
+	recordDir := t.TempDir()
+	logger := newAsciicastLogger(recordDir, "")
+
+	serverChannelID := uint32(22)
+	clientChannelID := uint32(11)
+
+	openConfirm := make([]byte, 9)
+	openConfirm[0] = msgChannelOpenConfirm
+	binary.BigEndian.PutUint32(openConfirm[1:5], clientChannelID)
+	binary.BigEndian.PutUint32(openConfirm[5:9], serverChannelID)
+	if err := logger.uphook(openConfirm); err != nil {
+		t.Fatalf("uphook open confirm: %v", err)
+	}
+
+	req := buildChannelRequest(serverChannelID, "shell", nil)
+	if err := logger.downhook(req); err != nil {
+		t.Fatalf("downhook shell: %v", err)
+	}
+
+	inputPayload := []byte("ls -la\n")
+	inputMsg := buildChannelData(serverChannelID, inputPayload)
+	if err := logger.downhook(inputMsg); err != nil {
+		t.Fatalf("downhook input: %v", err)
+	}
+
+	_ = logger.Close()
+
+	castPath := filepath.Join(recordDir, "shell-channel-11.cast")
+	content, err := os.ReadFile(castPath)
+	if err != nil {
+		t.Fatalf("read cast: %v", err)
+	}
+
+	if !bytes.Contains(content, []byte("\"i\"")) {
+		t.Fatalf("expected input event in cast, got: %s", string(content))
+	}
+}
+
+func TestAsciicastLogger_MarkerEvent(t *testing.T) {
+	recordDir := t.TempDir()
+	logger := newAsciicastLogger(recordDir, "")
+
+	serverChannelID := uint32(42)
+	clientChannelID := uint32(7)
+
+	openConfirm := make([]byte, 9)
+	openConfirm[0] = msgChannelOpenConfirm
+	binary.BigEndian.PutUint32(openConfirm[1:5], clientChannelID)
+	binary.BigEndian.PutUint32(openConfirm[5:9], serverChannelID)
+	if err := logger.uphook(openConfirm); err != nil {
+		t.Fatalf("uphook open confirm: %v", err)
+	}
+
+	execPayload := buildStringPayload("date")
+	req := buildChannelRequest(serverChannelID, "exec", execPayload)
+	if err := logger.downhook(req); err != nil {
+		t.Fatalf("downhook exec: %v", err)
+	}
+
+	_ = logger.Close()
+
+	castPath := filepath.Join(recordDir, "exec-channel-7.cast")
+	content, err := os.ReadFile(castPath)
+	if err != nil {
+		t.Fatalf("read cast: %v", err)
+	}
+
+	if !bytes.Contains(content, []byte("\"m\"")) {
+		t.Fatalf("expected marker event in cast, got: %s", string(content))
+	}
+}
+
+func buildChannelData(channelID uint32, payload []byte) []byte {
+	msg := make([]byte, 9+len(payload))
+	msg[0] = msgChannelData
+	binary.BigEndian.PutUint32(msg[1:5], channelID)
+	binary.BigEndian.PutUint32(msg[5:9], uint32(len(payload)))
+	copy(msg[9:], payload)
+	return msg
+}
+
+func buildChannelRequest(channelID uint32, reqType string, payload []byte) []byte {
+	typeBytes := buildStringPayload(reqType)
+	msg := make([]byte, 5+len(typeBytes)+1+len(payload))
+	msg[0] = msgChannelRequest
+	binary.BigEndian.PutUint32(msg[1:5], channelID)
+	copy(msg[5:], typeBytes)
+	msg[5+len(typeBytes)] = 0
+	copy(msg[6+len(typeBytes):], payload)
+	return msg
+}
+
+func buildStringPayload(value string) []byte {
+	buf := &bytes.Buffer{}
+	if err := binary.Write(buf, binary.BigEndian, uint32(len(value))); err != nil {
+		panic(err)
+	}
+	if _, err := buf.WriteString(value); err != nil {
+		panic(err)
+	}
+	return buf.Bytes()
+}