108 lines
2.2 KiB
Go
108 lines
2.2 KiB
Go
package main
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"slices"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func InitSse() *SseStream {
|
|
sse := &SseStream{
|
|
Clients: make([]Client, 0),
|
|
MsgId: map[string]int{},
|
|
Events: events,
|
|
Data: initPriority(),
|
|
}
|
|
|
|
for i := range events {
|
|
sse.MsgId[events[i]] = 0
|
|
}
|
|
|
|
return sse
|
|
}
|
|
|
|
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("event"),
|
|
}
|
|
|
|
log.Printf("events: %+v", client.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 - TOT %+v", client.Ip, sse.Clients)
|
|
|
|
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
|
|
}
|
|
}
|
|
}
|
|
|
|
func (sse *SseStream) Send(msg Message) {
|
|
sse.MsgId[msg.Event]++
|
|
msg.Id = strconv.Itoa(sse.MsgId[msg.Event])
|
|
for _, client := range sse.Clients {
|
|
if slices.Contains(client.Events, msg.Event) || slices.Contains(client.Events, "*") {
|
|
client.Chan <- msg
|
|
log.Printf("sent: %+v -> %+v", msg, client.Ip)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (sse *SseStream) SendMsg(c *gin.Context) {
|
|
var msg Message
|
|
c.ShouldBind(&msg)
|
|
sse.Send(msg)
|
|
c.JSON(http.StatusOK, gin.H{"status": "msg"})
|
|
}
|
|
|
|
func generateClientId(ip string, events []string, len int) string {
|
|
str := ip + strings.Join(events, "|")
|
|
id := sha256.Sum256([]byte(str))
|
|
|
|
return fmt.Sprintf("%X", id)[0:len]
|
|
}
|