state.go 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759
  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, merge in the new message contents.
  396. for _, m := range c.Messages {
  397. if m.ID == message.ID {
  398. if message.Content != "" {
  399. m.Content = message.Content
  400. }
  401. if message.EditedTimestamp != "" {
  402. m.EditedTimestamp = message.EditedTimestamp
  403. }
  404. if message.Mentions != nil {
  405. m.Mentions = message.Mentions
  406. }
  407. if message.Embeds != nil {
  408. m.Embeds = message.Embeds
  409. }
  410. if message.Attachments != nil {
  411. m.Attachments = message.Attachments
  412. }
  413. if message.Timestamp != "" {
  414. m.Timestamp = message.Timestamp
  415. }
  416. if message.Author != nil {
  417. m.Author = message.Author
  418. }
  419. return nil
  420. }
  421. }
  422. c.Messages = append(c.Messages, message)
  423. if len(c.Messages) > s.MaxMessageCount {
  424. c.Messages = c.Messages[len(c.Messages)-s.MaxMessageCount:]
  425. }
  426. return nil
  427. }
  428. // MessageRemove removes a message from the world state.
  429. func (s *State) MessageRemove(message *Message) error {
  430. if s == nil {
  431. return ErrNilState
  432. }
  433. c, err := s.Channel(message.ChannelID)
  434. if err != nil {
  435. return err
  436. }
  437. s.Lock()
  438. defer s.Unlock()
  439. for i, m := range c.Messages {
  440. if m.ID == message.ID {
  441. c.Messages = append(c.Messages[:i], c.Messages[i+1:]...)
  442. return nil
  443. }
  444. }
  445. return errors.New("Message not found.")
  446. }
  447. func (s *State) voiceStateUpdate(update *VoiceStateUpdate) error {
  448. guild, err := s.Guild(update.GuildID)
  449. if err != nil {
  450. return err
  451. }
  452. s.Lock()
  453. defer s.Unlock()
  454. // Handle Leaving Channel
  455. if update.ChannelID == "" {
  456. for i, state := range guild.VoiceStates {
  457. if state.UserID == update.UserID {
  458. guild.VoiceStates = append(guild.VoiceStates[:i], guild.VoiceStates[i+1:]...)
  459. return nil
  460. }
  461. }
  462. } else {
  463. for i, state := range guild.VoiceStates {
  464. if state.UserID == update.UserID {
  465. guild.VoiceStates[i] = update.VoiceState
  466. return nil
  467. }
  468. }
  469. guild.VoiceStates = append(guild.VoiceStates, update.VoiceState)
  470. }
  471. return nil
  472. }
  473. // Message gets a message by channel and message ID.
  474. func (s *State) Message(channelID, messageID string) (*Message, error) {
  475. if s == nil {
  476. return nil, ErrNilState
  477. }
  478. c, err := s.Channel(channelID)
  479. if err != nil {
  480. return nil, err
  481. }
  482. s.RLock()
  483. defer s.RUnlock()
  484. for _, m := range c.Messages {
  485. if m.ID == messageID {
  486. return m, nil
  487. }
  488. }
  489. return nil, errors.New("Message not found.")
  490. }
  491. // onInterface handles all events related to states.
  492. func (s *State) onInterface(se *Session, i interface{}) (err error) {
  493. if s == nil {
  494. return ErrNilState
  495. }
  496. if !se.StateEnabled {
  497. return nil
  498. }
  499. switch t := i.(type) {
  500. case *Ready:
  501. err = s.OnReady(t)
  502. case *GuildCreate:
  503. err = s.GuildAdd(t.Guild)
  504. case *GuildUpdate:
  505. err = s.GuildAdd(t.Guild)
  506. case *GuildDelete:
  507. err = s.GuildRemove(t.Guild)
  508. case *GuildMemberAdd:
  509. if s.TrackMembers {
  510. err = s.MemberAdd(t.Member)
  511. }
  512. case *GuildMemberUpdate:
  513. if s.TrackMembers {
  514. err = s.MemberAdd(t.Member)
  515. }
  516. case *GuildMemberRemove:
  517. if s.TrackMembers {
  518. err = s.MemberRemove(t.Member)
  519. }
  520. case *GuildRoleCreate:
  521. if s.TrackRoles {
  522. err = s.RoleAdd(t.GuildID, t.Role)
  523. }
  524. case *GuildRoleUpdate:
  525. if s.TrackRoles {
  526. err = s.RoleAdd(t.GuildID, t.Role)
  527. }
  528. case *GuildRoleDelete:
  529. if s.TrackRoles {
  530. err = s.RoleRemove(t.GuildID, t.RoleID)
  531. }
  532. case *GuildEmojisUpdate:
  533. if s.TrackEmojis {
  534. err = s.EmojisAdd(t.GuildID, t.Emojis)
  535. }
  536. case *ChannelCreate:
  537. if s.TrackChannels {
  538. err = s.ChannelAdd(t.Channel)
  539. }
  540. case *ChannelUpdate:
  541. if s.TrackChannels {
  542. err = s.ChannelAdd(t.Channel)
  543. }
  544. case *ChannelDelete:
  545. if s.TrackChannels {
  546. err = s.ChannelRemove(t.Channel)
  547. }
  548. case *MessageCreate:
  549. if s.MaxMessageCount != 0 {
  550. err = s.MessageAdd(t.Message)
  551. }
  552. case *MessageUpdate:
  553. if s.MaxMessageCount != 0 {
  554. err = s.MessageAdd(t.Message)
  555. }
  556. case *MessageDelete:
  557. if s.MaxMessageCount != 0 {
  558. err = s.MessageRemove(t.Message)
  559. }
  560. case *VoiceStateUpdate:
  561. if s.TrackVoice {
  562. err = s.voiceStateUpdate(t)
  563. }
  564. }
  565. return
  566. }
  567. // UserChannelPermissions returns the permission of a user in a channel.
  568. // userID : The ID of the user to calculate permissions for.
  569. // channelID : The ID of the channel to calculate permission for.
  570. func (s *State) UserChannelPermissions(userID, channelID string) (apermissions int, err error) {
  571. channel, err := s.Channel(channelID)
  572. if err != nil {
  573. return
  574. }
  575. guild, err := s.Guild(channel.GuildID)
  576. if err != nil {
  577. return
  578. }
  579. if userID == guild.OwnerID {
  580. apermissions = PermissionAll
  581. return
  582. }
  583. member, err := s.Member(guild.ID, userID)
  584. if err != nil {
  585. return
  586. }
  587. for _, role := range guild.Roles {
  588. if role.ID == guild.ID {
  589. apermissions |= role.Permissions
  590. break
  591. }
  592. }
  593. for _, role := range guild.Roles {
  594. for _, roleID := range member.Roles {
  595. if role.ID == roleID {
  596. apermissions |= role.Permissions
  597. break
  598. }
  599. }
  600. }
  601. if apermissions&PermissionManageRoles > 0 {
  602. apermissions |= PermissionAll
  603. }
  604. // Member overwrites can override role overrides, so do two passes
  605. for _, overwrite := range channel.PermissionOverwrites {
  606. for _, roleID := range member.Roles {
  607. if overwrite.Type == "role" && roleID == overwrite.ID {
  608. apermissions &= ^overwrite.Deny
  609. apermissions |= overwrite.Allow
  610. break
  611. }
  612. }
  613. }
  614. for _, overwrite := range channel.PermissionOverwrites {
  615. if overwrite.Type == "member" && overwrite.ID == userID {
  616. apermissions &= ^overwrite.Deny
  617. apermissions |= overwrite.Allow
  618. break
  619. }
  620. }
  621. if apermissions&PermissionManageRoles > 0 {
  622. apermissions |= PermissionAllChannel
  623. }
  624. return
  625. }