另客网go项目公用的代码库

sse-encoder.go 2.2KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  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"
  18. var contentType = []string{ContentType}
  19. var noCache = []string{"no-cache"}
  20. var fieldReplacer = strings.NewReplacer(
  21. "\n", "\\n",
  22. "\r", "\\r")
  23. var dataReplacer = strings.NewReplacer(
  24. "\n", "\ndata:",
  25. "\r", "\\r")
  26. type Event struct {
  27. Event string
  28. Id string
  29. Retry uint
  30. Data interface{}
  31. }
  32. func Encode(writer io.Writer, event Event) error {
  33. w := checkWriter(writer)
  34. writeId(w, event.Id)
  35. writeEvent(w, event.Event)
  36. writeRetry(w, event.Retry)
  37. return writeData(w, event.Data)
  38. }
  39. func writeId(w stringWriter, id string) {
  40. if len(id) > 0 {
  41. w.WriteString("id:")
  42. fieldReplacer.WriteString(w, id)
  43. w.WriteString("\n")
  44. }
  45. }
  46. func writeEvent(w stringWriter, event string) {
  47. if len(event) > 0 {
  48. w.WriteString("event:")
  49. fieldReplacer.WriteString(w, event)
  50. w.WriteString("\n")
  51. }
  52. }
  53. func writeRetry(w stringWriter, retry uint) {
  54. if retry > 0 {
  55. w.WriteString("retry:")
  56. w.WriteString(strconv.FormatUint(uint64(retry), 10))
  57. w.WriteString("\n")
  58. }
  59. }
  60. func writeData(w stringWriter, data interface{}) error {
  61. w.WriteString("data:")
  62. switch kindOfData(data) {
  63. case reflect.Struct, reflect.Slice, reflect.Map:
  64. err := json.NewEncoder(w).Encode(data)
  65. if err != nil {
  66. return err
  67. }
  68. w.WriteString("\n")
  69. default:
  70. dataReplacer.WriteString(w, fmt.Sprint(data))
  71. w.WriteString("\n\n")
  72. }
  73. return nil
  74. }
  75. func (r Event) Render(w http.ResponseWriter) error {
  76. r.WriteContentType(w)
  77. return Encode(w, r)
  78. }
  79. func (r Event) WriteContentType(w http.ResponseWriter) {
  80. header := w.Header()
  81. header["Content-Type"] = contentType
  82. if _, exist := header["Cache-Control"]; !exist {
  83. header["Cache-Control"] = noCache
  84. }
  85. }
  86. func kindOfData(data interface{}) reflect.Kind {
  87. value := reflect.ValueOf(data)
  88. valueType := value.Kind()
  89. if valueType == reflect.Ptr {
  90. valueType = value.Elem().Kind()
  91. }
  92. return valueType
  93. }