From 8c7ea173ad9dc23b8bff237f497904bebc2d9141 Mon Sep 17 00:00:00 2001
From: yumaokao <yumaokaoh@gmail.com>
Date: Sat, 21 Mar 2020 19:57:39 +0800
Subject: [PATCH] sshd/terminal: Add fullwidth check for CJK in visualLength

---
 go.mod                         |  1 +
 go.sum                         |  3 +++
 sshd/terminal/terminal.go      |  8 +++++++-
 sshd/terminal/terminal_test.go | 27 +++++++++++++++++++++++++++
 4 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index c86bcd7..447b14d 100644
--- a/go.mod
+++ b/go.mod
@@ -7,6 +7,7 @@ require (
 	github.com/shazow/rateio v0.0.0-20150116013248-e8e00881e5c1
 	golang.org/x/crypto v0.0.0-20190320223903-b7391e95e576
 	golang.org/x/sys v0.0.0-20190321052220-f7bb7a8bee54
+	golang.org/x/text v0.3.2
 )
 
 go 1.13
diff --git a/go.sum b/go.sum
index 5ca9258..52d7b84 100644
--- a/go.sum
+++ b/go.sum
@@ -11,3 +11,6 @@ golang.org/x/crypto v0.0.0-20190320223903-b7391e95e576/go.mod h1:djNgcEr1/C05ACk
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190321052220-f7bb7a8bee54 h1:xe1/2UUJRmA9iDglQSlkx8c5n3twv58+K0mPpC2zmhA=
 golang.org/x/sys v0.0.0-20190321052220-f7bb7a8bee54/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/text v0.3.2 h1:tW2bmiBqwgJj/UpqtC8EpXEZVYOwU0yG4iWbprSVAcs=
+golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
diff --git a/sshd/terminal/terminal.go b/sshd/terminal/terminal.go
index d6650bb..20af6d2 100644
--- a/sshd/terminal/terminal.go
+++ b/sshd/terminal/terminal.go
@@ -10,6 +10,7 @@ import (
 	"strconv"
 	"sync"
 	"unicode/utf8"
+	"golang.org/x/text/width"
 )
 
 // EscapeCodes contains escape sequences that can be written to the terminal in
@@ -262,7 +263,7 @@ func (t *Terminal) moveCursorToPos(pos int) {
 		return
 	}
 
-	x := visualLength(t.prompt) + pos
+	x := visualLength(t.prompt) + visualLength(t.line[:pos])
 	y := x / t.termWidth
 	x = x % t.termWidth
 
@@ -351,6 +352,7 @@ func (t *Terminal) setLine(newLine []rune, newPos int) {
 		for i := len(newLine); i < len(t.line); i++ {
 			t.writeLine(space)
 		}
+		t.line = newLine
 		t.moveCursorToPos(newPos)
 	}
 	t.line = newLine
@@ -462,6 +464,10 @@ func visualLength(runes []rune) int {
 			inEscapeSeq = true
 		default:
 			length++
+			kind := width.LookupRune(r).Kind()
+			if kind == width.EastAsianFullwidth || kind == width.EastAsianWide {
+				length++
+			}
 		}
 	}
 
diff --git a/sshd/terminal/terminal_test.go b/sshd/terminal/terminal_test.go
index f7eb958..13e9133 100644
--- a/sshd/terminal/terminal_test.go
+++ b/sshd/terminal/terminal_test.go
@@ -12,6 +12,7 @@ import (
 	"os"
 	"runtime"
 	"testing"
+	"unicode/utf8"
 )
 
 type MockTerminal struct {
@@ -406,3 +407,29 @@ func TestOutputNewlines(t *testing.T) {
 		t.Errorf("incorrect output: was %q, expected %q", output, expected)
 	}
 }
+
+func TestTerminalvisualLength(t *testing.T) {
+	var tests = []struct {
+		input string
+		want  int
+	}{
+		{"hello world", 11},
+		{"babalala", 8},
+		{"端子", 4},
+		{"を搭載", 6},
+		{"baba端子lalaを搭載", 18},
+	}
+	for _, test := range tests {
+		var runes []rune
+		for i, w := 0, 0; i < len(test.input); i += w {
+			runeValue, width := utf8.DecodeRuneInString(test.input[i:])
+			runes = append(runes, runeValue)
+			w = width
+		}
+		output := visualLength(runes)
+		if output != test.want {
+			t.Errorf("incorrect [%s] output: was %d, expected %d",
+				test.input, output, test.want)
+		}
+	}
+}