瀏覽代碼

Adds support for Server-Sent Events

Manu Mtz-Almeida 10 年之前
父節點
當前提交
470b7e1010
共有 2 個文件被更改,包括 78 次插入2 次删除
  1. 20 2
      context.go
  2. 58 0
      render/ssevent.go

+ 20 - 2
context.go

@@ -6,6 +6,7 @@ package gin
 
 
 import (
 import (
 	"errors"
 	"errors"
+	"io"
 	"math"
 	"math"
 	"net/http"
 	"net/http"
 	"strings"
 	"strings"
@@ -379,7 +380,21 @@ func (c *Context) File(filepath string) {
 	http.ServeFile(c.Writer, c.Request, filepath)
 	http.ServeFile(c.Writer, c.Request, filepath)
 }
 }
 
 
-func (c *Context) Stream(step func(w http.ResponseWriter)) {
+func (c *Context) SSEvent(name string, message interface{}) {
+	render.WriteSSEvent(c.Writer, name, message)
+}
+
+func (c *Context) Header(code int, headers map[string]string) {
+	if len(headers) > 0 {
+		header := c.Writer.Header()
+		for key, value := range headers {
+			header.Set(key, value)
+		}
+	}
+	c.Writer.WriteHeader(code)
+}
+
+func (c *Context) Stream(step func(w io.Writer) bool) {
 	w := c.Writer
 	w := c.Writer
 	clientGone := w.CloseNotify()
 	clientGone := w.CloseNotify()
 	for {
 	for {
@@ -387,8 +402,11 @@ func (c *Context) Stream(step func(w http.ResponseWriter)) {
 		case <-clientGone:
 		case <-clientGone:
 			return
 			return
 		default:
 		default:
-			step(w)
+			keepopen := step(w)
 			w.Flush()
 			w.Flush()
+			if !keepopen {
+				return
+			}
 		}
 		}
 	}
 	}
 }
 }

+ 58 - 0
render/ssevent.go

@@ -0,0 +1,58 @@
+package render
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"reflect"
+)
+
+type sseRender struct{}
+
+var SSEvent Render = sseRender{}
+
+func (_ sseRender) Render(w http.ResponseWriter, code int, data ...interface{}) error {
+	eventName := data[0].(string)
+	obj := data[1]
+	return WriteSSEvent(w, eventName, obj)
+}
+
+func WriteSSEvent(w http.ResponseWriter, eventName string, data interface{}) error {
+	header := w.Header()
+	if len(header.Get("Content-Type")) == 0 {
+		w.Header().Set("Content-Type", "text/event-stream")
+	}
+	var stringData string
+	switch typeOfData(data) {
+	case reflect.Struct, reflect.Slice:
+		if jsonBytes, err := json.Marshal(data); err == nil {
+			stringData = string(jsonBytes)
+		} else {
+			return err
+		}
+	case reflect.Ptr:
+		stringData = escape(fmt.Sprintf("%v", &data)) + "\n"
+	default:
+		stringData = escape(fmt.Sprintf("%v", data)) + "\n"
+	}
+	_, err := fmt.Fprintf(w, "event: %s\ndata: %s\n", escape(eventName), stringData)
+	return err
+}
+
+func typeOfData(data interface{}) reflect.Kind {
+	value := reflect.ValueOf(data)
+	valueType := value.Kind()
+	if valueType == reflect.Ptr {
+		newValue := value.Elem().Kind()
+		if newValue == reflect.Struct || newValue == reflect.Slice {
+			return newValue
+		} else {
+			return valueType
+		}
+	}
+	return valueType
+}
+
+func escape(str string) string {
+	return str
+}