From 9bf1f53445946fd8496cb1df9eaeb4bc88cd630f Mon Sep 17 00:00:00 2001 From: Andrey Petrov Date: Fri, 15 Jul 2016 16:22:25 -0400 Subject: [PATCH] Fixed /kick command to actually close target --- auth.go | 2 -- chat/message/message.go | 36 +++++++++++++++++----------------- chat/message/user.go | 43 +++++++++++++++++++++++++++-------------- chat/set.go | 4 +++- host.go | 12 +++++------- host_test.go | 2 +- set.go | 4 +++- 7 files changed, 59 insertions(+), 44 deletions(-) diff --git a/auth.go b/auth.go index 4d27ca4..7095260 100644 --- a/auth.go +++ b/auth.go @@ -3,7 +3,6 @@ package sshchat import ( "errors" "net" - "sync" "time" "github.com/shazow/ssh-chat/sshd" @@ -36,7 +35,6 @@ func newAuthAddr(addr net.Addr) string { // Auth stores lookups for bans, whitelists, and ops. It implements the sshd.Auth interface. type Auth struct { - sync.RWMutex bannedAddr *Set banned *Set whitelist *Set diff --git a/chat/message/message.go b/chat/message/message.go index 491e8b6..6716be5 100644 --- a/chat/message/message.go +++ b/chat/message/message.go @@ -48,21 +48,21 @@ func NewMsg(body string) *Msg { } // Render message based on a theme. -func (m *Msg) Render(t *Theme) string { +func (m Msg) Render(t *Theme) string { // TODO: Render based on theme // TODO: Cache based on theme return m.String() } -func (m *Msg) String() string { +func (m Msg) String() string { return m.body } -func (m *Msg) Command() string { +func (m Msg) Command() string { return "" } -func (m *Msg) Timestamp() time.Time { +func (m Msg) Timestamp() time.Time { return m.timestamp } @@ -72,8 +72,8 @@ type PublicMsg struct { from *User } -func NewPublicMsg(body string, from *User) *PublicMsg { - return &PublicMsg{ +func NewPublicMsg(body string, from *User) PublicMsg { + return PublicMsg{ Msg: Msg{ body: body, timestamp: time.Now(), @@ -82,11 +82,11 @@ func NewPublicMsg(body string, from *User) *PublicMsg { } } -func (m *PublicMsg) From() *User { +func (m PublicMsg) From() *User { return m.from } -func (m *PublicMsg) ParseCommand() (*CommandMsg, bool) { +func (m PublicMsg) ParseCommand() (*CommandMsg, bool) { // Check if the message is a command if !strings.HasPrefix(m.body, "/") { return nil, false @@ -104,7 +104,7 @@ func (m *PublicMsg) ParseCommand() (*CommandMsg, bool) { return &msg, true } -func (m *PublicMsg) Render(t *Theme) string { +func (m PublicMsg) Render(t *Theme) string { if t == nil { return m.String() } @@ -112,7 +112,7 @@ func (m *PublicMsg) Render(t *Theme) string { return fmt.Sprintf("%s: %s", t.ColorName(m.from), m.body) } -func (m *PublicMsg) RenderFor(cfg UserConfig) string { +func (m PublicMsg) RenderFor(cfg UserConfig) string { if cfg.Highlight == nil || cfg.Theme == nil { return m.Render(cfg.Theme) } @@ -128,7 +128,7 @@ func (m *PublicMsg) RenderFor(cfg UserConfig) string { return fmt.Sprintf("%s: %s", cfg.Theme.ColorName(m.from), body) } -func (m *PublicMsg) String() string { +func (m PublicMsg) String() string { return fmt.Sprintf("%s: %s", m.from.Name(), m.body) } @@ -164,9 +164,9 @@ type PrivateMsg struct { to *User } -func NewPrivateMsg(body string, from *User, to *User) *PrivateMsg { - return &PrivateMsg{ - PublicMsg: *NewPublicMsg(body, from), +func NewPrivateMsg(body string, from *User, to *User) PrivateMsg { + return PrivateMsg{ + PublicMsg: NewPublicMsg(body, from), to: to, } } @@ -242,19 +242,19 @@ func (m *AnnounceMsg) String() string { } type CommandMsg struct { - *PublicMsg + PublicMsg command string args []string } -func (m *CommandMsg) Command() string { +func (m CommandMsg) Command() string { return m.command } -func (m *CommandMsg) Args() []string { +func (m CommandMsg) Args() []string { return m.args } -func (m *CommandMsg) Body() string { +func (m CommandMsg) Body() string { return m.body } diff --git a/chat/message/user.go b/chat/message/user.go index 73d384e..25098dc 100644 --- a/chat/message/user.go +++ b/chat/message/user.go @@ -18,14 +18,16 @@ var ErrUserClosed = errors.New("user closed") // User definition, implemented set Item interface and io.Writer type User struct { Identifier - Config UserConfig - colorIdx int - joined time.Time - msg chan Message - done chan struct{} - replyTo *User // Set when user gets a /msg, for replying. - closed bool - closeOnce sync.Once + Config UserConfig + colorIdx int + joined time.Time + msg chan Message + done chan struct{} + + mu sync.Mutex + replyTo *User // Set when user gets a /msg, for replying. + screen io.Closer + closed bool } func NewUser(identity Identifier) *User { @@ -41,8 +43,9 @@ func NewUser(identity Identifier) *User { return &u } -func NewUserScreen(identity Identifier, screen io.Writer) *User { +func NewUserScreen(identity Identifier, screen io.WriteCloser) *User { u := NewUser(identity) + u.screen = screen go u.Consume(screen) return u @@ -82,11 +85,20 @@ func (u *User) Wait() { // Disconnect user, stop accepting messages func (u *User) Close() { - u.closeOnce.Do(func() { - u.closed = true - close(u.done) - close(u.msg) - }) + u.mu.Lock() + defer u.mu.Unlock() + + if u.closed { + return + } + + u.closed = true + close(u.done) + close(u.msg) + + if u.screen != nil { + u.screen.Close() + } } // Consume message buffer into an io.Writer. Will block, should be called in a @@ -136,6 +148,9 @@ func (u *User) HandleMsg(m Message, out io.Writer) { // Add message to consume by user func (u *User) Send(m Message) error { + u.mu.Lock() + defer u.mu.Unlock() + if u.closed { return ErrUserClosed } diff --git a/chat/set.go b/chat/set.go index b7ec5b3..438fe97 100644 --- a/chat/set.go +++ b/chat/set.go @@ -20,8 +20,8 @@ type identified interface { // Set with string lookup. // TODO: Add trie for efficient prefix lookup? type idSet struct { - lookup map[string]identified sync.RWMutex + lookup map[string]identified } // newIdSet creates a new set. @@ -42,6 +42,8 @@ func (s *idSet) Clear() int { // Len returns the size of the set right now. func (s *idSet) Len() int { + s.RLock() + defer s.RUnlock() return len(s.lookup) } diff --git a/host.go b/host.go index b5adbc2..4505dc1 100644 --- a/host.go +++ b/host.go @@ -90,12 +90,10 @@ func (h *Host) Connect(term *sshd.Terminal) { id := NewIdentity(term.Conn) user := message.NewUserScreen(id, term) user.Config.Theme = &h.theme - go func() { - // Close term once user is closed. - user.Wait() - term.Close() - }() + + // Close term once user is closed. defer user.Close() + defer term.Close() h.mu.Lock() motd := h.motd @@ -285,7 +283,7 @@ func (h *Host) InitCommands(c *chat.Commands) { } m := message.NewPrivateMsg(strings.Join(args[1:], " "), msg.From(), target) - room.Send(m) + room.Send(&m) return nil }, }) @@ -307,7 +305,7 @@ func (h *Host) InitCommands(c *chat.Commands) { } m := message.NewPrivateMsg(strings.Join(args, " "), msg.From(), target) - room.Send(m) + room.Send(&m) return nil }, }) diff --git a/host_test.go b/host_test.go index a869ebb..a34e683 100644 --- a/host_test.go +++ b/host_test.go @@ -206,7 +206,7 @@ func TestHostKick(t *testing.T) { <-connected // Consume while we're connected. Should break when kicked. - ioutil.ReadAll(r) + ioutil.ReadAll(r) // XXX? }) if err != nil { t.Fatal(err) diff --git a/set.go b/set.go index f0d607c..3e29a57 100644 --- a/set.go +++ b/set.go @@ -25,8 +25,8 @@ type setValue interface { // Set with expire-able keys type Set struct { - lookup map[string]setValue sync.Mutex + lookup map[string]setValue } // NewSet creates a new set. @@ -38,6 +38,8 @@ func NewSet() *Set { // Len returns the size of the set right now. func (s *Set) Len() int { + s.Lock() + defer s.Unlock() return len(s.lookup) }