From bd0d1aaebe6755ddbecbffb42a8f0c6599e020ed Mon Sep 17 00:00:00 2001
From: Laurence Withers <laurence.withers@yoti.com>
Date: Fri, 28 Apr 2023 12:22:51 +0100
Subject: [PATCH] Improvements: cope with extra input chars and replay

This still suffers a bit if the user does manage to type anything during the invocation, because the
simulated input characters get echoed to the terminal and mess up the nicely-aligned newline we just
made.
---
 main.go | 139 +++++++++++++++++++++++++++++++-------------------------
 1 file changed, 78 insertions(+), 61 deletions(-)

diff --git a/main.go b/main.go
index 0661720..1688c4c 100644
--- a/main.go
+++ b/main.go
@@ -2,105 +2,122 @@ package main
 
 import (
 	"bytes"
-	"fmt"
 	"os"
-	"strconv"
-	"strings"
 	"time"
+	"unsafe"
 
 	"golang.org/x/sys/unix"
 	"golang.org/x/term"
 )
 
+type dsrResult struct {
+	row        int
+	column     int
+	extraBytes []byte
+}
+
 func main() {
 	st, err := term.MakeRaw(0)
 	if err != nil {
 		os.Exit(1)
 	}
 
-	var column int
+	var res dsrResult
 	select {
 	case <-time.After(100 * time.Millisecond):
-	case column = <-awaitColumn():
+	case res = <-awaitDSR():
 	}
 
 	term.Restore(0, st)
-	if column > 1 {
+	if res.column > 1 {
 		unix.Write(1, []byte("🚫\n")) // U+1F6AB
 	}
+	for _, b := range res.extraBytes {
+		unix.Syscall(unix.SYS_IOCTL, 0, unix.TIOCSTI, uintptr(unsafe.Pointer(&b)))
+	}
 }
 
-func awaitColumn() <-chan int {
-	ch := make(chan int)
+func awaitDSR() <-chan dsrResult {
+	ch := make(chan dsrResult)
 	go func() {
-		row, col := readDSR()
-		_ = row
-		ch <- col
+		row, col, extra := readDSR()
+		ch <- dsrResult{
+			row:        row,
+			column:     col,
+			extraBytes: extra,
+		}
 	}()
 	return ch
 }
 
-func readDSR() (row, col int) {
+func readDSR() (row, col int, extra []byte) {
 	unix.Write(1, []byte{
 		'\033', '[', // ANSI escape code: CSI
 		'6', 'n', // DSR / device status report
 	})
 
-	buf := make([]byte, 64)
-	n, err := unix.Read(0, buf)
-	if err != nil {
-		return 0, 0
-	}
-	buf = buf[:n]
+	inputBuf := bytes.NewBuffer(make([]byte, 0, 128))
+	termIn := make([]byte, 64)
+	for {
+		// read some bytes from stdin and append to our input buffer
+		termIn = termIn[:64]
+		n, err := unix.Read(0, termIn)
+		inputBuf.Write(termIn[:n])
+		if err != nil {
+			return 0, 0, inputBuf.Bytes()
+		}
 
-	if len(buf) < 6 || buf[0] != '\033' || buf[1] != '[' || buf[len(buf)-1] != 'R' {
-		return 0, 0
+		// does the input buffer contain a valid DSR result sequence?
+		b := inputBuf.Bytes()
+		row, col, pos, count, ok := parseDSR(b)
+		if ok {
+			// replay the non-DSR sequence bytes
+			return row, col, append(b[:pos], b[pos+count:]...)
+		}
 	}
-
-	s := string(buf[2 : len(buf)-1])
-	pos := strings.IndexByte(s, ';')
-	if pos == -1 {
-		return 0, 0
-	}
-
-	row, err = strconv.Atoi(s[:pos])
-	if err != nil {
-		return 0, 0
-	}
-	col, err = strconv.Atoi(s[pos+1:])
-	if err != nil {
-		return 0, 0
-	}
-	return row, col
 }
 
-func processit(b []byte) {
-	if !bytes.HasPrefix(b, []byte{'\033', '['}) {
-		return
-	}
-	if !bytes.HasSuffix(b, []byte{'R'}) {
-		return
-	}
-	s := string(b[2 : len(b)-1])
-	pos := strings.IndexByte(s, ';')
+func parseDSR(input []byte) (row, col, pos, count int, ok bool) {
+	// locate the start of the result
+	pos = bytes.Index(input, []byte{'\033', '['})
 	if pos == -1 {
 		return
 	}
-	x, err := strconv.Atoi(s[pos+1:])
-	if err != nil {
-		return
-	}
-	if x != 1 {
-		fmt.Println("🚫") // U+1F6AB
-	}
-}
 
-func readit() chan []byte {
-	q := make(chan []byte)
-	go func() {
-		b := make([]byte, 1024)
-		n, _ := os.Stdin.Read(b)
-		q <- b[:n]
-	}()
-	return q
+	// parse the row and column numbers, separated by ';'
+	p := pos + 2
+rowLoop:
+	for {
+		if p == len(input) {
+			return
+		}
+		switch input[p] {
+		case ';':
+			break rowLoop
+		case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
+			row *= 10
+			row += int(input[p] - '0')
+		default:
+			return
+		}
+		p++
+	}
+	p++
+	for {
+		if p == len(input) {
+			return
+		}
+		switch input[p] {
+		case 'R':
+			count = p - pos + 1
+			ok = true
+			return
+		case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
+			col *= 10
+			col += int(input[p] - '0')
+		default:
+			return
+		}
+		p++
+	}
 }