sse-encoder.go 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. // Copyright 2014 Manu Martinez-Almeida. All rights reserved.
  2. // Use of this source code is governed by a MIT style
  3. // license that can be found in the LICENSE file.
  4. package sse
  5. import (
  6. "encoding/json"
  7. "fmt"
  8. "io"
  9. "net/http"
  10. "reflect"
  11. "strconv"
  12. "strings"
  13. )
  14. // Server-Sent Events
  15. // W3C Working Draft 29 October 2009
  16. // http://www.w3.org/TR/2009/WD-eventsource-20091029/
  17. const ContentType = "text/event-stream;charset=utf-8"
  18. var (
  19. contentType = []string{ContentType}
  20. noCache = []string{"no-cache"}
  21. )
  22. var fieldReplacer = strings.NewReplacer(
  23. "\n", "\\n",
  24. "\r", "\\r")
  25. var dataReplacer = strings.NewReplacer(
  26. "\n", "\ndata:",
  27. "\r", "\\r")
  28. type Event struct {
  29. Event string
  30. Id string
  31. Retry uint
  32. Data interface{}
  33. }
  34. func Encode(writer io.Writer, event Event) error {
  35. w := checkWriter(writer)
  36. writeId(w, event.Id)
  37. writeEvent(w, event.Event)
  38. writeRetry(w, event.Retry)
  39. return writeData(w, event.Data)
  40. }
  41. func writeId(w stringWriter, id string) {
  42. if len(id) > 0 {
  43. _, _ = w.WriteString("id:")
  44. _, _ = fieldReplacer.WriteString(w, id)
  45. _, _ = w.WriteString("\n")
  46. }
  47. }
  48. func writeEvent(w stringWriter, event string) {
  49. if len(event) > 0 {
  50. _, _ = w.WriteString("event:")
  51. _, _ = fieldReplacer.WriteString(w, event)
  52. _, _ = w.WriteString("\n")
  53. }
  54. }
  55. func writeRetry(w stringWriter, retry uint) {
  56. if retry > 0 {
  57. _, _ = w.WriteString("retry:")
  58. _, _ = w.WriteString(strconv.FormatUint(uint64(retry), 10))
  59. _, _ = w.WriteString("\n")
  60. }
  61. }
  62. func writeData(w stringWriter, data interface{}) error {
  63. _, _ = w.WriteString("data:")
  64. bData, ok := data.([]byte)
  65. if ok {
  66. _, _ = dataReplacer.WriteString(w, string(bData))
  67. _, _ = w.WriteString("\n\n")
  68. return nil
  69. }
  70. switch kindOfData(data) { //nolint:exhaustive
  71. case reflect.Struct, reflect.Slice, reflect.Map:
  72. err := json.NewEncoder(w).Encode(data)
  73. if err != nil {
  74. return err
  75. }
  76. _, _ = w.WriteString("\n")
  77. default:
  78. _, _ = dataReplacer.WriteString(w, fmt.Sprint(data))
  79. _, _ = w.WriteString("\n\n")
  80. }
  81. return nil
  82. }
  83. func (r Event) Render(w http.ResponseWriter) error {
  84. r.WriteContentType(w)
  85. return Encode(w, r)
  86. }
  87. func (r Event) WriteContentType(w http.ResponseWriter) {
  88. header := w.Header()
  89. header["Content-Type"] = contentType
  90. if _, exist := header["Cache-Control"]; !exist {
  91. header["Cache-Control"] = noCache
  92. }
  93. }
  94. func kindOfData(data interface{}) reflect.Kind {
  95. value := reflect.ValueOf(data)
  96. valueType := value.Kind()
  97. if valueType == reflect.Ptr {
  98. valueType = value.Elem().Kind()
  99. }
  100. return valueType
  101. }