sshchat: Echo command messages with the new timestamp code.

chat/message/user.go

@@ -170,6 +170,8 @@ func (u *User) render(m Message) string {
 		if cfg.Bell {
 			out += Bel
 		}
+	case *CommandMsg:
+		out += m.RenderSelf(cfg)
 	default:
 		out += m.Render(cfg.Theme)
 	}

host.go

@@ -162,6 +162,20 @@ func (h *Host) Connect(term *sshd.Terminal) {
 
 		m := message.ParseInput(line, user)
 
+		// Gross hack to override line echo in golang.org/x/crypto/ssh/terminal
+		// It needs to live before we render the resulting message.
+		term.Write([]byte{
+			27, '[', 'A', // Up
+			27, '[', '2', 'K', // Clear line
+		})
+		// May the gods have mercy on our souls.
+
+		if m, ok := m.(*message.CommandMsg); ok {
+			// Other messages render themselves by the room, commands we'll
+			// have to re-echo ourselves manually.
+			user.HandleMsg(m)
+		}
+
 		// FIXME: Any reason to use h.room.Send(m) instead?
 		h.HandleMsg(m)
 
@@ -175,15 +189,6 @@ func (h *Host) Connect(term *sshd.Terminal) {
 			term.SetPrompt(GetPrompt(user))
 			user.SetHighlight(user.Name())
 		}
-
-		// Gross hack to override line echo in golang.org/x/crypto/ssh/terminal
-		term.Write([]byte{
-			27,       // keyEscape
-			'[', 'A', // Up
-			27,            // keyEscape
-			'[', '2', 'K', // Clear line
-		})
-		// May the gods have mercy on our souls.
 	}
 
 	err = h.Leave(user)

host_test.go

@@ -6,6 +6,7 @@ import (
 	"crypto/rsa"
 	"errors"
 	"io"
+	mathRand "math/rand"
 	"strings"
 	"testing"
 
@@ -15,22 +16,53 @@ import (
 )
 
 func stripPrompt(s string) string {
-	pos := strings.LastIndex(s, "\033[K")
+	// FIXME: Is there a better way to do this?
-	if pos < 0 {
+	if endPos := strings.Index(s, "\x1b[K "); endPos > 0 {
+		return s[endPos+3:]
+	}
+	if endPos := strings.Index(s, "\x1b[2K "); endPos > 0 {
+		return s[endPos+4:]
+	}
+	if endPos := strings.Index(s, "] "); endPos > 0 {
+		return s[endPos+2:]
+	}
 	return s
 }
-	return s[pos+3:]
+
+func TestStripPrompt(t *testing.T) {
+	tests := []struct {
+		Input string
+		Want  string
+	}{
+		{
+			Input: "\x1b[A\x1b[2K[quux] hello",
+			Want:  "hello",
+		},
+		{
+			Input: "[foo] \x1b[D\x1b[D\x1b[D\x1b[D\x1b[D\x1b[D\x1b[K * Guest1 joined. (Connected: 2)\r",
+			Want:  " * Guest1 joined. (Connected: 2)\r",
+		},
+	}
+
+	for i, tc := range tests {
+		if got, want := stripPrompt(tc.Input), tc.Want; got != want {
+			t.Errorf("case #%d:\n got: %q\nwant: %q", i, got, want)
+		}
+	}
 }
 
 func TestHostGetPrompt(t *testing.T) {
 	var expected, actual string
 
+	// Make the random colors consistent across tests
+	mathRand.Seed(1)
+
 	u := message.NewUser(&Identity{id: "foo"})
 
 	actual = GetPrompt(u)
 	expected = "[foo] "
 	if actual != expected {
-		t.Errorf("Got: %q; Expected: %q", actual, expected)
+		t.Errorf("Invalid host prompt:\n Got: %q;\nWant: %q", actual, expected)
 	}
 
 	u.SetConfig(message.UserConfig{
@@ -39,7 +71,7 @@ func TestHostGetPrompt(t *testing.T) {
 	actual = GetPrompt(u)
 	expected = "[\033[38;05;88mfoo\033[0m] "
 	if actual != expected {
-		t.Errorf("Got: %q; Expected: %q", actual, expected)
+		t.Errorf("Invalid host prompt:\n Got: %q;\nWant: %q", actual, expected)
 	}
 }
 
@@ -205,18 +237,23 @@ func TestHostKick(t *testing.T) {
 			// Change nicks, make sure op sticks
 			w.Write([]byte("/nick quux\r\n"))
 			scanner.Scan() // Prompt
+			scanner.Scan() // Prompt echo
 			scanner.Scan() // Nick change response
 
+			// Signal for the second client to connect
+			connected <- struct{}{}
+
 			// Block until second client is here
 			connected <- struct{}{}
 			scanner.Scan() // Connected message
 
 			w.Write([]byte("/kick bar\r\n"))
 			scanner.Scan() // Prompt
+			scanner.Scan() // Prompt echo
 
-			scanner.Scan()
+			scanner.Scan() // Kick result
 			if actual, expected := stripPrompt(scanner.Text()), " * bar was kicked by quux.\r"; actual != expected {
-				t.Errorf("Got %q; expected %q", actual, expected)
+				t.Errorf("Failed to detect kick:\n Got: %q;\nWant: %q", actual, expected)
 			}
 
 			kicked <- struct{}{}
@@ -231,6 +268,8 @@ func TestHostKick(t *testing.T) {
 	}()
 
 	go func() {
+		<-connected
+
 		// Second client
 		err := sshd.ConnectShell(addr, "bar", func(r io.Reader, w io.WriteCloser) error {
 			scanner := bufio.NewScanner(r)