github.com/segakazzz/buffalo@v0.16.22-0.20210119082501-1f52048d3feb/render/sse.go (about)

     1  package render
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"net/http"
     7  )
     8  
     9  type sse struct {
    10  	Data interface{} `json:"data"`
    11  	Type string      `json:"type"`
    12  }
    13  
    14  func (s *sse) String() string {
    15  	b, _ := json.Marshal(s)
    16  	return fmt.Sprintf("data: %s\n\n", string(b))
    17  }
    18  
    19  func (s *sse) Bytes() []byte {
    20  	return []byte(s.String())
    21  }
    22  
    23  // EventSource is designed to work with JavaScript EventSource objects.
    24  // see https://developer.mozilla.org/en-US/docs/Web/API/EventSource for
    25  // more details
    26  type EventSource struct {
    27  	w  http.ResponseWriter
    28  	fl http.Flusher
    29  }
    30  
    31  func (es *EventSource) Write(t string, d interface{}) error {
    32  	s := &sse{Type: t, Data: d}
    33  	_, err := es.w.Write(s.Bytes())
    34  	if err != nil {
    35  		return err
    36  	}
    37  	es.Flush()
    38  	return nil
    39  }
    40  
    41  // Flush messages down the pipe. If messages aren't flushed they
    42  // won't be sent.
    43  func (es *EventSource) Flush() {
    44  	es.fl.Flush()
    45  }
    46  
    47  type closeNotifier interface {
    48  	CloseNotify() <-chan bool
    49  }
    50  
    51  // CloseNotify return true across the channel when the connection
    52  // in the browser has been severed.
    53  func (es *EventSource) CloseNotify() <-chan bool {
    54  	if cn, ok := es.w.(closeNotifier); ok {
    55  		return cn.CloseNotify()
    56  	}
    57  	return nil
    58  }
    59  
    60  // NewEventSource returns a new EventSource instance while ensuring
    61  // that the http.ResponseWriter is able to handle EventSource messages.
    62  // It also makes sure to set the proper response heads.
    63  func NewEventSource(w http.ResponseWriter) (*EventSource, error) {
    64  	es := &EventSource{w: w}
    65  	var ok bool
    66  	es.fl, ok = w.(http.Flusher)
    67  	if !ok {
    68  		return es, fmt.Errorf("streaming is not supported")
    69  	}
    70  
    71  	es.w.Header().Set("Content-Type", "text/event-stream")
    72  	es.w.Header().Set("Cache-Control", "no-cache")
    73  	es.w.Header().Set("Connection", "keep-alive")
    74  	es.w.Header().Set("Access-Control-Allow-Origin", "*")
    75  	return es, nil
    76  }