state.go 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667
  1. // Discordgo - Discord bindings for Go
  2. // Available at https://github.com/bwmarrin/discordgo
  3. // Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>. All rights reserved.
  4. // Use of this source code is governed by a BSD-style
  5. // license that can be found in the LICENSE file.
  6. // This file contains code related to state tracking. If enabled, state
  7. // tracking will capture the initial READY packet and many other websocket
  8. // events and maintain an in-memory state of of guilds, channels, users, and
  9. // so forth. This information can be accessed through the Session.State struct.
  10. package discordgo
  11. import (
  12. "errors"
  13. "sync"
  14. )
  15. // ErrNilState is returned when the state is nil.
  16. var ErrNilState = errors.New("State not instantiated, please use discordgo.New() or assign Session.State.")
  17. // A State contains the current known state.
  18. // As discord sends this in a READY blob, it seems reasonable to simply
  19. // use that struct as the data store.
  20. type State struct {
  21. sync.RWMutex
  22. Ready
  23. MaxMessageCount int
  24. TrackChannels bool
  25. TrackEmojis bool
  26. TrackMembers bool
  27. TrackRoles bool
  28. TrackVoice bool
  29. guildMap map[string]*Guild
  30. channelMap map[string]*Channel
  31. }
  32. // NewState creates an empty state.
  33. func NewState() *State {
  34. return &State{
  35. Ready: Ready{
  36. PrivateChannels: []*Channel{},
  37. Guilds: []*Guild{},
  38. },
  39. TrackChannels: true,
  40. TrackEmojis: true,
  41. TrackMembers: true,
  42. TrackRoles: true,
  43. TrackVoice: true,
  44. guildMap: make(map[string]*Guild),
  45. channelMap: make(map[string]*Channel),
  46. }
  47. }
  48. // OnReady takes a Ready event and updates all internal state.
  49. func (s *State) OnReady(r *Ready) error {
  50. if s == nil {
  51. return ErrNilState
  52. }
  53. s.Lock()
  54. defer s.Unlock()
  55. s.Ready = *r
  56. for _, g := range s.Guilds {
  57. s.guildMap[g.ID] = g
  58. for _, c := range g.Channels {
  59. c.GuildID = g.ID
  60. s.channelMap[c.ID] = c
  61. }
  62. }
  63. for _, c := range s.PrivateChannels {
  64. s.channelMap[c.ID] = c
  65. }
  66. return nil
  67. }
  68. // GuildAdd adds a guild to the current world state, or
  69. // updates it if it already exists.
  70. func (s *State) GuildAdd(guild *Guild) error {
  71. if s == nil {
  72. return ErrNilState
  73. }
  74. s.Lock()
  75. defer s.Unlock()
  76. // Update the channels to point to the right guild, adding them to the channelMap as we go
  77. for _, c := range guild.Channels {
  78. c.GuildID = guild.ID
  79. s.channelMap[c.ID] = c
  80. }
  81. // If the guild exists, replace it.
  82. if g, ok := s.guildMap[guild.ID]; ok {
  83. // If this guild already exists with data, don't stomp on props.
  84. if g.Unavailable != nil && !*g.Unavailable {
  85. guild.Members = g.Members
  86. guild.Presences = g.Presences
  87. guild.Channels = g.Channels
  88. guild.VoiceStates = g.VoiceStates
  89. }
  90. *g = *guild
  91. return nil
  92. }
  93. s.Guilds = append(s.Guilds, guild)
  94. s.guildMap[guild.ID] = guild
  95. return nil
  96. }
  97. // GuildRemove removes a guild from current world state.
  98. func (s *State) GuildRemove(guild *Guild) error {
  99. if s == nil {
  100. return ErrNilState
  101. }
  102. _, err := s.Guild(guild.ID)
  103. if err != nil {
  104. return err
  105. }
  106. s.Lock()
  107. defer s.Unlock()
  108. delete(s.guildMap, guild.ID)
  109. for i, g := range s.Guilds {
  110. if g.ID == guild.ID {
  111. s.Guilds = append(s.Guilds[:i], s.Guilds[i+1:]...)
  112. return nil
  113. }
  114. }
  115. return nil
  116. }
  117. // Guild gets a guild by ID.
  118. // Useful for querying if @me is in a guild:
  119. // _, err := discordgo.Session.State.Guild(guildID)
  120. // isInGuild := err == nil
  121. func (s *State) Guild(guildID string) (*Guild, error) {
  122. if s == nil {
  123. return nil, ErrNilState
  124. }
  125. s.RLock()
  126. defer s.RUnlock()
  127. if g, ok := s.guildMap[guildID]; ok {
  128. return g, nil
  129. }
  130. return nil, errors.New("Guild not found.")
  131. }
  132. // TODO: Consider moving Guild state update methods onto *Guild.
  133. // MemberAdd adds a member to the current world state, or
  134. // updates it if it already exists.
  135. func (s *State) MemberAdd(member *Member) error {
  136. if s == nil {
  137. return ErrNilState
  138. }
  139. guild, err := s.Guild(member.GuildID)
  140. if err != nil {
  141. return err
  142. }
  143. s.Lock()
  144. defer s.Unlock()
  145. for i, m := range guild.Members {
  146. if m.User.ID == member.User.ID {
  147. guild.Members[i] = member
  148. return nil
  149. }
  150. }
  151. guild.Members = append(guild.Members, member)
  152. return nil
  153. }
  154. // MemberRemove removes a member from current world state.
  155. func (s *State) MemberRemove(member *Member) error {
  156. if s == nil {
  157. return ErrNilState
  158. }
  159. guild, err := s.Guild(member.GuildID)
  160. if err != nil {
  161. return err
  162. }
  163. s.Lock()
  164. defer s.Unlock()
  165. for i, m := range guild.Members {
  166. if m.User.ID == member.User.ID {
  167. guild.Members = append(guild.Members[:i], guild.Members[i+1:]...)
  168. return nil
  169. }
  170. }
  171. return errors.New("Member not found.")
  172. }
  173. // Member gets a member by ID from a guild.
  174. func (s *State) Member(guildID, userID string) (*Member, error) {
  175. if s == nil {
  176. return nil, ErrNilState
  177. }
  178. guild, err := s.Guild(guildID)
  179. if err != nil {
  180. return nil, err
  181. }
  182. s.RLock()
  183. defer s.RUnlock()
  184. for _, m := range guild.Members {
  185. if m.User.ID == userID {
  186. return m, nil
  187. }
  188. }
  189. return nil, errors.New("Member not found.")
  190. }
  191. // RoleAdd adds a role to the current world state, or
  192. // updates it if it already exists.
  193. func (s *State) RoleAdd(guildID string, role *Role) error {
  194. if s == nil {
  195. return ErrNilState
  196. }
  197. guild, err := s.Guild(guildID)
  198. if err != nil {
  199. return err
  200. }
  201. s.Lock()
  202. defer s.Unlock()
  203. for i, r := range guild.Roles {
  204. if r.ID == role.ID {
  205. guild.Roles[i] = role
  206. return nil
  207. }
  208. }
  209. guild.Roles = append(guild.Roles, role)
  210. return nil
  211. }
  212. // RoleRemove removes a role from current world state by ID.
  213. func (s *State) RoleRemove(guildID, roleID string) error {
  214. if s == nil {
  215. return ErrNilState
  216. }
  217. guild, err := s.Guild(guildID)
  218. if err != nil {
  219. return err
  220. }
  221. s.Lock()
  222. defer s.Unlock()
  223. for i, r := range guild.Roles {
  224. if r.ID == roleID {
  225. guild.Roles = append(guild.Roles[:i], guild.Roles[i+1:]...)
  226. return nil
  227. }
  228. }
  229. return errors.New("Role not found.")
  230. }
  231. // Role gets a role by ID from a guild.
  232. func (s *State) Role(guildID, roleID string) (*Role, error) {
  233. if s == nil {
  234. return nil, ErrNilState
  235. }
  236. guild, err := s.Guild(guildID)
  237. if err != nil {
  238. return nil, err
  239. }
  240. s.RLock()
  241. defer s.RUnlock()
  242. for _, r := range guild.Roles {
  243. if r.ID == roleID {
  244. return r, nil
  245. }
  246. }
  247. return nil, errors.New("Role not found.")
  248. }
  249. // ChannelAdd adds a guild to the current world state, or
  250. // updates it if it already exists.
  251. // Channels may exist either as PrivateChannels or inside
  252. // a guild.
  253. func (s *State) ChannelAdd(channel *Channel) error {
  254. if s == nil {
  255. return ErrNilState
  256. }
  257. s.Lock()
  258. defer s.Unlock()
  259. // If the channel exists, replace it
  260. if c, ok := s.channelMap[channel.ID]; ok {
  261. channel.Messages = c.Messages
  262. channel.PermissionOverwrites = c.PermissionOverwrites
  263. *c = *channel
  264. return nil
  265. }
  266. if channel.IsPrivate {
  267. s.PrivateChannels = append(s.PrivateChannels, channel)
  268. } else {
  269. guild, ok := s.guildMap[channel.GuildID]
  270. if !ok {
  271. return errors.New("Guild for channel not found.")
  272. }
  273. guild.Channels = append(guild.Channels, channel)
  274. }
  275. s.channelMap[channel.ID] = channel
  276. return nil
  277. }
  278. // ChannelRemove removes a channel from current world state.
  279. func (s *State) ChannelRemove(channel *Channel) error {
  280. if s == nil {
  281. return ErrNilState
  282. }
  283. _, err := s.Channel(channel.ID)
  284. if err != nil {
  285. return err
  286. }
  287. if channel.IsPrivate {
  288. s.Lock()
  289. defer s.Unlock()
  290. for i, c := range s.PrivateChannels {
  291. if c.ID == channel.ID {
  292. s.PrivateChannels = append(s.PrivateChannels[:i], s.PrivateChannels[i+1:]...)
  293. break
  294. }
  295. }
  296. } else {
  297. guild, err := s.Guild(channel.GuildID)
  298. if err != nil {
  299. return err
  300. }
  301. s.Lock()
  302. defer s.Unlock()
  303. for i, c := range guild.Channels {
  304. if c.ID == channel.ID {
  305. guild.Channels = append(guild.Channels[:i], guild.Channels[i+1:]...)
  306. break
  307. }
  308. }
  309. }
  310. delete(s.channelMap, channel.ID)
  311. return nil
  312. }
  313. // GuildChannel gets a channel by ID from a guild.
  314. // This method is Deprecated, use Channel(channelID)
  315. func (s *State) GuildChannel(guildID, channelID string) (*Channel, error) {
  316. return s.Channel(channelID)
  317. }
  318. // PrivateChannel gets a private channel by ID.
  319. // This method is Deprecated, use Channel(channelID)
  320. func (s *State) PrivateChannel(channelID string) (*Channel, error) {
  321. return s.Channel(channelID)
  322. }
  323. // Channel gets a channel by ID, it will look in all guilds an private channels.
  324. func (s *State) Channel(channelID string) (*Channel, error) {
  325. if s == nil {
  326. return nil, ErrNilState
  327. }
  328. s.RLock()
  329. defer s.RUnlock()
  330. if c, ok := s.channelMap[channelID]; ok {
  331. return c, nil
  332. }
  333. return nil, errors.New("Channel not found.")
  334. }
  335. // Emoji returns an emoji for a guild and emoji id.
  336. func (s *State) Emoji(guildID, emojiID string) (*Emoji, error) {
  337. if s == nil {
  338. return nil, ErrNilState
  339. }
  340. guild, err := s.Guild(guildID)
  341. if err != nil {
  342. return nil, err
  343. }
  344. s.RLock()
  345. defer s.RUnlock()
  346. for _, e := range guild.Emojis {
  347. if e.ID == emojiID {
  348. return e, nil
  349. }
  350. }
  351. return nil, errors.New("Emoji not found.")
  352. }
  353. // EmojiAdd adds an emoji to the current world state.
  354. func (s *State) EmojiAdd(guildID string, emoji *Emoji) error {
  355. if s == nil {
  356. return ErrNilState
  357. }
  358. guild, err := s.Guild(guildID)
  359. if err != nil {
  360. return err
  361. }
  362. s.Lock()
  363. defer s.Unlock()
  364. for i, e := range guild.Emojis {
  365. if e.ID == emoji.ID {
  366. guild.Emojis[i] = emoji
  367. return nil
  368. }
  369. }
  370. guild.Emojis = append(guild.Emojis, emoji)
  371. return nil
  372. }
  373. // EmojisAdd adds multiple emojis to the world state.
  374. func (s *State) EmojisAdd(guildID string, emojis []*Emoji) error {
  375. for _, e := range emojis {
  376. if err := s.EmojiAdd(guildID, e); err != nil {
  377. return err
  378. }
  379. }
  380. return nil
  381. }
  382. // MessageAdd adds a message to the current world state, or updates it if it exists.
  383. // If the channel cannot be found, the message is discarded.
  384. // Messages are kept in state up to s.MaxMessageCount
  385. func (s *State) MessageAdd(message *Message) error {
  386. if s == nil {
  387. return ErrNilState
  388. }
  389. c, err := s.Channel(message.ChannelID)
  390. if err != nil {
  391. return err
  392. }
  393. s.Lock()
  394. defer s.Unlock()
  395. // If the message exists, replace it.
  396. for i, m := range c.Messages {
  397. if m.ID == message.ID {
  398. c.Messages[i] = message
  399. return nil
  400. }
  401. }
  402. c.Messages = append(c.Messages, message)
  403. if len(c.Messages) > s.MaxMessageCount {
  404. c.Messages = c.Messages[len(c.Messages)-s.MaxMessageCount:]
  405. }
  406. return nil
  407. }
  408. // MessageRemove removes a message from the world state.
  409. func (s *State) MessageRemove(message *Message) error {
  410. if s == nil {
  411. return ErrNilState
  412. }
  413. c, err := s.Channel(message.ChannelID)
  414. if err != nil {
  415. return err
  416. }
  417. s.Lock()
  418. defer s.Unlock()
  419. for i, m := range c.Messages {
  420. if m.ID == message.ID {
  421. c.Messages = append(c.Messages[:i], c.Messages[i+1:]...)
  422. return nil
  423. }
  424. }
  425. return errors.New("Message not found.")
  426. }
  427. func (s *State) voiceStateUpdate(update *VoiceStateUpdate) error {
  428. guild, err := s.Guild(update.GuildID)
  429. if err != nil {
  430. return err
  431. }
  432. s.Lock()
  433. defer s.Unlock()
  434. // Handle Leaving Channel
  435. if update.ChannelID == "" {
  436. for i, state := range guild.VoiceStates {
  437. if state.UserID == update.UserID {
  438. guild.VoiceStates = append(guild.VoiceStates[:i], guild.VoiceStates[i+1:]...)
  439. return nil
  440. }
  441. }
  442. } else {
  443. for i, state := range guild.VoiceStates {
  444. if state.UserID == update.UserID {
  445. guild.VoiceStates[i] = update.VoiceState
  446. return nil
  447. }
  448. }
  449. guild.VoiceStates = append(guild.VoiceStates, update.VoiceState)
  450. }
  451. return nil
  452. }
  453. // Message gets a message by channel and message ID.
  454. func (s *State) Message(channelID, messageID string) (*Message, error) {
  455. if s == nil {
  456. return nil, ErrNilState
  457. }
  458. c, err := s.Channel(channelID)
  459. if err != nil {
  460. return nil, err
  461. }
  462. s.RLock()
  463. defer s.RUnlock()
  464. for _, m := range c.Messages {
  465. if m.ID == messageID {
  466. return m, nil
  467. }
  468. }
  469. return nil, errors.New("Message not found.")
  470. }
  471. // onInterface handles all events related to states.
  472. func (s *State) onInterface(se *Session, i interface{}) (err error) {
  473. if s == nil {
  474. return ErrNilState
  475. }
  476. if !se.StateEnabled {
  477. return nil
  478. }
  479. switch t := i.(type) {
  480. case *Ready:
  481. err = s.OnReady(t)
  482. case *GuildCreate:
  483. err = s.GuildAdd(t.Guild)
  484. case *GuildUpdate:
  485. err = s.GuildAdd(t.Guild)
  486. case *GuildDelete:
  487. err = s.GuildRemove(t.Guild)
  488. case *GuildMemberAdd:
  489. if s.TrackMembers {
  490. err = s.MemberAdd(t.Member)
  491. }
  492. case *GuildMemberUpdate:
  493. if s.TrackMembers {
  494. err = s.MemberAdd(t.Member)
  495. }
  496. case *GuildMemberRemove:
  497. if s.TrackMembers {
  498. err = s.MemberRemove(t.Member)
  499. }
  500. case *GuildRoleCreate:
  501. if s.TrackRoles {
  502. err = s.RoleAdd(t.GuildID, t.Role)
  503. }
  504. case *GuildRoleUpdate:
  505. if s.TrackRoles {
  506. err = s.RoleAdd(t.GuildID, t.Role)
  507. }
  508. case *GuildRoleDelete:
  509. if s.TrackRoles {
  510. err = s.RoleRemove(t.GuildID, t.RoleID)
  511. }
  512. case *GuildEmojisUpdate:
  513. if s.TrackEmojis {
  514. err = s.EmojisAdd(t.GuildID, t.Emojis)
  515. }
  516. case *ChannelCreate:
  517. if s.TrackChannels {
  518. err = s.ChannelAdd(t.Channel)
  519. }
  520. case *ChannelUpdate:
  521. if s.TrackChannels {
  522. err = s.ChannelAdd(t.Channel)
  523. }
  524. case *ChannelDelete:
  525. if s.TrackChannels {
  526. err = s.ChannelRemove(t.Channel)
  527. }
  528. case *MessageCreate:
  529. if s.MaxMessageCount != 0 {
  530. err = s.MessageAdd(t.Message)
  531. }
  532. case *MessageUpdate:
  533. if s.MaxMessageCount != 0 {
  534. err = s.MessageAdd(t.Message)
  535. }
  536. case *MessageDelete:
  537. if s.MaxMessageCount != 0 {
  538. err = s.MessageRemove(t.Message)
  539. }
  540. case *VoiceStateUpdate:
  541. if s.TrackVoice {
  542. err = s.voiceStateUpdate(t)
  543. }
  544. }
  545. return
  546. }