Add commands.

client.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"strings"
 
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/terminal"
@@ -9,6 +10,24 @@ import (
 
 const MSG_BUFFER = 10
 
+const HELP_TEXT = `-> Available commands:
+  /about
+  /exit
+  /help
+  /list
+  /nick $NAME
+`
+
+const ABOUT_TEXT = `-> ssh-chat is made by @shazow.
+
+  It is a custom ssh server built in Go to serve a chat experience
+  instead of a shell.
+
+  Source: https://github.com/shazow/ssh-chat
+
+  For more, visit shazow.net or follow at twitter.com/shazow
+`
+
 type Client struct {
 	Server     *Server
 	Conn       *ssh.ServerConn
@@ -19,15 +38,11 @@ type Client struct {
 	termHeight int
 }
 
-func NewClient(server *Server, conn *ssh.ServerConn, name string) *Client {
-	if name == "" {
-		name = "Anonymoose"
-	}
-
+func NewClient(server *Server, conn *ssh.ServerConn) *Client {
 	return &Client{
 		Server: server,
 		Conn:   conn,
-		Name:   name,
+		Name:   conn.User(),
 		Msg:    make(chan string, MSG_BUFFER),
 	}
 }
@@ -42,10 +57,9 @@ func (c *Client) Resize(width int, height int) error {
 	return nil
 }
 
-func (c *Client) sendWelcome() {
-	msg := fmt.Sprintf("Welcome to ssh-chat. Enter /help for more.\r\n")
-	c.Msg <- msg
-
+func (c *Client) Rename(name string) {
+	c.Name = name
+	c.term.SetPrompt(fmt.Sprintf("[%s] ", name))
 }
 
 func (c *Client) handleShell(channel ssh.Channel) {
@@ -63,12 +77,31 @@ func (c *Client) handleShell(channel ssh.Channel) {
 			break
 		}
 
-		switch line {
-		case "/exit":
-			channel.Close()
+		parts := strings.SplitN(line, " ", 2)
+		isCmd := strings.HasPrefix(parts[0], "/")
+
+		if isCmd {
+			switch parts[0] {
+			case "/exit":
+				channel.Close()
+			case "/help":
+				c.Msg <- HELP_TEXT
+			case "/about":
+				c.Msg <- ABOUT_TEXT
+			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 "/list":
+				c.Msg <- fmt.Sprintf("-> Connected: %s\r\n", strings.Join(c.Server.List(nil), ","))
+			default:
+				c.Msg <- fmt.Sprintf("-> Invalid command: %s\r\n", line)
+			}
+			continue
 		}
 
-		//c.term.Write(c.term.Escape.Reset)
 		msg := fmt.Sprintf("%s: %s\r\n", c.Name, line)
 		c.Server.Broadcast(msg, c)
 	}

server.go

@@ -9,12 +9,15 @@ import (
 	"golang.org/x/crypto/ssh"
 )
 
+type Clients map[string]*Client
+
 type Server struct {
 	sshConfig *ssh.ServerConfig
 	sshSigner *ssh.Signer
 	done      chan struct{}
-	clients   map[*Client]struct{}
+	clients   Clients
 	lock      sync.Mutex
+	count     int
 }
 
 func NewServer(privateKey []byte) (*Server, error) {
@@ -39,7 +42,8 @@ func NewServer(privateKey []byte) (*Server, error) {
 		sshConfig: &config,
 		sshSigner: &signer,
 		done:      make(chan struct{}),
-		clients:   map[*Client]struct{}{},
+		clients:   Clients{},
+		count:     0,
 	}
 
 	return &server, nil
@@ -47,7 +51,8 @@ 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"))
-	for client := range s.clients {
+
+	for _, client := range s.clients {
 		if except != nil && client == except {
 			continue
 		}
@@ -55,6 +60,65 @@ func (s *Server) Broadcast(msg string, except *Client) {
 	}
 }
 
+func (s *Server) Add(client *Client) {
+	client.Msg <- fmt.Sprintf("-> Welcome to ssh-chat. Enter /help for more.\r\n")
+
+	s.lock.Lock()
+	s.count++
+
+	_, collision := s.clients[client.Name]
+	if collision {
+		newName := fmt.Sprintf("Guest%d", s.count)
+		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
+	}
+
+	s.clients[client.Name] = client
+	num := len(s.clients)
+	s.lock.Unlock()
+
+	s.Broadcast(fmt.Sprintf("* %s joined. (Total connected: %d)\r\n", client.Name, num), nil)
+}
+
+func (s *Server) Remove(client *Client) {
+	s.lock.Lock()
+	delete(s.clients, client.Name)
+	s.lock.Unlock()
+
+	s.Broadcast(fmt.Sprintf("* %s left.\r\n", client.Name), nil)
+}
+
+func (s *Server) Rename(client *Client, newName string) {
+	s.lock.Lock()
+
+	_, collision := s.clients[newName]
+	if collision {
+		client.Msg <- fmt.Sprintf("-> %s is not available.\r\n", newName)
+		s.lock.Unlock()
+		return
+	}
+	delete(s.clients, client.Name)
+	oldName := client.Name
+	client.Rename(newName)
+	s.clients[client.Name] = client
+	s.lock.Unlock()
+
+	s.Broadcast(fmt.Sprintf("* %s is now known as %s.\r\n", oldName, newName), nil)
+}
+
+func (s *Server) List(prefix *string) []string {
+	r := []string{}
+
+	for name, _ := range s.clients {
+		if prefix != nil && !strings.HasPrefix(name, *prefix) {
+			continue
+		}
+		r = append(r, name)
+	}
+
+	return r
+}
+
 func (s *Server) Start(laddr string) error {
 	// Once a ServerConfig has been configured, connections can be
 	// accepted.
@@ -88,25 +152,13 @@ func (s *Server) Start(laddr string) error {
 
 				go ssh.DiscardRequests(requests)
 
-				client := NewClient(s, sshConn, sshConn.User())
-				// TODO: mutex this
-
-				s.lock.Lock()
-				s.clients[client] = struct{}{}
-				num := len(s.clients)
-				s.lock.Unlock()
-
-				client.sendWelcome()
-
-				s.Broadcast(fmt.Sprintf("* %s joined. (Total connected: %d)\r\n", client.Name, num), nil)
+				client := NewClient(s, sshConn)
+				s.Add(client)
 
 				go func() {
+					// Block until done, then remove.
 					sshConn.Wait()
-					s.lock.Lock()
-					delete(s.clients, client)
-					s.lock.Unlock()
-
-					s.Broadcast(fmt.Sprintf("* %s left.\r\n", client.Name), nil)
+					s.Remove(client)
 				}()
 
 				go client.handleChannels(channels)
@@ -123,7 +175,7 @@ func (s *Server) Start(laddr string) error {
 }
 
 func (s *Server) Stop() {
-	for client := range s.clients {
+	for _, client := range s.clients {
 		client.Conn.Close()
 	}