state.go 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790
  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. // GuildAdd adds a guild to the current world state, or
  49. // updates it if it already exists.
  50. func (s *State) GuildAdd(guild *Guild) error {
  51. if s == nil {
  52. return ErrNilState
  53. }
  54. s.Lock()
  55. defer s.Unlock()
  56. // Update the channels to point to the right guild, adding them to the channelMap as we go
  57. for _, c := range guild.Channels {
  58. s.channelMap[c.ID] = c
  59. }
  60. if g, ok := s.guildMap[guild.ID]; ok {
  61. // We are about to replace `g` in the state with `guild`, but first we need to
  62. // make sure we preserve any fields that the `guild` doesn't contain from `g`.
  63. if guild.Roles == nil {
  64. guild.Roles = g.Roles
  65. }
  66. if guild.Emojis == nil {
  67. guild.Emojis = g.Emojis
  68. }
  69. if guild.Members == nil {
  70. guild.Members = g.Members
  71. }
  72. if guild.Presences == nil {
  73. guild.Presences = g.Presences
  74. }
  75. if guild.Channels == nil {
  76. guild.Channels = g.Channels
  77. }
  78. if guild.VoiceStates == nil {
  79. guild.VoiceStates = g.VoiceStates
  80. }
  81. *g = *guild
  82. return nil
  83. }
  84. s.Guilds = append(s.Guilds, guild)
  85. s.guildMap[guild.ID] = guild
  86. return nil
  87. }
  88. // GuildRemove removes a guild from current world state.
  89. func (s *State) GuildRemove(guild *Guild) error {
  90. if s == nil {
  91. return ErrNilState
  92. }
  93. _, err := s.Guild(guild.ID)
  94. if err != nil {
  95. return err
  96. }
  97. s.Lock()
  98. defer s.Unlock()
  99. delete(s.guildMap, guild.ID)
  100. for i, g := range s.Guilds {
  101. if g.ID == guild.ID {
  102. s.Guilds = append(s.Guilds[:i], s.Guilds[i+1:]...)
  103. return nil
  104. }
  105. }
  106. return nil
  107. }
  108. // Guild gets a guild by ID.
  109. // Useful for querying if @me is in a guild:
  110. // _, err := discordgo.Session.State.Guild(guildID)
  111. // isInGuild := err == nil
  112. func (s *State) Guild(guildID string) (*Guild, error) {
  113. if s == nil {
  114. return nil, ErrNilState
  115. }
  116. s.RLock()
  117. defer s.RUnlock()
  118. if g, ok := s.guildMap[guildID]; ok {
  119. return g, nil
  120. }
  121. return nil, errors.New("Guild not found.")
  122. }
  123. // TODO: Consider moving Guild state update methods onto *Guild.
  124. // MemberAdd adds a member to the current world state, or
  125. // updates it if it already exists.
  126. func (s *State) MemberAdd(member *Member) error {
  127. if s == nil {
  128. return ErrNilState
  129. }
  130. guild, err := s.Guild(member.GuildID)
  131. if err != nil {
  132. return err
  133. }
  134. s.Lock()
  135. defer s.Unlock()
  136. for i, m := range guild.Members {
  137. if m.User.ID == member.User.ID {
  138. guild.Members[i] = member
  139. return nil
  140. }
  141. }
  142. guild.Members = append(guild.Members, member)
  143. return nil
  144. }
  145. // MemberRemove removes a member from current world state.
  146. func (s *State) MemberRemove(member *Member) error {
  147. if s == nil {
  148. return ErrNilState
  149. }
  150. guild, err := s.Guild(member.GuildID)
  151. if err != nil {
  152. return err
  153. }
  154. s.Lock()
  155. defer s.Unlock()
  156. for i, m := range guild.Members {
  157. if m.User.ID == member.User.ID {
  158. guild.Members = append(guild.Members[:i], guild.Members[i+1:]...)
  159. return nil
  160. }
  161. }
  162. return errors.New("Member not found.")
  163. }
  164. // Member gets a member by ID from a guild.
  165. func (s *State) Member(guildID, userID string) (*Member, error) {
  166. if s == nil {
  167. return nil, ErrNilState
  168. }
  169. guild, err := s.Guild(guildID)
  170. if err != nil {
  171. return nil, err
  172. }
  173. s.RLock()
  174. defer s.RUnlock()
  175. for _, m := range guild.Members {
  176. if m.User.ID == userID {
  177. return m, nil
  178. }
  179. }
  180. return nil, errors.New("Member not found.")
  181. }
  182. // RoleAdd adds a role to the current world state, or
  183. // updates it if it already exists.
  184. func (s *State) RoleAdd(guildID string, role *Role) error {
  185. if s == nil {
  186. return ErrNilState
  187. }
  188. guild, err := s.Guild(guildID)
  189. if err != nil {
  190. return err
  191. }
  192. s.Lock()
  193. defer s.Unlock()
  194. for i, r := range guild.Roles {
  195. if r.ID == role.ID {
  196. guild.Roles[i] = role
  197. return nil
  198. }
  199. }
  200. guild.Roles = append(guild.Roles, role)
  201. return nil
  202. }
  203. // RoleRemove removes a role from current world state by ID.
  204. func (s *State) RoleRemove(guildID, roleID string) error {
  205. if s == nil {
  206. return ErrNilState
  207. }
  208. guild, err := s.Guild(guildID)
  209. if err != nil {
  210. return err
  211. }
  212. s.Lock()
  213. defer s.Unlock()
  214. for i, r := range guild.Roles {
  215. if r.ID == roleID {
  216. guild.Roles = append(guild.Roles[:i], guild.Roles[i+1:]...)
  217. return nil
  218. }
  219. }
  220. return errors.New("Role not found.")
  221. }
  222. // Role gets a role by ID from a guild.
  223. func (s *State) Role(guildID, roleID string) (*Role, error) {
  224. if s == nil {
  225. return nil, ErrNilState
  226. }
  227. guild, err := s.Guild(guildID)
  228. if err != nil {
  229. return nil, err
  230. }
  231. s.RLock()
  232. defer s.RUnlock()
  233. for _, r := range guild.Roles {
  234. if r.ID == roleID {
  235. return r, nil
  236. }
  237. }
  238. return nil, errors.New("Role not found.")
  239. }
  240. // ChannelAdd adds a guild to the current world state, or
  241. // updates it if it already exists.
  242. // Channels may exist either as PrivateChannels or inside
  243. // a guild.
  244. func (s *State) ChannelAdd(channel *Channel) error {
  245. if s == nil {
  246. return ErrNilState
  247. }
  248. s.Lock()
  249. defer s.Unlock()
  250. // If the channel exists, replace it
  251. if c, ok := s.channelMap[channel.ID]; ok {
  252. channel.Messages = c.Messages
  253. channel.PermissionOverwrites = c.PermissionOverwrites
  254. *c = *channel
  255. return nil
  256. }
  257. if channel.IsPrivate {
  258. s.PrivateChannels = append(s.PrivateChannels, channel)
  259. } else {
  260. guild, ok := s.guildMap[channel.GuildID]
  261. if !ok {
  262. return errors.New("Guild for channel not found.")
  263. }
  264. guild.Channels = append(guild.Channels, channel)
  265. }
  266. s.channelMap[channel.ID] = channel
  267. return nil
  268. }
  269. // ChannelRemove removes a channel from current world state.
  270. func (s *State) ChannelRemove(channel *Channel) error {
  271. if s == nil {
  272. return ErrNilState
  273. }
  274. _, err := s.Channel(channel.ID)
  275. if err != nil {
  276. return err
  277. }
  278. if channel.IsPrivate {
  279. s.Lock()
  280. defer s.Unlock()
  281. for i, c := range s.PrivateChannels {
  282. if c.ID == channel.ID {
  283. s.PrivateChannels = append(s.PrivateChannels[:i], s.PrivateChannels[i+1:]...)
  284. break
  285. }
  286. }
  287. } else {
  288. guild, err := s.Guild(channel.GuildID)
  289. if err != nil {
  290. return err
  291. }
  292. s.Lock()
  293. defer s.Unlock()
  294. for i, c := range guild.Channels {
  295. if c.ID == channel.ID {
  296. guild.Channels = append(guild.Channels[:i], guild.Channels[i+1:]...)
  297. break
  298. }
  299. }
  300. }
  301. delete(s.channelMap, channel.ID)
  302. return nil
  303. }
  304. // GuildChannel gets a channel by ID from a guild.
  305. // This method is Deprecated, use Channel(channelID)
  306. func (s *State) GuildChannel(guildID, channelID string) (*Channel, error) {
  307. return s.Channel(channelID)
  308. }
  309. // PrivateChannel gets a private channel by ID.
  310. // This method is Deprecated, use Channel(channelID)
  311. func (s *State) PrivateChannel(channelID string) (*Channel, error) {
  312. return s.Channel(channelID)
  313. }
  314. // Channel gets a channel by ID, it will look in all guilds an private channels.
  315. func (s *State) Channel(channelID string) (*Channel, error) {
  316. if s == nil {
  317. return nil, ErrNilState
  318. }
  319. s.RLock()
  320. defer s.RUnlock()
  321. if c, ok := s.channelMap[channelID]; ok {
  322. return c, nil
  323. }
  324. return nil, errors.New("Channel not found.")
  325. }
  326. // Emoji returns an emoji for a guild and emoji id.
  327. func (s *State) Emoji(guildID, emojiID string) (*Emoji, error) {
  328. if s == nil {
  329. return nil, ErrNilState
  330. }
  331. guild, err := s.Guild(guildID)
  332. if err != nil {
  333. return nil, err
  334. }
  335. s.RLock()
  336. defer s.RUnlock()
  337. for _, e := range guild.Emojis {
  338. if e.ID == emojiID {
  339. return e, nil
  340. }
  341. }
  342. return nil, errors.New("Emoji not found.")
  343. }
  344. // EmojiAdd adds an emoji to the current world state.
  345. func (s *State) EmojiAdd(guildID string, emoji *Emoji) error {
  346. if s == nil {
  347. return ErrNilState
  348. }
  349. guild, err := s.Guild(guildID)
  350. if err != nil {
  351. return err
  352. }
  353. s.Lock()
  354. defer s.Unlock()
  355. for i, e := range guild.Emojis {
  356. if e.ID == emoji.ID {
  357. guild.Emojis[i] = emoji
  358. return nil
  359. }
  360. }
  361. guild.Emojis = append(guild.Emojis, emoji)
  362. return nil
  363. }
  364. // EmojisAdd adds multiple emojis to the world state.
  365. func (s *State) EmojisAdd(guildID string, emojis []*Emoji) error {
  366. for _, e := range emojis {
  367. if err := s.EmojiAdd(guildID, e); err != nil {
  368. return err
  369. }
  370. }
  371. return nil
  372. }
  373. // MessageAdd adds a message to the current world state, or updates it if it exists.
  374. // If the channel cannot be found, the message is discarded.
  375. // Messages are kept in state up to s.MaxMessageCount
  376. func (s *State) MessageAdd(message *Message) error {
  377. if s == nil {
  378. return ErrNilState
  379. }
  380. c, err := s.Channel(message.ChannelID)
  381. if err != nil {
  382. return err
  383. }
  384. s.Lock()
  385. defer s.Unlock()
  386. // If the message exists, merge in the new message contents.
  387. for _, m := range c.Messages {
  388. if m.ID == message.ID {
  389. if message.Content != "" {
  390. m.Content = message.Content
  391. }
  392. if message.EditedTimestamp != "" {
  393. m.EditedTimestamp = message.EditedTimestamp
  394. }
  395. if message.Mentions != nil {
  396. m.Mentions = message.Mentions
  397. }
  398. if message.Embeds != nil {
  399. m.Embeds = message.Embeds
  400. }
  401. if message.Attachments != nil {
  402. m.Attachments = message.Attachments
  403. }
  404. if message.Timestamp != "" {
  405. m.Timestamp = message.Timestamp
  406. }
  407. if message.Author != nil {
  408. m.Author = message.Author
  409. }
  410. return nil
  411. }
  412. }
  413. c.Messages = append(c.Messages, message)
  414. if len(c.Messages) > s.MaxMessageCount {
  415. c.Messages = c.Messages[len(c.Messages)-s.MaxMessageCount:]
  416. }
  417. return nil
  418. }
  419. // MessageRemove removes a message from the world state.
  420. func (s *State) MessageRemove(message *Message) error {
  421. if s == nil {
  422. return ErrNilState
  423. }
  424. c, err := s.Channel(message.ChannelID)
  425. if err != nil {
  426. return err
  427. }
  428. s.Lock()
  429. defer s.Unlock()
  430. for i, m := range c.Messages {
  431. if m.ID == message.ID {
  432. c.Messages = append(c.Messages[:i], c.Messages[i+1:]...)
  433. return nil
  434. }
  435. }
  436. return errors.New("Message not found.")
  437. }
  438. func (s *State) voiceStateUpdate(update *VoiceStateUpdate) error {
  439. guild, err := s.Guild(update.GuildID)
  440. if err != nil {
  441. return err
  442. }
  443. s.Lock()
  444. defer s.Unlock()
  445. // Handle Leaving Channel
  446. if update.ChannelID == "" {
  447. for i, state := range guild.VoiceStates {
  448. if state.UserID == update.UserID {
  449. guild.VoiceStates = append(guild.VoiceStates[:i], guild.VoiceStates[i+1:]...)
  450. return nil
  451. }
  452. }
  453. } else {
  454. for i, state := range guild.VoiceStates {
  455. if state.UserID == update.UserID {
  456. guild.VoiceStates[i] = update.VoiceState
  457. return nil
  458. }
  459. }
  460. guild.VoiceStates = append(guild.VoiceStates, update.VoiceState)
  461. }
  462. return nil
  463. }
  464. // Message gets a message by channel and message ID.
  465. func (s *State) Message(channelID, messageID string) (*Message, error) {
  466. if s == nil {
  467. return nil, ErrNilState
  468. }
  469. c, err := s.Channel(channelID)
  470. if err != nil {
  471. return nil, err
  472. }
  473. s.RLock()
  474. defer s.RUnlock()
  475. for _, m := range c.Messages {
  476. if m.ID == messageID {
  477. return m, nil
  478. }
  479. }
  480. return nil, errors.New("Message not found.")
  481. }
  482. // OnReady takes a Ready event and updates all internal state.
  483. func (s *State) onReady(se *Session, r *Ready) (err error) {
  484. if s == nil {
  485. return ErrNilState
  486. }
  487. s.Lock()
  488. defer s.Unlock()
  489. // We must track at least the current user for Voice, even
  490. // if state is disabled, store the bare essentials.
  491. if !se.StateEnabled {
  492. ready := Ready{
  493. Version: r.Version,
  494. SessionID: r.SessionID,
  495. HeartbeatInterval: r.HeartbeatInterval,
  496. User: r.User,
  497. }
  498. s.Ready = ready
  499. return nil
  500. }
  501. s.Ready = *r
  502. for _, g := range s.Guilds {
  503. s.guildMap[g.ID] = g
  504. for _, c := range g.Channels {
  505. s.channelMap[c.ID] = c
  506. }
  507. }
  508. for _, c := range s.PrivateChannels {
  509. s.channelMap[c.ID] = c
  510. }
  511. return nil
  512. }
  513. // onInterface handles all events related to states.
  514. func (s *State) onInterface(se *Session, i interface{}) (err error) {
  515. if s == nil {
  516. return ErrNilState
  517. }
  518. r, ok := i.(*Ready)
  519. if ok {
  520. return s.onReady(se, r)
  521. }
  522. if !se.StateEnabled {
  523. return nil
  524. }
  525. switch t := i.(type) {
  526. case *GuildCreate:
  527. err = s.GuildAdd(t.Guild)
  528. case *GuildUpdate:
  529. err = s.GuildAdd(t.Guild)
  530. case *GuildDelete:
  531. err = s.GuildRemove(t.Guild)
  532. case *GuildMemberAdd:
  533. if s.TrackMembers {
  534. err = s.MemberAdd(t.Member)
  535. }
  536. case *GuildMemberUpdate:
  537. if s.TrackMembers {
  538. err = s.MemberAdd(t.Member)
  539. }
  540. case *GuildMemberRemove:
  541. if s.TrackMembers {
  542. err = s.MemberRemove(t.Member)
  543. }
  544. case *GuildRoleCreate:
  545. if s.TrackRoles {
  546. err = s.RoleAdd(t.GuildID, t.Role)
  547. }
  548. case *GuildRoleUpdate:
  549. if s.TrackRoles {
  550. err = s.RoleAdd(t.GuildID, t.Role)
  551. }
  552. case *GuildRoleDelete:
  553. if s.TrackRoles {
  554. err = s.RoleRemove(t.GuildID, t.RoleID)
  555. }
  556. case *GuildEmojisUpdate:
  557. if s.TrackEmojis {
  558. err = s.EmojisAdd(t.GuildID, t.Emojis)
  559. }
  560. case *ChannelCreate:
  561. if s.TrackChannels {
  562. err = s.ChannelAdd(t.Channel)
  563. }
  564. case *ChannelUpdate:
  565. if s.TrackChannels {
  566. err = s.ChannelAdd(t.Channel)
  567. }
  568. case *ChannelDelete:
  569. if s.TrackChannels {
  570. err = s.ChannelRemove(t.Channel)
  571. }
  572. case *MessageCreate:
  573. if s.MaxMessageCount != 0 {
  574. err = s.MessageAdd(t.Message)
  575. }
  576. case *MessageUpdate:
  577. if s.MaxMessageCount != 0 {
  578. err = s.MessageAdd(t.Message)
  579. }
  580. case *MessageDelete:
  581. if s.MaxMessageCount != 0 {
  582. err = s.MessageRemove(t.Message)
  583. }
  584. case *VoiceStateUpdate:
  585. if s.TrackVoice {
  586. err = s.voiceStateUpdate(t)
  587. }
  588. }
  589. return
  590. }
  591. // UserChannelPermissions returns the permission of a user in a channel.
  592. // userID : The ID of the user to calculate permissions for.
  593. // channelID : The ID of the channel to calculate permission for.
  594. func (s *State) UserChannelPermissions(userID, channelID string) (apermissions int, err error) {
  595. if s == nil {
  596. return 0, ErrNilState
  597. }
  598. channel, err := s.Channel(channelID)
  599. if err != nil {
  600. return
  601. }
  602. guild, err := s.Guild(channel.GuildID)
  603. if err != nil {
  604. return
  605. }
  606. if userID == guild.OwnerID {
  607. apermissions = PermissionAll
  608. return
  609. }
  610. member, err := s.Member(guild.ID, userID)
  611. if err != nil {
  612. return
  613. }
  614. for _, role := range guild.Roles {
  615. if role.ID == guild.ID {
  616. apermissions |= role.Permissions
  617. break
  618. }
  619. }
  620. for _, role := range guild.Roles {
  621. for _, roleID := range member.Roles {
  622. if role.ID == roleID {
  623. apermissions |= role.Permissions
  624. break
  625. }
  626. }
  627. }
  628. if apermissions&PermissionAdministrator > 0 {
  629. apermissions |= PermissionAll
  630. }
  631. // Member overwrites can override role overrides, so do two passes
  632. for _, overwrite := range channel.PermissionOverwrites {
  633. for _, roleID := range member.Roles {
  634. if overwrite.Type == "role" && roleID == overwrite.ID {
  635. apermissions &= ^overwrite.Deny
  636. apermissions |= overwrite.Allow
  637. break
  638. }
  639. }
  640. }
  641. for _, overwrite := range channel.PermissionOverwrites {
  642. if overwrite.Type == "member" && overwrite.ID == userID {
  643. apermissions &= ^overwrite.Deny
  644. apermissions |= overwrite.Allow
  645. break
  646. }
  647. }
  648. if apermissions&PermissionAdministrator > 0 {
  649. apermissions |= PermissionAllChannel
  650. }
  651. return
  652. }