diff --git a/.gitignore b/.gitignore index d4d9525..cacc2aa 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ .idea .vscode .DS_Store +xtest \ No newline at end of file diff --git a/app.go b/app.go index 8e7fe8f..8dc5930 100644 --- a/app.go +++ b/app.go @@ -19,8 +19,6 @@ var ( regSafePrefix = regexp.MustCompile("[^a-zA-Z0-9/-]+") regRemoveRepeatedChar = regexp.MustCompile("/{2,}") - - mimePlain = []string{"text/plain"} ) type App struct { @@ -176,7 +174,6 @@ func (a *App) handleHTTPRequest(c *Ctx) { serveError(c, errorHandler) } - c.writermem.WriteHeaderNow() return } if httpMethod != http.MethodConnect && rPath != "/" { diff --git a/ctx.go b/ctx.go index 7ad8566..1b1ea50 100644 --- a/ctx.go +++ b/ctx.go @@ -3,7 +3,9 @@ package nf import ( "bytes" "encoding/json" + "errors" "fmt" + "github.com/loveuer/nf/internal/sse" "io" "mime/multipart" "net" @@ -281,7 +283,7 @@ func (c *Ctx) SendString(data string) error { func (c *Ctx) Writef(format string, values ...interface{}) (int, error) { c.SetHeader("Content-Type", "text/plain") - return c.writer.Write([]byte(fmt.Sprintf(format, values...))) + return c.Write([]byte(fmt.Sprintf(format, values...))) } func (c *Ctx) JSON(data interface{}) error { @@ -296,6 +298,23 @@ func (c *Ctx) JSON(data interface{}) error { return nil } +func (c *Ctx) SSEvent(event string, data interface{}) error { + c.Set("Content-Type", "text/event-stream") + c.Set("Cache-Control", "no-cache") + c.Set("Transfer-Encoding", "chunked") + + return sse.Encode(c.writer, sse.Event{Event: event, Data: data}) +} + +func (c *Ctx) Flush() error { + if f, ok := c.writer.(http.Flusher); ok { + f.Flush() + return nil + } + + return errors.New("http.Flusher is not implemented") +} + func (c *Ctx) RawWriter() http.ResponseWriter { return c.writer } diff --git a/internal/sse/sse-encoder.go b/internal/sse/sse-encoder.go new file mode 100644 index 0000000..26100e7 --- /dev/null +++ b/internal/sse/sse-encoder.go @@ -0,0 +1,106 @@ +package sse + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "reflect" + "strconv" + "strings" +) + +// Server-Sent Events +// W3C Working Draft 29 October 2009 +// http://www.w3.org/TR/2009/WD-eventsource-20091029/ + +const ContentType = "text/event-stream" + +var contentType = []string{ContentType} +var noCache = []string{"no-cache"} + +var fieldReplacer = strings.NewReplacer( + "\n", "\\n", + "\r", "\\r") + +var dataReplacer = strings.NewReplacer( + "\n", "\ndata:", + "\r", "\\r") + +type Event struct { + Event string + Id string + Retry uint + Data interface{} +} + +func Encode(writer io.Writer, event Event) error { + w := checkWriter(writer) + writeId(w, event.Id) + writeEvent(w, event.Event) + writeRetry(w, event.Retry) + return writeData(w, event.Data) +} + +func writeId(w stringWriter, id string) { + if len(id) > 0 { + w.WriteString("id:") + fieldReplacer.WriteString(w, id) + w.WriteString("\n") + } +} + +func writeEvent(w stringWriter, event string) { + if len(event) > 0 { + w.WriteString("event:") + fieldReplacer.WriteString(w, event) + w.WriteString("\n") + } +} + +func writeRetry(w stringWriter, retry uint) { + if retry > 0 { + w.WriteString("retry:") + w.WriteString(strconv.FormatUint(uint64(retry), 10)) + w.WriteString("\n") + } +} + +func writeData(w stringWriter, data interface{}) error { + w.WriteString("data:") + switch kindOfData(data) { + case reflect.Struct, reflect.Slice, reflect.Map: + err := json.NewEncoder(w).Encode(data) + if err != nil { + return err + } + w.WriteString("\n") + default: + dataReplacer.WriteString(w, fmt.Sprint(data)) + w.WriteString("\n\n") + } + return nil +} + +func (r Event) Render(w http.ResponseWriter) error { + r.WriteContentType(w) + return Encode(w, r) +} + +func (r Event) WriteContentType(w http.ResponseWriter) { + header := w.Header() + header["Content-Type"] = contentType + + if _, exist := header["Cache-Control"]; !exist { + header["Cache-Control"] = noCache + } +} + +func kindOfData(data interface{}) reflect.Kind { + value := reflect.ValueOf(data) + valueType := value.Kind() + if valueType == reflect.Ptr { + valueType = value.Elem().Kind() + } + return valueType +} diff --git a/internal/sse/writer.go b/internal/sse/writer.go new file mode 100644 index 0000000..6f9806c --- /dev/null +++ b/internal/sse/writer.go @@ -0,0 +1,24 @@ +package sse + +import "io" + +type stringWriter interface { + io.Writer + WriteString(string) (int, error) +} + +type stringWrapper struct { + io.Writer +} + +func (w stringWrapper) WriteString(str string) (int, error) { + return w.Writer.Write([]byte(str)) +} + +func checkWriter(writer io.Writer) stringWriter { + if w, ok := writer.(stringWriter); ok { + return w + } else { + return stringWrapper{writer} + } +} diff --git a/response_writer.go b/response_writer.go index 185cddc..62a1fcf 100644 --- a/response_writer.go +++ b/response_writer.go @@ -3,6 +3,7 @@ package nf import ( "bufio" "io" + "log" "net" "net/http" ) @@ -60,7 +61,7 @@ func (w *responseWriter) reset(writer http.ResponseWriter) { func (w *responseWriter) WriteHeader(code int) { if code > 0 && w.status != code { if w.Written() { - // todo: debugPrint("[WARNING] Headers were already written. Wanted to override status code %d with %d", w.status, code) + log.Printf("[NF] WARNING: Headers were already written. Wanted to override status code %d with %d", w.status, code) return } w.status = code @@ -102,7 +103,7 @@ func (w *responseWriter) Size() int { } func (w *responseWriter) Written() bool { - return w.size != noWritten + return w.size != noWritten || w.status != 0 } // Hijack implements the http.Hijacker interface.