93 lines
1.8 KiB
Go
93 lines
1.8 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"io"
|
||
|
|
"log"
|
||
|
|
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
)
|
||
|
|
|
||
|
|
var events = []string{"message", "timer", "priority"}
|
||
|
|
|
||
|
|
type Message struct {
|
||
|
|
Event string `json:"event"`
|
||
|
|
Data string `json:"data"`
|
||
|
|
Cmd string `json:"cmd"`
|
||
|
|
Id string `json:"id"`
|
||
|
|
}
|
||
|
|
|
||
|
|
type ClientChan chan Message
|
||
|
|
|
||
|
|
type Client struct {
|
||
|
|
Id string `json:"id"`
|
||
|
|
Ip string `json:"ip"`
|
||
|
|
Chan ClientChan `json:"chan"`
|
||
|
|
Events []string `json:"events"`
|
||
|
|
}
|
||
|
|
|
||
|
|
type SseStream struct {
|
||
|
|
Clients []Client `json:"clients"`
|
||
|
|
MsgId int `json:"msgid"`
|
||
|
|
Events []string `json:"events"`
|
||
|
|
}
|
||
|
|
|
||
|
|
func InitSse() *SseStream {
|
||
|
|
return &SseStream{
|
||
|
|
Clients: make([]Client, 0),
|
||
|
|
MsgId: 0,
|
||
|
|
Events: events,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (sse *SseStream) SseHeadersMiddleware() gin.HandlerFunc {
|
||
|
|
return func(c *gin.Context) {
|
||
|
|
c.Writer.Header().Set("Content-Type", "text/event-stream")
|
||
|
|
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||
|
|
c.Writer.Header().Set("Connection", "keep-alive")
|
||
|
|
c.Writer.Header().Set("Transfer-Encoding", "chunked")
|
||
|
|
c.Next()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (sse *SseStream) Stream(c *gin.Context) {
|
||
|
|
client := &Client{
|
||
|
|
Ip: c.Request.RemoteAddr,
|
||
|
|
Chan: make(ClientChan),
|
||
|
|
Events: c.QueryArray("events"),
|
||
|
|
}
|
||
|
|
|
||
|
|
client.Id = generateClientId(client.Ip, client.Events, 8)
|
||
|
|
|
||
|
|
sse.AddClient(*client)
|
||
|
|
|
||
|
|
defer func() {
|
||
|
|
sse.RemoveClient(client)
|
||
|
|
close(client.Chan)
|
||
|
|
log.Printf("Client %s disconnected", client.Ip)
|
||
|
|
}()
|
||
|
|
|
||
|
|
log.Printf("Client %s connected", client.Ip)
|
||
|
|
|
||
|
|
c.Stream(func(w io.Writer) bool {
|
||
|
|
msg, ok := <-client.Chan
|
||
|
|
if !ok {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
c.SSEvent(msg.Event, msg)
|
||
|
|
return true
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func (sse *SseStream) AddClient(client Client) {
|
||
|
|
sse.Clients = append(sse.Clients, client)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (sse *SseStream) RemoveClient(client *Client) {
|
||
|
|
for i, c := range sse.Clients {
|
||
|
|
if c.Id == client.Id {
|
||
|
|
sse.Clients = append(sse.Clients[:i], sse.Clients[i+1:]...)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|