diff --git a/chat/command.go b/chat/command.go index 10dbdc3..060bfee 100644 --- a/chat/command.go +++ b/chat/command.go @@ -8,7 +8,7 @@ import ( "strings" "github.com/shazow/ssh-chat/chat/message" - "github.com/shazow/ssh-chat/common" + "github.com/shazow/ssh-chat/set" ) // The error returned when an invalid command is issued. @@ -250,8 +250,9 @@ func InitCommands(c *Commands) { id := strings.TrimSpace(strings.TrimLeft(msg.Body(), "/ignore")) if id == "" { var names []string - msg.From().Ignored.Each(func(i common.Identified) { - names = append(names, i.Id()) + msg.From().Ignored.Each(func(_ string, item set.Item) error { + names = append(names, item.Key()) + return nil }) var systemMsg string diff --git a/chat/message/user.go b/chat/message/user.go index d8934d3..53180b6 100644 --- a/chat/message/user.go +++ b/chat/message/user.go @@ -9,7 +9,7 @@ import ( "sync" "time" - "github.com/shazow/ssh-chat/common" + "github.com/shazow/ssh-chat/set" ) const messageBuffer = 5 @@ -25,7 +25,7 @@ type User struct { joined time.Time msg chan Message done chan struct{} - Ignored *common.IdSet + Ignored *set.Set screen io.WriteCloser closeOnce sync.Once @@ -42,7 +42,7 @@ func NewUser(identity Identifier) *User { joined: time.Now(), msg: make(chan Message, messageBuffer), done: make(chan struct{}), - Ignored: common.NewIdSet(), + Ignored: set.New(), } u.SetColorIdx(rand.Int()) @@ -189,20 +189,20 @@ func (u *User) Send(m Message) error { return nil } -func (u *User) Ignore(identified common.Identified) error { - if identified == nil { +func (u *User) Ignore(other Identifier) error { + if other == nil { return errors.New("user is nil.") } - if identified.Id() == u.Id() { + if other.Id() == u.Id() { return errors.New("cannot ignore self.") } - if u.Ignored.In(identified) { + if u.Ignored.In(other.Id()) { return errors.New("user already ignored.") } - u.Ignored.Add(identified) + u.Ignored.Add(set.Itemize(other.Id(), other)) return nil } @@ -211,12 +211,7 @@ func (u *User) Unignore(id string) error { return errors.New("user is nil.") } - identified, err := u.Ignored.Get(id) - if err != nil { - return err - } - - return u.Ignored.Remove(identified) + return u.Ignored.Remove(id) } // Container for per-user configurations. diff --git a/chat/room.go b/chat/room.go index 8e3ebba..a0ba7cf 100644 --- a/chat/room.go +++ b/chat/room.go @@ -7,7 +7,7 @@ import ( "sync" "github.com/shazow/ssh-chat/chat/message" - "github.com/shazow/ssh-chat/common" + "github.com/shazow/ssh-chat/set" ) const historyLen = 20 @@ -35,8 +35,8 @@ type Room struct { closed bool closeOnce sync.Once - Members *common.IdSet - Ops *common.IdSet + Members *set.Set + Ops *set.Set } // NewRoom creates a new room. @@ -48,8 +48,8 @@ func NewRoom() *Room { history: message.NewHistory(historyLen), commands: *defaultCommands, - Members: common.NewIdSet(), - Ops: common.NewIdSet(), + Members: set.New(), + Ops: set.New(), } } @@ -62,8 +62,9 @@ func (r *Room) SetCommands(commands Commands) { func (r *Room) Close() { r.closeOnce.Do(func() { r.closed = true - r.Members.Each(func(m common.Identified) { - m.(*Member).Close() + r.Members.Each(func(_ string, item set.Item) error { + item.Value().(*Member).Close() + return nil }) r.Members.Clear() close(r.broadcast) @@ -96,10 +97,10 @@ func (r *Room) HandleMsg(m message.Message) { } r.history.Add(m) - r.Members.Each(func(u common.Identified) { - user := u.(*Member).User + r.Members.Each(func(_ string, item set.Item) (err error) { + user := item.Value().(*Member).User - if fromMsg != nil && user.Ignored.In(fromMsg.From()) { + if fromMsg != nil && user.Ignored.In(fromMsg.From().Id()) { // Skip because ignored return } @@ -115,6 +116,7 @@ func (r *Room) HandleMsg(m message.Message) { } } user.Send(m) + return }) } } @@ -145,40 +147,40 @@ func (r *Room) Join(u *message.User) (*Member, error) { if u.Id() == "" { return nil, ErrInvalidName } - member := Member{u} - err := r.Members.Add(&member) + member := &Member{u} + err := r.Members.Add(set.Itemize(u.Id(), member)) if err != nil { return nil, err } r.History(u) s := fmt.Sprintf("%s joined. (Connected: %d)", u.Name(), r.Members.Len()) r.Send(message.NewAnnounceMsg(s)) - return &member, nil + return member, nil } // Leave the room as a user, will announce. Mostly used during setup. func (r *Room) Leave(u message.Identifier) error { - err := r.Members.Remove(u) + err := r.Members.Remove(u.Id()) if err != nil { return err } - r.Ops.Remove(u) + r.Ops.Remove(u.Id()) s := fmt.Sprintf("%s left.", u.Name()) r.Send(message.NewAnnounceMsg(s)) return nil } // Rename member with a new identity. This will not call rename on the member. -func (r *Room) Rename(oldId string, identity message.Identifier) error { - if identity.Id() == "" { +func (r *Room) Rename(oldId string, u message.Identifier) error { + if u.Id() == "" { return ErrInvalidName } - err := r.Members.Replace(oldId, identity) + err := r.Members.Replace(oldId, set.Itemize(u.Id(), u)) if err != nil { return err } - s := fmt.Sprintf("%s is now known as %s.", oldId, identity.Id()) + s := fmt.Sprintf("%s is now known as %s.", oldId, u.Id()) r.Send(message.NewAnnounceMsg(s)) return nil } @@ -202,12 +204,12 @@ func (r *Room) MemberById(id string) (*Member, bool) { if err != nil { return nil, false } - return m.(*Member), true + return m.Value().(*Member), true } // IsOp returns whether a user is an operator in this room. func (r *Room) IsOp(u *message.User) bool { - return r.Ops.In(u) + return r.Ops.In(u.Id()) } // Topic of the room. @@ -223,10 +225,10 @@ func (r *Room) SetTopic(s string) { // NamesPrefix lists all members' names with a given prefix, used to query // for autocompletion purposes. func (r *Room) NamesPrefix(prefix string) []string { - members := r.Members.ListPrefix(prefix) - names := make([]string, len(members)) - for i, u := range members { - names[i] = u.(*Member).User.Name() + items := r.Members.ListPrefix(prefix) + names := make([]string, len(items)) + for i, item := range items { + names[i] = item.Value().(*Member).User.Name() } return names } diff --git a/chat/set_test.go b/chat/set_test.go index cf8135c..da944fc 100644 --- a/chat/set_test.go +++ b/chat/set_test.go @@ -4,35 +4,35 @@ import ( "testing" "github.com/shazow/ssh-chat/chat/message" - "github.com/shazow/ssh-chat/common" + "github.com/shazow/ssh-chat/set" ) func TestSet(t *testing.T) { var err error - s := common.NewIdSet() + s := set.New() u := message.NewUser(message.SimpleId("foo")) - if s.In(u) { + if s.In(u.Id()) { t.Errorf("Set should be empty.") } - err = s.Add(u) + err = s.Add(set.Itemize(u.Id(), u)) if err != nil { t.Error(err) } - if !s.In(u) { + if !s.In(u.Id()) { t.Errorf("Set should contain user.") } u2 := message.NewUser(message.SimpleId("bar")) - err = s.Add(u2) + err = s.Add(set.Itemize(u2.Id(), u2)) if err != nil { t.Error(err) } - err = s.Add(u2) - if err != common.ErrIdTaken { + err = s.Add(set.Itemize(u2.Id(), u2)) + if err != set.ErrCollision { t.Error(err) } diff --git a/common/set.go b/common/set.go deleted file mode 100644 index 8888e9d..0000000 --- a/common/set.go +++ /dev/null @@ -1,158 +0,0 @@ -package common - -import ( - "errors" - "strings" - "sync" -) - -// The error returned when an added id already exists in the set. -var ErrIdTaken = errors.New("id already taken") - -// The error returned when a requested item does not exist in the set. -var ErrIdentifiedMissing = errors.New("item does not exist") - -// Interface for an item storeable in the set -type Identified interface { - Id() string -} - -// Set with string lookup. -// TODO: Add trie for efficient prefix lookup? -type IdSet struct { - sync.RWMutex - lookup map[string]Identified -} - -// newIdSet creates a new set. -func NewIdSet() *IdSet { - return &IdSet{ - lookup: map[string]Identified{}, - } -} - -// Clear removes all items and returns the number removed. -func (s *IdSet) Clear() int { - s.Lock() - n := len(s.lookup) - s.lookup = map[string]Identified{} - s.Unlock() - return n -} - -// Len returns the size of the set right now. -func (s *IdSet) Len() int { - s.RLock() - defer s.RUnlock() - return len(s.lookup) -} - -// In checks if an item exists in this set. -func (s *IdSet) In(item Identified) bool { - id := s.normalize(item.Id()) - s.RLock() - _, ok := s.lookup[id] - s.RUnlock() - return ok -} - -// Get returns an item with the given Id. -func (s *IdSet) Get(id string) (Identified, error) { - s.RLock() - item, ok := s.lookup[s.normalize(id)] - s.RUnlock() - - if !ok { - return nil, ErrIdentifiedMissing - } - - return item, nil -} - -// Add item to this set if it does not exist already. -func (s *IdSet) Add(item Identified) error { - s.Lock() - defer s.Unlock() - - id := s.normalize(item.Id()) - _, found := s.lookup[id] - if found { - return ErrIdTaken - } - - s.lookup[id] = item - return nil -} - -// Remove item from this set. -func (s *IdSet) Remove(item Identified) error { - s.Lock() - defer s.Unlock() - id := s.normalize(item.Id()) - _, found := s.lookup[id] - if !found { - return ErrIdentifiedMissing - } - delete(s.lookup, id) - return nil -} - -// Replace item from old id with new Identified. -// Used for moving the same Identified to a new Id, such as a rename. -func (s *IdSet) Replace(oldId string, item Identified) error { - id := s.normalize(item.Id()) - oldId = s.normalize(oldId) - - s.Lock() - defer s.Unlock() - - // Check if it already exists - _, found := s.lookup[id] - if found { - return ErrIdTaken - } - - // Remove oldId - _, found = s.lookup[oldId] - if !found { - return ErrIdentifiedMissing - } - delete(s.lookup, oldId) - - // Add new Identified - s.lookup[id] = item - - return nil -} - -// Each loops over every item while holding a read lock and applies fn to each -// element. -func (s *IdSet) Each(fn func(item Identified)) { - s.RLock() - for _, item := range s.lookup { - fn(item) - } - s.RUnlock() -} - -// ListPrefix returns a list of items with a prefix, case insensitive. -func (s *IdSet) ListPrefix(prefix string) []Identified { - r := []Identified{} - prefix = strings.ToLower(prefix) - - s.RLock() - defer s.RUnlock() - - for id, item := range s.lookup { - if !strings.HasPrefix(string(id), prefix) { - continue - } - r = append(r, item) - } - - return r -} - -func (s *IdSet) normalize(id string) string { - return strings.ToLower(id) -} diff --git a/host.go b/host.go index 472f3d7..0e138e8 100644 --- a/host.go +++ b/host.go @@ -12,6 +12,7 @@ import ( "github.com/shazow/rateio" "github.com/shazow/ssh-chat/chat" "github.com/shazow/ssh-chat/chat/message" + "github.com/shazow/ssh-chat/set" "github.com/shazow/ssh-chat/sshd" ) @@ -126,7 +127,7 @@ func (h *Host) Connect(term *sshd.Terminal) { // Should the user be op'd on join? if h.isOp(term.Conn) { - h.Room.Ops.Add(member) + h.Room.Ops.Add(set.Itemize(member.Id(), member)) } ratelimit := rateio.NewSimpleLimiter(3, time.Second*3) @@ -493,7 +494,8 @@ func (h *Host) InitCommands(c *chat.Commands) { if !ok { return errors.New("user not found") } - room.Ops.Add(member) + room.Ops.Add(set.Itemize(member.Id(), member)) + id := member.Identifier.(*Identity) h.auth.Op(id.PublicKey(), until) diff --git a/host_test.go b/host_test.go index e340267..b83216c 100644 --- a/host_test.go +++ b/host_test.go @@ -11,6 +11,7 @@ import ( "testing" "github.com/shazow/ssh-chat/chat/message" + "github.com/shazow/ssh-chat/set" "github.com/shazow/ssh-chat/sshd" "golang.org/x/crypto/ssh" ) @@ -193,7 +194,7 @@ func TestHostKick(t *testing.T) { if member == nil { return errors.New("failed to load MemberById") } - host.Room.Ops.Add(member) + host.Room.Ops.Add(set.Itemize(member.Id(), member)) // Block until second client is here connected <- struct{}{} diff --git a/set/item.go b/set/item.go index a59fa13..25edf66 100644 --- a/set/item.go +++ b/set/item.go @@ -8,6 +8,23 @@ type Item interface { Value() interface{} } +type item struct { + key string + value interface{} +} + +func (item *item) Key() string { + return item.key +} + +func (item *item) Value() interface{} { + return item.value +} + +func Itemize(key string, value interface{}) Item { + return &item{key, value} +} + type StringItem string func (item StringItem) Key() string {