More commands, history, tests.

client.go

@@ -3,39 +3,43 @@ package main
 import (
 	"fmt"
 	"strings"
+	"time"
 
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/terminal"
 )
 
-const MSG_BUFFER = 10
+const MSG_BUFFER int = 10
 
-const HELP_TEXT = `-> Available commands:
+const HELP_TEXT string = `-> Available commands:
-  /about
+   /about
-  /exit
+   /exit
-  /help
+   /help
-  /list
+   /list
-  /nick $NAME
+   /nick $NAME
+   /whois $NAME
 `
 
-const ABOUT_TEXT = `-> ssh-chat is made by @shazow.
+const ABOUT_TEXT string = `-> ssh-chat is made by @shazow.
 
-  It is a custom ssh server built in Go to serve a chat experience
+   It is a custom ssh server built in Go to serve a chat experience
-  instead of a shell.
+   instead of a shell.
 
-  Source: https://github.com/shazow/ssh-chat
+   Source: https://github.com/shazow/ssh-chat
 
-  For more, visit shazow.net or follow at twitter.com/shazow
+   For more, visit shazow.net or follow at twitter.com/shazow
 `
 
 type Client struct {
-	Server     *Server
+	Server        *Server
-	Conn       *ssh.ServerConn
+	Conn          *ssh.ServerConn
-	Msg        chan string
+	Msg           chan string
-	Name       string
+	Name          string
-	term       *terminal.Terminal
+	Op            bool
-	termWidth  int
+	term          *terminal.Terminal
-	termHeight int
+	termWidth     int
+	termHeight    int
+	silencedUntil time.Time
 }
 
 func NewClient(server *Server, conn *ssh.ServerConn) *Client {
@@ -47,6 +51,24 @@ func NewClient(server *Server, conn *ssh.ServerConn) *Client {
 	}
 }
 
+func (c *Client) Write(msg string) {
+	c.term.Write([]byte(msg))
+}
+
+func (c *Client) WriteLines(msg []string) {
+	for _, line := range msg {
+		c.Write(line + "\r\n")
+	}
+}
+
+func (c *Client) IsSilenced() bool {
+	return c.silencedUntil.After(time.Now())
+}
+
+func (c *Client) Silence(d time.Duration) {
+	c.silencedUntil = time.Now().Add(d)
+}
+
 func (c *Client) Resize(width int, height int) error {
 	err := c.term.SetSize(width, height)
 	if err != nil {
@@ -67,7 +89,7 @@ func (c *Client) handleShell(channel ssh.Channel) {
 
 	go func() {
 		for msg := range c.Msg {
-			c.term.Write([]byte(msg))
+			c.Write(msg)
 		}
 	}()
 
@@ -85,15 +107,22 @@ func (c *Client) handleShell(channel ssh.Channel) {
 			case "/exit":
 				channel.Close()
 			case "/help":
-				c.Msg <- HELP_TEXT
+				c.WriteLines(strings.Split(HELP_TEXT, "\n"))
 			case "/about":
-				c.Msg <- ABOUT_TEXT
+				c.WriteLines(strings.Split(ABOUT_TEXT, "\n"))
 			case "/nick":
 				if len(parts) == 2 {
 					c.Server.Rename(c, parts[1])
 				} else {
 					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /nick $NAME\r\n")
 				}
+			case "/whois":
+				if len(parts) == 2 {
+					client := c.Server.Who(parts[1])
+					c.Msg <- fmt.Sprintf("-> %s is %s via %s\r\n", client.Name, client.Conn.RemoteAddr(), client.Conn.ClientVersion())
+				} else {
+					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /whois $NAME\r\n")
+				}
 			case "/list":
 				c.Msg <- fmt.Sprintf("-> Connected: %s\r\n", strings.Join(c.Server.List(nil), ","))
 			default:
@@ -103,6 +132,10 @@ func (c *Client) handleShell(channel ssh.Channel) {
 		}
 
 		msg := fmt.Sprintf("%s: %s\r\n", c.Name, line)
+		if c.IsSilenced() {
+			c.Msg <- fmt.Sprintf("-> Message rejected, silenced.")
+			continue
+		}
 		c.Server.Broadcast(msg, c)
 	}
 

history.go

@@ -0,0 +1,44 @@
+package main
+
+type History struct {
+	entries []string
+	head    int
+	size    int
+}
+
+func NewHistory(size int) *History {
+	return &History{
+		entries: make([]string, size),
+	}
+}
+
+func (h *History) Add(entry string) {
+	max := cap(h.entries)
+	h.head = (h.head + 1) % max
+	h.entries[h.head] = entry
+	if h.size < max {
+		h.size++
+	}
+}
+
+func (h *History) Len() int {
+	return h.size
+}
+
+func (h *History) Get(num int) []string {
+	max := cap(h.entries)
+	if num > h.size {
+		num = h.size
+	}
+
+	r := make([]string, num)
+	for i := 0; i < num; i++ {
+		idx := (h.head - i) % max
+		if idx < 0 {
+			idx += max
+		}
+		r[num-i-1] = h.entries[idx]
+	}
+
+	return r
+}

history_test.go

@@ -0,0 +1,53 @@
+package main
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestHistory(t *testing.T) {
+	var r, expected []string
+	var size int
+
+	h := NewHistory(5)
+
+	r = h.Get(10)
+	expected = []string{}
+	if !reflect.DeepEqual(r, expected) {
+		t.Errorf("Got: %v, Expected: %v", r, expected)
+	}
+
+	h.Add("1")
+
+	if size = h.Len(); size != 1 {
+		t.Errorf("Wrong size: %v", size)
+	}
+
+	r = h.Get(1)
+	expected = []string{"1"}
+	if !reflect.DeepEqual(r, expected) {
+		t.Errorf("Got: %v, Expected: %v", r, expected)
+	}
+
+	h.Add("2")
+	h.Add("3")
+	h.Add("4")
+	h.Add("5")
+	h.Add("6")
+
+	if size = h.Len(); size != 5 {
+		t.Errorf("Wrong size: %v", size)
+	}
+
+	r = h.Get(2)
+	expected = []string{"5", "6"}
+	if !reflect.DeepEqual(r, expected) {
+		t.Errorf("Got: %v, Expected: %v", r, expected)
+	}
+
+	r = h.Get(10)
+	expected = []string{"2", "3", "4", "5", "6"}
+	if !reflect.DeepEqual(r, expected) {
+		t.Errorf("Got: %v, Expected: %v", r, expected)
+	}
+}

server.go

@@ -3,12 +3,18 @@ package main
 import (
 	"fmt"
 	"net"
+	"regexp"
 	"strings"
 	"sync"
 
 	"golang.org/x/crypto/ssh"
 )
 
+const MAX_NAME_LENGTH = 32
+const HISTORY_LEN = 20
+
+var RE_STRIP_NAME = regexp.MustCompile("[[:^alpha:]]")
+
 type Clients map[string]*Client
 
 type Server struct {
@@ -18,6 +24,7 @@ type Server struct {
 	clients   Clients
 	lock      sync.Mutex
 	count     int
+	history   *History
 }
 
 func NewServer(privateKey []byte) (*Server, error) {
@@ -32,7 +39,7 @@ func NewServer(privateKey []byte) (*Server, error) {
 			return nil, nil
 		},
 		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
-			// fingerprint := md5.Sum(key.Marshal()
+			// fingerprint := md5.Sum(key.Marshal())
 			return nil, nil
 		},
 	}
@@ -44,6 +51,7 @@ func NewServer(privateKey []byte) (*Server, error) {
 		done:      make(chan struct{}),
 		clients:   Clients{},
 		count:     0,
+		history:   NewHistory(HISTORY_LEN),
 	}
 
 	return &server, nil
@@ -51,6 +59,7 @@ func NewServer(privateKey []byte) (*Server, error) {
 
 func (s *Server) Broadcast(msg string, except *Client) {
 	logger.Debugf("Broadcast to %d: %s", len(s.clients), strings.TrimRight(msg, "\r\n"))
+	s.history.Add(msg)
 
 	for _, client := range s.clients {
 		if except != nil && client == except {
@@ -61,18 +70,23 @@ func (s *Server) Broadcast(msg string, except *Client) {
 }
 
 func (s *Server) Add(client *Client) {
+	for _, msg := range s.history.Get(10) {
+		if msg == "" {
+			continue
+		}
+		client.Msg <- msg
+	}
 	client.Msg <- fmt.Sprintf("-> Welcome to ssh-chat. Enter /help for more.\r\n")
 
 	s.lock.Lock()
 	s.count++
 
-	_, collision := s.clients[client.Name]
+	newName, err := s.proposeName(client.Name)
-	if collision {
+	if err != nil {
-		newName := fmt.Sprintf("Guest%d", s.count)
+		client.Msg <- fmt.Sprintf("-> Your name '%s' is not avaialble, renamed to '%s'. Use /nick <name> to change it.\r\n", client.Name, newName)
-		client.Msg <- fmt.Sprintf("-> Your name '%s' was taken, renamed to '%s'. Use /nick <name> to change it.\r\n", client.Name, newName)
-		client.Name = newName
 	}
 
+	client.Name = newName
 	s.clients[client.Name] = client
 	num := len(s.clients)
 	s.lock.Unlock()
@@ -88,15 +102,36 @@ func (s *Server) Remove(client *Client) {
 	s.Broadcast(fmt.Sprintf("* %s left.\r\n", client.Name), nil)
 }
 
+func (s *Server) proposeName(name string) (string, error) {
+	// Assumes caller holds lock.
+	var err error
+	name = RE_STRIP_NAME.ReplaceAllString(name, "")
+
+	if len(name) > MAX_NAME_LENGTH {
+		name = name[:MAX_NAME_LENGTH]
+	} else if len(name) == 0 {
+		name = fmt.Sprintf("Guest%d", s.count)
+	}
+
+	_, collision := s.clients[name]
+	if collision {
+		err = fmt.Errorf("%s is not available.", name)
+		name = fmt.Sprintf("Guest%d", s.count)
+	}
+
+	return name, err
+}
+
 func (s *Server) Rename(client *Client, newName string) {
 	s.lock.Lock()
 
-	_, collision := s.clients[newName]
+	newName, err := s.proposeName(newName)
-	if collision {
+	if err != nil {
-		client.Msg <- fmt.Sprintf("-> %s is not available.\r\n", newName)
+		client.Msg <- fmt.Sprintf("-> %s\r\n", err)
 		s.lock.Unlock()
 		return
 	}
+
 	delete(s.clients, client.Name)
 	oldName := client.Name
 	client.Rename(newName)
@@ -119,6 +154,10 @@ func (s *Server) List(prefix *string) []string {
 	return r
 }
 
+func (s *Server) Who(name string) *Client {
+	return s.clients[name]
+}
+
 func (s *Server) Start(laddr string) error {
 	// Once a ServerConfig has been configured, connections can be
 	// accepted.