bvbej/base-golang
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 17 Wiki Activity

[🚀] sse

Browse Source
This commit is contained in:
bvbej 2024-09-07 11:25:30 +08:00
parent bbf89b2f08
commit be7b2bf15e
1 changed files with 37 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
pkg/sse/server.go
View File

@ -1,6 +1,7 @@
package sse package sse
import ( import (
"gitea.bvbej.com/bvbej/base-golang/pkg/mux"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"io" "io"
"net/http" "net/http"
@ -11,13 +12,14 @@ import (
var _ Server = (*event)(nil) var _ Server = (*event)(nil)
type Server interface { type Server interface {
HandlerFunc(auth func(c *gin.Context) (string, error)) gin.HandlerFunc GinHandle(ctx *gin.Context, user any)
HandlerFunc() mux.HandlerFunc
Push(user, name, msg string) bool Push(user, name, msg string) bool
Broadcast(name, msg string) Broadcast(name, msg string)
} }
type clientChan struct { type clientChan struct {
User string User any
Chan chan msgChan Chan chan msgChan
} }
@ -31,7 +33,7 @@ type event struct {
Count atomic.Int32 Count atomic.Int32
Register chan clientChan Register chan clientChan
Unregister chan string Unregister chan any
} }
func NewServer() Server { func NewServer() Server {
@ -39,7 +41,7 @@ func NewServer() Server {
SessionList: sync.Map{}, SessionList: sync.Map{},
Count: atomic.Int32{}, Count: atomic.Int32{},
Register: make(chan clientChan), Register: make(chan clientChan),
Unregister: make(chan string), Unregister: make(chan any),
} }
go e.listen() go e.listen()
@ -65,14 +67,11 @@ func (stream *event) listen() {
} }
} }
func (stream *event) HandlerFunc(auth func(c *gin.Context) (string, error)) gin.HandlerFunc { func (stream *event) GinHandle(ctx *gin.Context, user any) {
return func(c *gin.Context) { if user == nil {
user, err := auth(c) ctx.AbortWithStatus(http.StatusUnauthorized)
if err != nil {
c.AbortWithStatus(http.StatusBadRequest)
return return
} }
e := make(chan msgChan) e := make(chan msgChan)
client := clientChan{ client := clientChan{
User: user, User: user,
@ -83,20 +82,25 @@ func (stream *event) HandlerFunc(auth func(c *gin.Context) (string, error)) gin.
stream.Unregister <- user stream.Unregister <- user
}() }()
c.Writer.Header().Set("Content-Type", "text/event-stream") ctx.Writer.Header().Set("Content-Type", "text/event-stream")
c.Writer.Header().Set("Cache-Control", "no-cache") ctx.Writer.Header().Set("Cache-Control", "no-cache")
c.Writer.Header().Set("Connection", "keep-alive") ctx.Writer.Header().Set("Connection", "keep-alive")
c.Writer.Header().Set("Transfer-Encoding", "chunked") ctx.Writer.Header().Set("Transfer-Encoding", "chunked")
c.Stream(func(w io.Writer) bool { ctx.Stream(func(w io.Writer) bool {
if msg, ok := <-e; ok { if msg, ok := <-e; ok {
c.SSEvent(msg.Name, msg.Message) ctx.SSEvent(msg.Name, msg.Message)
return true return true
} }
return false return false
}) })
c.Next() ctx.Next()
}
func (stream *event) HandlerFunc() mux.HandlerFunc {
return func(c mux.Context) {
stream.GinHandle(c.Context(), c.Auth())
} }
} }