discord.go 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218
  1. // Copyright 2018 Unknwon
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License"): you may
  4. // not use this file except in compliance with the License. You may obtain
  5. // a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  11. // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  12. // License for the specific language governing permissions and limitations
  13. // under the License.
  14. package clog
  15. import (
  16. "bytes"
  17. "encoding/json"
  18. "errors"
  19. "fmt"
  20. "io"
  21. "io/ioutil"
  22. "net/http"
  23. "time"
  24. )
  25. type (
  26. discordEmbed struct {
  27. Title string `json:"title"`
  28. Description string `json:"description"`
  29. Timestamp string `json:"timestamp"`
  30. Color int `json:"color"`
  31. }
  32. discordPayload struct {
  33. Username string `json:"username,omitempty"`
  34. Embeds []*discordEmbed `json:"embeds"`
  35. }
  36. )
  37. var (
  38. discordTitles = []string{
  39. "Tracing",
  40. "Information",
  41. "Warning",
  42. "Error",
  43. "Fatal",
  44. }
  45. discordColors = []int{
  46. 0, // Trace
  47. 3843043, // Info
  48. 16761600, // Warn
  49. 13041721, // Error
  50. 9440319, // Fatal
  51. }
  52. )
  53. type DiscordConfig struct {
  54. // Minimum level of messages to be processed.
  55. Level LEVEL
  56. // Buffer size defines how many messages can be queued before hangs.
  57. BufferSize int64
  58. // Discord webhook URL.
  59. URL string
  60. // Username to be shown for the message.
  61. // Leave empty to use default as set in the Discord.
  62. Username string
  63. }
  64. type discord struct {
  65. Adapter
  66. url string
  67. username string
  68. }
  69. func newDiscord() Logger {
  70. return &discord{
  71. Adapter: Adapter{
  72. quitChan: make(chan struct{}),
  73. },
  74. }
  75. }
  76. func (d *discord) Level() LEVEL { return d.level }
  77. func (d *discord) Init(v interface{}) error {
  78. cfg, ok := v.(DiscordConfig)
  79. if !ok {
  80. return ErrConfigObject{"DiscordConfig", v}
  81. }
  82. if !isValidLevel(cfg.Level) {
  83. return ErrInvalidLevel{}
  84. }
  85. d.level = cfg.Level
  86. if len(cfg.URL) == 0 {
  87. return errors.New("URL cannot be empty")
  88. }
  89. d.url = cfg.URL
  90. d.username = cfg.Username
  91. d.msgChan = make(chan *Message, cfg.BufferSize)
  92. return nil
  93. }
  94. func (d *discord) ExchangeChans(errorChan chan<- error) chan *Message {
  95. d.errorChan = errorChan
  96. return d.msgChan
  97. }
  98. func buildDiscordPayload(username string, msg *Message) (string, error) {
  99. payload := discordPayload{
  100. Username: username,
  101. Embeds: []*discordEmbed{
  102. {
  103. Title: discordTitles[msg.Level],
  104. Description: msg.Body[8:],
  105. Timestamp: time.Now().Format(time.RFC3339),
  106. Color: discordColors[msg.Level],
  107. },
  108. },
  109. }
  110. p, err := json.Marshal(&payload)
  111. if err != nil {
  112. return "", err
  113. }
  114. return string(p), nil
  115. }
  116. type rateLimitMsg struct {
  117. RetryAfter int64 `json:"retry_after"`
  118. }
  119. func (d *discord) postMessage(r io.Reader) (int64, error) {
  120. resp, err := http.Post(d.url, "application/json", r)
  121. if err != nil {
  122. return -1, fmt.Errorf("HTTP Post: %v", err)
  123. }
  124. defer resp.Body.Close()
  125. if resp.StatusCode == 429 {
  126. rlMsg := &rateLimitMsg{}
  127. if err = json.NewDecoder(resp.Body).Decode(&rlMsg); err != nil {
  128. return -1, fmt.Errorf("decode rate limit message: %v", err)
  129. }
  130. return rlMsg.RetryAfter, nil
  131. } else if resp.StatusCode/100 != 2 {
  132. data, _ := ioutil.ReadAll(resp.Body)
  133. return -1, fmt.Errorf("%s", data)
  134. }
  135. return -1, nil
  136. }
  137. func (d *discord) write(msg *Message) {
  138. payload, err := buildDiscordPayload(d.username, msg)
  139. if err != nil {
  140. d.errorChan <- fmt.Errorf("discord: builddiscordPayload: %v", err)
  141. return
  142. }
  143. const RETRY_TIMES = 3
  144. // Due to discord limit, try at most x times with respect to "retry_after" parameter.
  145. for i := 1; i <= 3; i++ {
  146. retryAfter, err := d.postMessage(bytes.NewReader([]byte(payload)))
  147. if err != nil {
  148. d.errorChan <- fmt.Errorf("discord: postMessage: %v", err)
  149. return
  150. }
  151. if retryAfter > 0 {
  152. time.Sleep(time.Duration(retryAfter) * time.Millisecond)
  153. continue
  154. }
  155. return
  156. }
  157. d.errorChan <- fmt.Errorf("discord: failed to send message after %d retries", RETRY_TIMES)
  158. }
  159. func (d *discord) Start() {
  160. LOOP:
  161. for {
  162. select {
  163. case msg := <-d.msgChan:
  164. d.write(msg)
  165. case <-d.quitChan:
  166. break LOOP
  167. }
  168. }
  169. for {
  170. if len(d.msgChan) == 0 {
  171. break
  172. }
  173. d.write(<-d.msgChan)
  174. }
  175. d.quitChan <- struct{}{} // Notify the cleanup is done.
  176. }
  177. func (d *discord) Destroy() {
  178. d.quitChan <- struct{}{}
  179. <-d.quitChan
  180. close(d.msgChan)
  181. close(d.quitChan)
  182. }
  183. func init() {
  184. Register(DISCORD, newDiscord)
  185. }