From 4e23e5fc35e9dc4328930d5cfb92b361b471c86f Mon Sep 17 00:00:00 2001 From: Chris Rhodes Date: Thu, 31 Dec 2015 17:29:16 -0800 Subject: [PATCH] Implement state tracking. Currently maintains a list of Guilds (Members/Channels) and PrivateChannels. --- discord.go | 4 +- state.go | 195 +++++++++++++++++++++++++++++++++++++++++++++++++++++ structs.go | 22 +++--- wsapi.go | 83 +++++++++++++---------- 4 files changed, 257 insertions(+), 47 deletions(-) create mode 100644 state.go diff --git a/discord.go b/discord.go index 2f7c97d..799e725 100644 --- a/discord.go +++ b/discord.go @@ -61,7 +61,9 @@ type voiceUDP struct { func New(args ...interface{}) (s *Session, err error) { // Create an empty Session interface. - s = &Session{} + s = &Session{ + State: NewState(), + } // If no arguments are passed return the empty Session interface. // Later I will add default values, if appropriate. diff --git a/state.go b/state.go new file mode 100644 index 0000000..5645aa3 --- /dev/null +++ b/state.go @@ -0,0 +1,195 @@ +package discordgo + +import "errors" + +// NewState creates an empty state. +func NewState() *State { + return &State{ + Ready: Ready{ + PrivateChannels: []Channel{}, + Guilds: []Guild{}, + }, + } +} + +// OnReady takes a Ready event and updates all internal state. +func (s *State) OnReady(r *Ready) { + s.Ready = *r +} + +// AddGuild adds a guild to the current world state, or +// updates it if it already exists. +func (s *State) AddGuild(guild *Guild) { + for _, g := range s.Guilds { + if g.ID == guild.ID { + // This could be a little faster ;) + for _, m := range guild.Members { + s.AddMember(&m) + } + for _, c := range guild.Channels { + s.AddChannel(&c) + } + return + } + } + s.Guilds = append(s.Guilds, *guild) +} + +// RemoveGuild removes a guild from current world state. +func (s *State) RemoveGuild(guild *Guild) error { + for i, g := range s.Guilds { + if g.ID == guild.ID { + s.Guilds = append(s.Guilds[:i], s.Guilds[i+1:]...) + return nil + } + } + return errors.New("Guild not found.") +} + +// GetGuildByID gets a guild by ID. +// Useful for querying if @me is in a guild: +// _, err := discordgo.Session.State.GetGuildById(guildID) +// isInGuild := err == nil +func (s *State) GetGuildByID(guildID string) (*Guild, error) { + for _, g := range s.Guilds { + if g.ID == guildID { + return &g, nil + } + } + return nil, errors.New("Guild not found.") +} + +// TODO: Consider moving Guild state update methods onto *Guild. + +// AddMember adds a member to the current world state, or +// updates it if it already exists. +func (s *State) AddMember(member *Member) error { + guild, err := s.GetGuildByID(member.GuildID) + if err != nil { + return err + } + + for i, m := range guild.Members { + if m.User.ID == member.User.ID { + guild.Members[i] = *member + return nil + } + } + + guild.Members = append(guild.Members, *member) + return nil +} + +// RemoveMember removes a member from current world state. +func (s *State) RemoveMember(member *Member) error { + guild, err := s.GetGuildByID(member.GuildID) + if err != nil { + return err + } + + for i, m := range guild.Members { + if m.User.ID == member.User.ID { + guild.Members = append(guild.Members[:i], guild.Members[i+1:]...) + return nil + } + } + return errors.New("Member not found.") +} + +// GetMemberByID gets a member by ID from a guild. +func (s *State) GetMemberByID(guildID string, userID string) (*Member, error) { + guild, err := s.GetGuildByID(guildID) + if err != nil { + return nil, err + } + + for _, m := range guild.Members { + if m.User.ID == userID { + return &m, nil + } + } + return nil, errors.New("Member not found.") +} + +// AddChannel adds a guild to the current world state, or +// updates it if it already exists. +// Channels may exist either as PrivateChannels or inside +// a guild. +func (s *State) AddChannel(channel *Channel) error { + if channel.IsPrivate { + for i, c := range s.PrivateChannels { + if c.ID == channel.ID { + s.PrivateChannels[i] = *channel + return nil + } + } + + s.PrivateChannels = append(s.PrivateChannels, *channel) + } else { + guild, err := s.GetGuildByID(channel.GuildID) + if err != nil { + return err + } + + for i, c := range guild.Channels { + if c.ID == channel.ID { + guild.Channels[i] = *channel + return nil + } + } + + guild.Channels = append(guild.Channels, *channel) + } + return nil +} + +// RemoveChannel removes a channel from current world state. +func (s *State) RemoveChannel(channel *Channel) error { + if channel.IsPrivate { + for i, c := range s.PrivateChannels { + if c.ID == channel.ID { + s.PrivateChannels = append(s.PrivateChannels[:i], s.PrivateChannels[i+1:]...) + return nil + } + } + } else { + guild, err := s.GetGuildByID(channel.GuildID) + if err != nil { + return err + } + + for i, c := range guild.Channels { + if c.ID == channel.ID { + guild.Channels = append(guild.Channels[:i], guild.Channels[i+1:]...) + return nil + } + } + } + + return errors.New("Channel not found.") +} + +// GetGuildChannelById gets a channel by ID from a guild. +func (s *State) GetGuildChannelByID(guildID string, channelID string) (*Channel, error) { + guild, err := s.GetGuildByID(guildID) + if err != nil { + return nil, err + } + + for _, c := range guild.Channels { + if c.ID == channelID { + return &c, nil + } + } + return nil, errors.New("Channel not found.") +} + +// GetPrivateChannelByID gets a private channel by ID. +func (s *State) GetPrivateChannelByID(channelID string) (*Channel, error) { + for _, c := range s.PrivateChannels { + if c.ID == channelID { + return &c, nil + } + } + return nil, errors.New("Channel not found.") +} diff --git a/structs.go b/structs.go index d51fcf2..4d78151 100644 --- a/structs.go +++ b/structs.go @@ -84,6 +84,9 @@ type Session struct { VChannelID string Vop2 VoiceOP2 UDPConn *net.UDPConn + + // Managed state object, updated with events. + State *State } // A Message stores all data related to a specific Discord message. @@ -262,14 +265,6 @@ type User struct { // it just doesn't seem able to handle this one // field correctly. Need to research this more. -// A PrivateChannel stores all data for a specific user private channel. -type PrivateChannel struct { - ID string `json:"id"` - IsPrivate bool `json:"is_private"` - LastMessageID string `json:"last_message_id"` - Recipient User `json:"recipient"` -} // merge with channel? - // A Settings stores data for a specific users Discord client settings. type Settings struct { RenderEmbeds bool `json:"render_embeds"` @@ -298,8 +293,8 @@ type Ready struct { HeartbeatInterval time.Duration `json:"heartbeat_interval"` User User `json:"user"` ReadState []ReadState - PrivateChannels []PrivateChannel - Guilds []Guild + PrivateChannels []Channel `json:"private_channels"` + Guilds []Guild `json:"guilds"` } // A ReadState stores data on the read state of channels. @@ -360,3 +355,10 @@ type GuildBan struct { User User `json:"user"` GuildID string `json:"guild_id"` } + +// A State contains the current known state, as discord sends this in a +// READY blob, it seems reasonable to simply use that message type as the +// data store. +type State struct { + Ready +} diff --git a/wsapi.go b/wsapi.go index 2070df4..3a7374b 100644 --- a/wsapi.go +++ b/wsapi.go @@ -153,6 +153,7 @@ func (s *Session) event(messageType int, message []byte) (err error) { case "READY": var st Ready if err = unmarshalEvent(e, &st); err == nil { + s.State.OnReady(&st) if s.OnReady != nil { s.OnReady(s, st) } @@ -235,77 +236,87 @@ func (s *Session) event(messageType int, message []byte) (err error) { return } case "CHANNEL_CREATE": - if s.OnChannelCreate != nil { - var st Channel - if err = unmarshalEvent(e, &st); err == nil { + var st Channel + if err = unmarshalEvent(e, &st); err == nil { + fmt.Println("channel create", st) + s.State.AddChannel(&st) + if s.OnChannelCreate != nil { s.OnChannelCreate(s, st) } - return } + return case "CHANNEL_UPDATE": - if s.OnChannelUpdate != nil { - var st Channel - if err = unmarshalEvent(e, &st); err == nil { + var st Channel + if err = unmarshalEvent(e, &st); err == nil { + s.State.AddChannel(&st) + if s.OnChannelUpdate != nil { s.OnChannelUpdate(s, st) } - return } + return case "CHANNEL_DELETE": - if s.OnChannelDelete != nil { - var st Channel - if err = unmarshalEvent(e, &st); err == nil { + var st Channel + if err = unmarshalEvent(e, &st); err == nil { + s.State.RemoveChannel(&st) + if s.OnChannelDelete != nil { s.OnChannelDelete(s, st) } - return } + return case "GUILD_CREATE": - if s.OnGuildCreate != nil { - var st Guild - if err = unmarshalEvent(e, &st); err == nil { + var st Guild + if err = unmarshalEvent(e, &st); err == nil { + s.State.AddGuild(&st) + if s.OnGuildCreate != nil { s.OnGuildCreate(s, st) } - return } + return case "GUILD_UPDATE": - if s.OnGuildCreate != nil { - var st Guild - if err = unmarshalEvent(e, &st); err == nil { + var st Guild + if err = unmarshalEvent(e, &st); err == nil { + s.State.AddGuild(&st) + if s.OnGuildCreate != nil { s.OnGuildUpdate(s, st) } - return } + return case "GUILD_DELETE": - if s.OnGuildDelete != nil { - var st Guild - if err = unmarshalEvent(e, &st); err == nil { + var st Guild + if err = unmarshalEvent(e, &st); err == nil { + s.State.RemoveGuild(&st) + if s.OnGuildDelete != nil { s.OnGuildDelete(s, st) } - return } + return case "GUILD_MEMBER_ADD": - if s.OnGuildMemberAdd != nil { - var st Member - if err = unmarshalEvent(e, &st); err == nil { + var st Member + if err = unmarshalEvent(e, &st); err == nil { + s.State.AddMember(&st) + if s.OnGuildMemberAdd != nil { s.OnGuildMemberAdd(s, st) } - return } + return case "GUILD_MEMBER_REMOVE": - if s.OnGuildMemberRemove != nil { - var st Member - if err = unmarshalEvent(e, &st); err == nil { + var st Member + if err = unmarshalEvent(e, &st); err == nil { + s.State.RemoveMember(&st) + if s.OnGuildMemberRemove != nil { s.OnGuildMemberRemove(s, st) } - return } + return case "GUILD_MEMBER_UPDATE": - if s.OnGuildMemberUpdate != nil { - var st Member - if err = unmarshalEvent(e, &st); err == nil { + var st Member + if err = unmarshalEvent(e, &st); err == nil { + s.State.AddMember(&st) + if s.OnGuildMemberUpdate != nil { s.OnGuildMemberUpdate(s, st) } - return } + return case "GUILD_ROLE_CREATE": if s.OnGuildRoleCreate != nil { var st GuildRole