From 601a95c1cd2d82e6429bc21483569ff885a1ffc9 Mon Sep 17 00:00:00 2001
From: Andrey Petrov <andrey.petrov@shazow.net>
Date: Mon, 22 Dec 2014 22:54:29 -0800
Subject: [PATCH] Progress trying to make things less buggy, not much.

---
 chat/channel.go | 10 +++++++++-
 chat/set.go     |  9 +++++++++
 chat/theme.go   |  3 +++
 chat/user.go    |  8 +++++++-
 cmd.go          | 11 ++++++++---
 5 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/chat/channel.go b/chat/channel.go
index b52fb47..d9ace6e 100644
--- a/chat/channel.go
+++ b/chat/channel.go
@@ -32,7 +32,11 @@ func NewChannel() *Channel {
 					// Skip
 					return
 				}
-				user.Send(m)
+				err := user.Send(m)
+				if err != nil {
+					ch.Leave(user)
+					user.Close()
+				}
 			})
 		}
 	}()
@@ -41,6 +45,10 @@ func NewChannel() *Channel {
 }
 
 func (ch *Channel) Close() {
+	ch.users.Each(func(u Item) {
+		u.(*User).Close()
+	})
+	ch.users.Clear()
 	close(ch.broadcast)
 }
 
diff --git a/chat/set.go b/chat/set.go
index 8c01f87..c5aef93 100644
--- a/chat/set.go
+++ b/chat/set.go
@@ -34,6 +34,15 @@ func NewSet() *Set {
 	}
 }
 
+// Remove all items and return the number removed
+func (s *Set) Clear() int {
+	s.Lock()
+	n := len(s.lookup)
+	s.lookup = map[Id]Item{}
+	s.Unlock()
+	return n
+}
+
 // Size of the set right now
 func (s *Set) Len() int {
 	return len(s.lookup)
diff --git a/chat/theme.go b/chat/theme.go
index d187256..f592327 100644
--- a/chat/theme.go
+++ b/chat/theme.go
@@ -23,6 +23,9 @@ const (
 
 	// Invert inverts the following text
 	Invert = "\033[7m"
+
+	// Newline
+	Newline = "\r\n"
 )
 
 // Interface for Colors
diff --git a/chat/user.go b/chat/user.go
index ffc4e61..bd59981 100644
--- a/chat/user.go
+++ b/chat/user.go
@@ -19,6 +19,7 @@ type User struct {
 	joined   time.Time
 	msg      chan Message
 	done     chan struct{}
+	closed   bool
 	Config   UserConfig
 }
 
@@ -74,6 +75,7 @@ func (u *User) Wait() {
 
 // Disconnect user, stop accepting messages
 func (u *User) Close() {
+	u.closed = true
 	close(u.done)
 	close(u.msg)
 }
@@ -94,7 +96,7 @@ func (u *User) ConsumeOne(out io.Writer) {
 
 func (u *User) consumeMsg(m Message, out io.Writer) {
 	s := m.Render(u.Config.Theme)
-	_, err := out.Write([]byte(s))
+	_, err := out.Write([]byte(s + Newline))
 	if err != nil {
 		logger.Printf("Write failed to %s, closing: %s", u.Name(), err)
 		u.Close()
@@ -103,6 +105,10 @@ func (u *User) consumeMsg(m Message, out io.Writer) {
 
 // Add message to consume by user
 func (u *User) Send(m Message) error {
+	if u.closed {
+		return ErrUserClosed
+	}
+
 	select {
 	case u.msg <- m:
 	default:
diff --git a/cmd.go b/cmd.go
index 6491cdc..9d242f5 100644
--- a/cmd.go
+++ b/cmd.go
@@ -104,11 +104,18 @@ func main() {
 			go func() {
 				defer term.Close()
 				name := term.Conn.User()
-				term.SetPrompt(fmt.Sprintf("[%s]", name))
+				term.SetPrompt(fmt.Sprintf("[%s] ", name))
 				// TODO: term.AutoCompleteCallback = ...
 				user := chat.NewUserScreen(name, term)
+				defer user.Close()
 				channel.Join(user)
 
+				go func() {
+					// FIXME: This isn't working.
+					user.Wait()
+					channel.Leave(user)
+				}()
+
 				for {
 					// TODO: Handle commands etc?
 					line, err := term.ReadLine()
@@ -120,8 +127,6 @@ func main() {
 				}
 
 				// TODO: Handle disconnect sooner (currently closes channel before removing)
-				channel.Leave(user)
-				user.Close()
 			}()
 		}
 	}()