322 lines
8.4 KiB
Go
322 lines
8.4 KiB
Go
package stream
|
|
|
|
import (
|
|
"net/http"
|
|
"net/url"
|
|
"regexp"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/gorilla/websocket"
|
|
"github.com/gotify/server/v2/auth"
|
|
"github.com/gotify/server/v2/mode"
|
|
"github.com/gotify/server/v2/model"
|
|
)
|
|
|
|
// shard holds clients for a subset of users to reduce lock contention.
|
|
type shard struct {
|
|
clients map[uint]map[string][]*client // userID -> token -> []client (multiple clients can share same token)
|
|
lock sync.RWMutex
|
|
}
|
|
|
|
// The API provides a handler for a WebSocket stream API.
|
|
type API struct {
|
|
shards []*shard
|
|
shardCount int
|
|
pingPeriod time.Duration
|
|
pongTimeout time.Duration
|
|
upgrader *websocket.Upgrader
|
|
channelBuf int // Buffer size for client write channels
|
|
}
|
|
|
|
// New creates a new instance of API.
|
|
// pingPeriod: is the interval, in which is server sends the a ping to the client.
|
|
// pongTimeout: is the duration after the connection will be terminated, when the client does not respond with the
|
|
// pong command.
|
|
// shardCount: number of shards for client storage (should be power of 2 for optimal distribution).
|
|
// readBufferSize: WebSocket read buffer size in bytes.
|
|
// writeBufferSize: WebSocket write buffer size in bytes.
|
|
// channelBufferSize: buffer size for client write channels in messages.
|
|
func New(pingPeriod, pongTimeout time.Duration, allowedWebSocketOrigins []string, shardCount, readBufferSize, writeBufferSize, channelBufferSize int) *API {
|
|
// Ensure shardCount is at least 1 and is a power of 2 for optimal hashing
|
|
if shardCount < 1 {
|
|
shardCount = 256
|
|
}
|
|
// Round up to next power of 2
|
|
shardCount = nextPowerOf2(shardCount)
|
|
|
|
shards := make([]*shard, shardCount)
|
|
for i := range shards {
|
|
shards[i] = &shard{
|
|
clients: make(map[uint]map[string][]*client),
|
|
}
|
|
}
|
|
|
|
return &API{
|
|
shards: shards,
|
|
shardCount: shardCount,
|
|
pingPeriod: pingPeriod,
|
|
pongTimeout: pingPeriod + pongTimeout,
|
|
upgrader: newUpgrader(allowedWebSocketOrigins, readBufferSize, writeBufferSize),
|
|
channelBuf: channelBufferSize,
|
|
}
|
|
}
|
|
|
|
// nextPowerOf2 returns the next power of 2 >= n.
|
|
func nextPowerOf2(n int) int {
|
|
if n <= 0 {
|
|
return 1
|
|
}
|
|
n--
|
|
n |= n >> 1
|
|
n |= n >> 2
|
|
n |= n >> 4
|
|
n |= n >> 8
|
|
n |= n >> 16
|
|
n++
|
|
return n
|
|
}
|
|
|
|
// getShard returns the shard for a given userID using fast modulo operation.
|
|
func (a *API) getShard(userID uint) *shard {
|
|
return a.shards[userID%uint(a.shardCount)]
|
|
}
|
|
|
|
// CollectConnectedClientTokens returns all tokens of the connected clients.
|
|
func (a *API) CollectConnectedClientTokens() []string {
|
|
var allClients []string
|
|
for _, shard := range a.shards {
|
|
shard.lock.RLock()
|
|
for _, userClients := range shard.clients {
|
|
for _, tokenClients := range userClients {
|
|
for _, c := range tokenClients {
|
|
allClients = append(allClients, c.token)
|
|
}
|
|
}
|
|
}
|
|
shard.lock.RUnlock()
|
|
}
|
|
return uniq(allClients)
|
|
}
|
|
|
|
// NotifyDeletedUser closes existing connections for the given user.
|
|
func (a *API) NotifyDeletedUser(userID uint) error {
|
|
shard := a.getShard(userID)
|
|
shard.lock.Lock()
|
|
defer shard.lock.Unlock()
|
|
if userClients, ok := shard.clients[userID]; ok {
|
|
for _, tokenClients := range userClients {
|
|
for _, client := range tokenClients {
|
|
client.Close()
|
|
}
|
|
}
|
|
delete(shard.clients, userID)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// NotifyDeletedClient closes existing connections with the given token.
|
|
func (a *API) NotifyDeletedClient(userID uint, token string) {
|
|
shard := a.getShard(userID)
|
|
shard.lock.Lock()
|
|
defer shard.lock.Unlock()
|
|
if userClients, ok := shard.clients[userID]; ok {
|
|
if tokenClients, exists := userClients[token]; exists {
|
|
for _, client := range tokenClients {
|
|
client.Close()
|
|
}
|
|
delete(userClients, token)
|
|
// Clean up empty user map
|
|
if len(userClients) == 0 {
|
|
delete(shard.clients, userID)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Notify notifies the clients with the given userID that a new messages was created.
|
|
func (a *API) Notify(userID uint, msg *model.MessageExternal) {
|
|
shard := a.getShard(userID)
|
|
shard.lock.RLock()
|
|
userClients, ok := shard.clients[userID]
|
|
if !ok {
|
|
shard.lock.RUnlock()
|
|
return
|
|
}
|
|
// Create a snapshot of clients to avoid holding lock during send
|
|
clients := make([]*client, 0)
|
|
for _, tokenClients := range userClients {
|
|
for _, c := range tokenClients {
|
|
clients = append(clients, c)
|
|
}
|
|
}
|
|
shard.lock.RUnlock()
|
|
|
|
// Send messages without holding the lock to prevent blocking other shards
|
|
// The channel buffer (default 10) helps prevent blocking in most cases
|
|
for _, c := range clients {
|
|
c.write <- msg
|
|
}
|
|
}
|
|
|
|
func (a *API) remove(c *client) {
|
|
shard := a.getShard(c.userID)
|
|
shard.lock.Lock()
|
|
defer shard.lock.Unlock()
|
|
if userClients, ok := shard.clients[c.userID]; ok {
|
|
if tokenClients, exists := userClients[c.token]; exists {
|
|
// Remove the specific client from the slice
|
|
for i, client := range tokenClients {
|
|
if client == c {
|
|
userClients[c.token] = append(tokenClients[:i], tokenClients[i+1:]...)
|
|
// Clean up empty token slice
|
|
if len(userClients[c.token]) == 0 {
|
|
delete(userClients, c.token)
|
|
}
|
|
// Clean up empty user map
|
|
if len(userClients) == 0 {
|
|
delete(shard.clients, c.userID)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (a *API) register(c *client) {
|
|
shard := a.getShard(c.userID)
|
|
shard.lock.Lock()
|
|
defer shard.lock.Unlock()
|
|
if shard.clients[c.userID] == nil {
|
|
shard.clients[c.userID] = make(map[string][]*client)
|
|
}
|
|
if shard.clients[c.userID][c.token] == nil {
|
|
shard.clients[c.userID][c.token] = make([]*client, 0, 1)
|
|
}
|
|
shard.clients[c.userID][c.token] = append(shard.clients[c.userID][c.token], c)
|
|
}
|
|
|
|
// Handle handles incoming requests. First it upgrades the protocol to the WebSocket protocol and then starts listening
|
|
// for read and writes.
|
|
// swagger:operation GET /stream message streamMessages
|
|
//
|
|
// Websocket, return newly created messages.
|
|
//
|
|
// ---
|
|
// schema: ws, wss
|
|
// produces: [application/json]
|
|
// security: [clientTokenAuthorizationHeader: [], clientTokenHeader: [], clientTokenQuery: [], basicAuth: []]
|
|
// responses:
|
|
// 200:
|
|
// description: Ok
|
|
// schema:
|
|
// $ref: "#/definitions/Message"
|
|
// 400:
|
|
// description: Bad Request
|
|
// schema:
|
|
// $ref: "#/definitions/Error"
|
|
// 401:
|
|
// description: Unauthorized
|
|
// schema:
|
|
// $ref: "#/definitions/Error"
|
|
// 403:
|
|
// description: Forbidden
|
|
// schema:
|
|
// $ref: "#/definitions/Error"
|
|
// 500:
|
|
// description: Server Error
|
|
// schema:
|
|
// $ref: "#/definitions/Error"
|
|
func (a *API) Handle(ctx *gin.Context) {
|
|
conn, err := a.upgrader.Upgrade(ctx.Writer, ctx.Request, nil)
|
|
if err != nil {
|
|
ctx.Error(err)
|
|
return
|
|
}
|
|
|
|
client := newClient(conn, auth.GetUserID(ctx), auth.GetTokenID(ctx), a.remove, a.channelBuf)
|
|
a.register(client)
|
|
go client.startReading(a.pongTimeout)
|
|
go client.startWriteHandler(a.pingPeriod)
|
|
}
|
|
|
|
// Close closes all client connections and stops answering new connections.
|
|
func (a *API) Close() {
|
|
for _, shard := range a.shards {
|
|
shard.lock.Lock()
|
|
for _, userClients := range shard.clients {
|
|
for _, tokenClients := range userClients {
|
|
for _, client := range tokenClients {
|
|
client.Close()
|
|
}
|
|
}
|
|
}
|
|
for k := range shard.clients {
|
|
delete(shard.clients, k)
|
|
}
|
|
shard.lock.Unlock()
|
|
}
|
|
}
|
|
|
|
func uniq[T comparable](s []T) []T {
|
|
m := make(map[T]struct{}, len(s))
|
|
r := make([]T, 0, len(s))
|
|
for _, v := range s {
|
|
if _, ok := m[v]; !ok {
|
|
m[v] = struct{}{}
|
|
r = append(r, v)
|
|
}
|
|
}
|
|
return r
|
|
}
|
|
|
|
func isAllowedOrigin(r *http.Request, allowedOrigins []*regexp.Regexp) bool {
|
|
origin := r.Header.Get("origin")
|
|
if origin == "" {
|
|
return true
|
|
}
|
|
|
|
u, err := url.Parse(origin)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
|
|
if strings.EqualFold(u.Host, r.Host) {
|
|
return true
|
|
}
|
|
|
|
for _, allowedOrigin := range allowedOrigins {
|
|
if allowedOrigin.MatchString(strings.ToLower(u.Hostname())) {
|
|
return true
|
|
}
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
func newUpgrader(allowedWebSocketOrigins []string, readBufferSize, writeBufferSize int) *websocket.Upgrader {
|
|
compiledAllowedOrigins := compileAllowedWebSocketOrigins(allowedWebSocketOrigins)
|
|
return &websocket.Upgrader{
|
|
ReadBufferSize: readBufferSize,
|
|
WriteBufferSize: writeBufferSize,
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
if mode.IsDev() {
|
|
return true
|
|
}
|
|
return isAllowedOrigin(r, compiledAllowedOrigins)
|
|
},
|
|
}
|
|
}
|
|
|
|
func compileAllowedWebSocketOrigins(allowedOrigins []string) []*regexp.Regexp {
|
|
var compiledAllowedOrigins []*regexp.Regexp
|
|
for _, origin := range allowedOrigins {
|
|
compiledAllowedOrigins = append(compiledAllowedOrigins, regexp.MustCompile(origin))
|
|
}
|
|
|
|
return compiledAllowedOrigins
|
|
}
|