-
Notifications
You must be signed in to change notification settings - Fork 0
/
writer.go
96 lines (82 loc) · 2.03 KB
/
writer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
package sse
import (
"encoding/json"
"fmt"
"net/http"
)
const (
// NonceMax is the maximum value of the `id` field in SSE before it resets to 0.
NonceMax = 1<<63 - 1
)
// Writer is the interface for writing Server-Sent Events.
type Writer interface {
Write(event string, data interface{}) error
}
// Options holds configuration for the SSE writer.
type Options struct {
ResponseStatus int
Encoding string
}
// NewResponseWriter creates a new Writer for Server-Sent Events.
func NewResponseWriter(w http.ResponseWriter, opts Options) Writer {
rw := &responseWriter{
writer: w,
nonce: 0,
options: opts,
}
rw.sendHeaders()
return rw
}
type responseWriter struct {
writer http.ResponseWriter
nonce uint64
options Options
}
// Write sends a message to the client.
func (rw *responseWriter) Write(event string, data interface{}) error {
rw.nonce = (rw.nonce + 1) % NonceMax
output := fmt.Sprintf("id: %d\n", rw.nonce)
if event != "" {
output += fmt.Sprintf("event: %s\n", event)
}
if data != nil {
encodedData, err := json.Marshal(data)
if err != nil {
return err
}
output += fmt.Sprintf("data: %s\n", encodedData)
}
output += "\n"
encodedOutput, err := encode(rw.options.Encoding, output)
if err != nil {
return err
}
if _, err := rw.writer.Write(encodedOutput); err != nil {
return err
}
return rw.flush()
}
// sendHeaders sends the headers for Server-Sent Events.
func (rw *responseWriter) sendHeaders() {
headers := rw.writer.Header()
headers.Set("Content-Type", "text/event-stream")
headers.Set("Cache-Control", "no-store")
headers.Set("Connection", "keep-alive")
if rw.options.Encoding != EncodeNone {
headers.Set("Content-Encoding", rw.options.Encoding)
}
status := rw.options.ResponseStatus
if status == 0 {
status = http.StatusOK
}
rw.writer.WriteHeader(status)
rw.flush()
}
// flush flushes the response.
func (rw *responseWriter) flush() error {
if flusher, ok := rw.writer.(http.Flusher); ok {
flusher.Flush()
return nil
}
return fmt.Errorf("ResponseWriter is not a Flusher")
}