This commit is contained in:
ITWorx Pulse release export
@@ -0,0 +1,511 @@
|
||||
package live
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/coder/websocket"
|
||||
"github.com/itworx/pulse/internal/auth"
|
||||
"github.com/itworx/pulse/internal/correlation"
|
||||
"github.com/itworx/pulse/internal/queryplan"
|
||||
)
|
||||
|
||||
const (
|
||||
schemaVersion = 1
|
||||
defaultMaxMessage = 64 << 10
|
||||
defaultMaxSubs = 50
|
||||
defaultRateWindow = time.Minute
|
||||
defaultRateMessages = 120
|
||||
defaultHeartbeat = 15 * time.Second
|
||||
defaultIdleTimeout = 45 * time.Second
|
||||
defaultWriteTimeout = 5 * time.Second
|
||||
maxSubscriptionID = 128
|
||||
maxDetailLength = 500
|
||||
defaultOutboundMessages = 32
|
||||
)
|
||||
|
||||
var (
|
||||
ErrMessageTooLarge = errors.New("live message exceeds size limit")
|
||||
ErrRateLimited = errors.New("live message rate limit exceeded")
|
||||
)
|
||||
|
||||
type Incoming struct {
|
||||
SchemaVersion int `json:"schemaVersion"`
|
||||
Type string `json:"type"`
|
||||
SubscriptionID string `json:"subscriptionId,omitempty"`
|
||||
Query json.RawMessage `json:"query,omitempty"`
|
||||
IntervalSeconds int `json:"intervalSeconds,omitempty"`
|
||||
Nonce string `json:"nonce,omitempty"`
|
||||
}
|
||||
|
||||
type Sample struct {
|
||||
Series string `json:"series"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Value *float64 `json:"value"`
|
||||
Freshness string `json:"freshness"`
|
||||
Labels map[string]string `json:"labels,omitempty"`
|
||||
}
|
||||
|
||||
type SamplesMessage struct {
|
||||
SchemaVersion int `json:"schemaVersion"`
|
||||
Type string `json:"type"`
|
||||
SubscriptionID string `json:"subscriptionId"`
|
||||
Sequence uint64 `json:"sequence"`
|
||||
ServerTime time.Time `json:"serverTime"`
|
||||
Samples []Sample `json:"samples"`
|
||||
Coalesced bool `json:"coalesced,omitempty"`
|
||||
}
|
||||
|
||||
type StatusMessage struct {
|
||||
SchemaVersion int `json:"schemaVersion"`
|
||||
Type string `json:"type"`
|
||||
SubscriptionID string `json:"subscriptionId"`
|
||||
State string `json:"state"`
|
||||
Detail string `json:"detail,omitempty"`
|
||||
}
|
||||
|
||||
type ErrorMessage struct {
|
||||
SchemaVersion int `json:"schemaVersion"`
|
||||
Type string `json:"type"`
|
||||
SubscriptionID string `json:"subscriptionId,omitempty"`
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
CorrelationID string `json:"correlationId"`
|
||||
}
|
||||
|
||||
type PongMessage struct {
|
||||
SchemaVersion int `json:"schemaVersion"`
|
||||
Type string `json:"type"`
|
||||
Nonce string `json:"nonce"`
|
||||
}
|
||||
|
||||
type Sampler interface {
|
||||
Sample(context.Context, queryplan.Request) ([]Sample, error)
|
||||
}
|
||||
|
||||
type Handler struct {
|
||||
Planner *queryplan.Planner
|
||||
Sampler Sampler
|
||||
AllowedOrigins []string
|
||||
MaxSubscriptions int
|
||||
MaxMessages int
|
||||
RateWindow time.Duration
|
||||
HeartbeatInterval time.Duration
|
||||
IdleTimeout time.Duration
|
||||
WriteTimeout time.Duration
|
||||
Now func() time.Time
|
||||
Registry *Registry
|
||||
}
|
||||
|
||||
type subscription struct {
|
||||
id string
|
||||
request queryplan.Request
|
||||
interval time.Duration
|
||||
nextAt time.Time
|
||||
sequence uint64
|
||||
lease *Lease
|
||||
}
|
||||
|
||||
type session struct {
|
||||
conn *websocket.Conn
|
||||
principal auth.Principal
|
||||
planner *queryplan.Planner
|
||||
sampler Sampler
|
||||
maxSubs int
|
||||
maxMessages int
|
||||
rateWindow time.Duration
|
||||
heartbeat time.Duration
|
||||
idleTimeout time.Duration
|
||||
writeTimeout time.Duration
|
||||
registry *Registry
|
||||
now func() time.Time
|
||||
correlationID string
|
||||
mu sync.Mutex
|
||||
subscriptions map[string]*subscription
|
||||
rateStarted time.Time
|
||||
rateCount int
|
||||
outbound chan []byte
|
||||
}
|
||||
|
||||
func (h Handler) ServeHTTP(response http.ResponseWriter, request *http.Request) {
|
||||
principal, ok := auth.PrincipalFromContext(request.Context())
|
||||
if !ok {
|
||||
http.Error(response, "unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
if !auth.Allows(principal.Role, auth.PermissionView) {
|
||||
http.Error(response, "forbidden", http.StatusForbidden)
|
||||
return
|
||||
}
|
||||
conn, err := websocket.Accept(response, request, &websocket.AcceptOptions{OriginPatterns: h.AllowedOrigins})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
conn.SetReadLimit(int64(defaultMaxMessage))
|
||||
registry := h.Registry
|
||||
if registry == nil {
|
||||
registry = NewRegistry(h.Sampler, RegistryOptions{})
|
||||
}
|
||||
// Keep the upgraded transport attached to the authenticated request. The
|
||||
// session middleware cancels this context on logout and at the finite
|
||||
// absolute session deadline; server/request cancellation must also release
|
||||
// subscriptions and the socket.
|
||||
ctx, cancel := context.WithCancel(request.Context())
|
||||
correlationID := correlation.FromContext(request.Context())
|
||||
if correlationID == "" {
|
||||
correlationID = correlation.New()
|
||||
}
|
||||
s := &session{
|
||||
conn: conn, principal: principal, planner: h.Planner, sampler: h.Sampler,
|
||||
maxSubs: positiveOr(h.MaxSubscriptions, defaultMaxSubs),
|
||||
maxMessages: positiveOr(h.MaxMessages, defaultRateMessages),
|
||||
rateWindow: durationOr(h.RateWindow, defaultRateWindow),
|
||||
heartbeat: durationOr(h.HeartbeatInterval, defaultHeartbeat),
|
||||
idleTimeout: durationOr(h.IdleTimeout, defaultIdleTimeout),
|
||||
writeTimeout: durationOr(h.WriteTimeout, defaultWriteTimeout),
|
||||
registry: registry,
|
||||
now: h.Now, correlationID: correlationID, subscriptions: make(map[string]*subscription), outbound: make(chan []byte, defaultOutboundMessages),
|
||||
}
|
||||
if s.now == nil {
|
||||
s.now = func() time.Time { return time.Now().UTC() }
|
||||
}
|
||||
defer func() {
|
||||
cancel()
|
||||
s.releaseAll()
|
||||
conn.CloseNow()
|
||||
}()
|
||||
go s.writeLoop(ctx, cancel)
|
||||
go s.heartbeatLoop(ctx, cancel)
|
||||
go s.sampleLoop(ctx, cancel)
|
||||
s.readLoop(ctx)
|
||||
}
|
||||
|
||||
func positiveOr(value, fallback int) int {
|
||||
if value > 0 {
|
||||
return value
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func durationOr(value, fallback time.Duration) time.Duration {
|
||||
if value > 0 {
|
||||
return value
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func (s *session) readLoop(ctx context.Context) {
|
||||
for {
|
||||
readCtx, cancel := context.WithTimeout(ctx, s.idleTimeout)
|
||||
_, payload, err := s.conn.Read(readCtx)
|
||||
cancel()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if err := s.acceptMessage(ctx, payload); err != nil {
|
||||
code, detail := "LIVE_MESSAGE_INVALID", "Live bericht is ongeldig."
|
||||
if errors.Is(err, ErrMessageTooLarge) {
|
||||
code, detail = "LIVE_MESSAGE_LIMIT", "Live bericht overschrijdt de maximale berichtgrootte."
|
||||
} else if errors.Is(err, ErrRateLimited) {
|
||||
code, detail = "LIVE_RATE_LIMIT", "Te veel live berichten; verbind opnieuw na een korte pauze."
|
||||
}
|
||||
_ = s.send(ErrorMessage{SchemaVersion: schemaVersion, Type: "error", Code: code, Message: detail, CorrelationID: s.correlationID})
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *session) acceptMessage(ctx context.Context, payload []byte) error {
|
||||
if len(payload) > defaultMaxMessage {
|
||||
return ErrMessageTooLarge
|
||||
}
|
||||
s.mu.Lock()
|
||||
now := s.now()
|
||||
if s.rateStarted.IsZero() || now.Sub(s.rateStarted) >= s.rateWindow {
|
||||
s.rateStarted = now
|
||||
s.rateCount = 0
|
||||
}
|
||||
s.rateCount++
|
||||
limited := s.rateCount > s.maxMessages
|
||||
s.mu.Unlock()
|
||||
if limited {
|
||||
return ErrRateLimited
|
||||
}
|
||||
var message Incoming
|
||||
decoder := json.NewDecoder(bytes.NewReader(payload))
|
||||
decoder.DisallowUnknownFields()
|
||||
if err := decoder.Decode(&message); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := decoder.Decode(&struct{}{}); err != io.EOF {
|
||||
return errors.New("live message has trailing data")
|
||||
}
|
||||
if message.SchemaVersion != schemaVersion || !validTypeFields(payload, message.Type) {
|
||||
return errors.New("live message schema is invalid")
|
||||
}
|
||||
switch message.Type {
|
||||
case "subscribe":
|
||||
return s.subscribe(ctx, message)
|
||||
case "unsubscribe":
|
||||
return s.unsubscribe(message.SubscriptionID)
|
||||
case "ping":
|
||||
if !validNonce(message.Nonce) {
|
||||
return errors.New("ping nonce is invalid")
|
||||
}
|
||||
return s.send(PongMessage{SchemaVersion: schemaVersion, Type: "pong", Nonce: message.Nonce})
|
||||
case "pong":
|
||||
if !validNonce(message.Nonce) {
|
||||
return errors.New("pong nonce is invalid")
|
||||
}
|
||||
return nil
|
||||
default:
|
||||
return errors.New("live message type is unsupported")
|
||||
}
|
||||
}
|
||||
|
||||
func validTypeFields(payload []byte, typ string) bool {
|
||||
var fields map[string]json.RawMessage
|
||||
if json.Unmarshal(payload, &fields) != nil {
|
||||
return false
|
||||
}
|
||||
allowed := map[string]map[string]struct{}{
|
||||
"subscribe": {"schemaVersion": {}, "type": {}, "subscriptionId": {}, "query": {}, "intervalSeconds": {}},
|
||||
"unsubscribe": {"schemaVersion": {}, "type": {}, "subscriptionId": {}},
|
||||
"ping": {"schemaVersion": {}, "type": {}, "nonce": {}},
|
||||
"pong": {"schemaVersion": {}, "type": {}, "nonce": {}},
|
||||
}
|
||||
known, ok := allowed[typ]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
for key := range fields {
|
||||
if _, ok := known[key]; !ok {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func validSubscriptionID(id string) bool {
|
||||
return id != "" && utf8.RuneCountInString(id) <= maxSubscriptionID
|
||||
}
|
||||
|
||||
func validNonce(nonce string) bool {
|
||||
return utf8.RuneCountInString(nonce) <= maxSubscriptionID
|
||||
}
|
||||
|
||||
func (s *session) subscribe(ctx context.Context, message Incoming) error {
|
||||
if !validSubscriptionID(message.SubscriptionID) || len(message.Query) == 0 || message.IntervalSeconds < 1 || message.IntervalSeconds > 300 {
|
||||
return s.sendError(message.SubscriptionID, "LIVE_SUBSCRIPTION_INVALID", "Abonnement, query of interval is ongeldig.")
|
||||
}
|
||||
var fields map[string]json.RawMessage
|
||||
if err := json.Unmarshal(message.Query, &fields); err != nil || len(fields) > 20 {
|
||||
return s.sendError(message.SubscriptionID, "LIVE_QUERY_INVALID", "Live query moet een begrensd JSON-object zijn.")
|
||||
}
|
||||
var request queryplan.Request
|
||||
decoder := json.NewDecoder(bytes.NewReader(message.Query))
|
||||
decoder.DisallowUnknownFields()
|
||||
if err := decoder.Decode(&request); err != nil {
|
||||
return s.sendError(message.SubscriptionID, "LIVE_QUERY_INVALID", "Live query bevat onbekende of ongeldige velden.")
|
||||
}
|
||||
if s.planner == nil {
|
||||
return s.sendError(message.SubscriptionID, "LIVE_QUERY_UNAVAILABLE", "Live queryvalidatie is tijdelijk niet beschikbaar.")
|
||||
}
|
||||
plan, err := s.planner.Plan(auth.WithPrincipal(ctx, s.principal), request)
|
||||
if err != nil {
|
||||
return s.sendError(message.SubscriptionID, "LIVE_SUBSCRIPTION_DENIED", queryErrorDetail(err))
|
||||
}
|
||||
now := s.now()
|
||||
interval := time.Duration(message.IntervalSeconds) * time.Second
|
||||
s.mu.Lock()
|
||||
duplicate := false
|
||||
sessionLimit := len(s.subscriptions) >= s.maxSubs
|
||||
if _, exists := s.subscriptions[message.SubscriptionID]; exists {
|
||||
duplicate = true
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if duplicate {
|
||||
return s.sendError(message.SubscriptionID, "LIVE_SUBSCRIPTION_CONFLICT", "Subscription-ID bestaat al.")
|
||||
}
|
||||
if sessionLimit {
|
||||
return s.sendError(message.SubscriptionID, "LIVE_SUBSCRIPTION_LIMIT", "Maximum aantal live abonnementen bereikt.")
|
||||
}
|
||||
lease, err := s.registry.Acquire(plan.Request, interval)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrRegistryLimit) {
|
||||
return s.sendError(message.SubscriptionID, "LIVE_REGISTRY_LIMIT", "Maximum aantal gedeelde live queries bereikt.")
|
||||
}
|
||||
return s.sendError(message.SubscriptionID, "LIVE_QUERY_INVALID", "Live query kon niet worden genormaliseerd.")
|
||||
}
|
||||
s.mu.Lock()
|
||||
if _, exists := s.subscriptions[message.SubscriptionID]; exists || len(s.subscriptions) >= s.maxSubs {
|
||||
s.mu.Unlock()
|
||||
lease.Release()
|
||||
return s.sendError(message.SubscriptionID, "LIVE_SUBSCRIPTION_CONFLICT", "Subscription-ID bestaat al.")
|
||||
}
|
||||
s.subscriptions[message.SubscriptionID] = &subscription{id: message.SubscriptionID, request: plan.Request, interval: interval, nextAt: now.Add(interval), lease: lease}
|
||||
s.mu.Unlock()
|
||||
return s.send(StatusMessage{SchemaVersion: schemaVersion, Type: "status", SubscriptionID: message.SubscriptionID, State: "subscribed"})
|
||||
}
|
||||
|
||||
func queryErrorDetail(err error) string {
|
||||
var plannerError queryplan.Error
|
||||
if errors.As(err, &plannerError) {
|
||||
return plannerError.Detail
|
||||
}
|
||||
return "Live query kon niet worden geautoriseerd."
|
||||
}
|
||||
|
||||
func (s *session) unsubscribe(id string) error {
|
||||
if !validSubscriptionID(id) {
|
||||
return s.sendError(id, "LIVE_SUBSCRIPTION_INVALID", "Subscription-ID is ongeldig.")
|
||||
}
|
||||
s.mu.Lock()
|
||||
sub, exists := s.subscriptions[id]
|
||||
delete(s.subscriptions, id)
|
||||
s.mu.Unlock()
|
||||
if exists && sub.lease != nil {
|
||||
sub.lease.Release()
|
||||
}
|
||||
return s.send(StatusMessage{SchemaVersion: schemaVersion, Type: "status", SubscriptionID: id, State: "unsubscribed"})
|
||||
}
|
||||
|
||||
func (s *session) releaseAll() {
|
||||
s.mu.Lock()
|
||||
subscriptions := make([]*subscription, 0, len(s.subscriptions))
|
||||
for _, sub := range s.subscriptions {
|
||||
subscriptions = append(subscriptions, sub)
|
||||
}
|
||||
s.subscriptions = make(map[string]*subscription)
|
||||
s.mu.Unlock()
|
||||
for _, sub := range subscriptions {
|
||||
if sub.lease != nil {
|
||||
sub.lease.Release()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *session) sendError(subscriptionID, code, message string) error {
|
||||
if len(message) > maxDetailLength {
|
||||
message = message[:maxDetailLength]
|
||||
}
|
||||
return s.send(ErrorMessage{SchemaVersion: schemaVersion, Type: "error", SubscriptionID: subscriptionID, Code: code, Message: message, CorrelationID: s.correlationID})
|
||||
}
|
||||
|
||||
func (s *session) send(message any) error {
|
||||
payload, err := json.Marshal(message)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(payload) > defaultMaxMessage {
|
||||
return ErrMessageTooLarge
|
||||
}
|
||||
sample := false
|
||||
switch message.(type) {
|
||||
case SamplesMessage, *SamplesMessage:
|
||||
sample = true
|
||||
}
|
||||
if !sample {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), s.writeTimeout)
|
||||
defer cancel()
|
||||
return s.conn.Write(ctx, websocket.MessageText, payload)
|
||||
}
|
||||
if s.outbound == nil {
|
||||
return ErrOutboundBackpressure
|
||||
}
|
||||
select {
|
||||
case s.outbound <- payload:
|
||||
return nil
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (s *session) writeLoop(ctx context.Context, cancel context.CancelFunc) {
|
||||
for {
|
||||
select {
|
||||
case payload := <-s.outbound:
|
||||
writeCtx, writeCancel := context.WithTimeout(ctx, s.writeTimeout)
|
||||
err := s.conn.Write(writeCtx, websocket.MessageText, payload)
|
||||
writeCancel()
|
||||
if err != nil {
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *session) heartbeatLoop(ctx context.Context, cancel context.CancelFunc) {
|
||||
ticker := time.NewTicker(s.heartbeat)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
pingCtx, pingCancel := context.WithTimeout(ctx, s.writeTimeout)
|
||||
err := s.conn.Ping(pingCtx)
|
||||
pingCancel()
|
||||
if err != nil {
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *session) sampleLoop(ctx context.Context, cancel context.CancelFunc) {
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
now := s.now()
|
||||
s.mu.Lock()
|
||||
due := make([]*subscription, 0, len(s.subscriptions))
|
||||
for _, sub := range s.subscriptions {
|
||||
if !now.Before(sub.nextAt) {
|
||||
sub.nextAt = now.Add(sub.interval)
|
||||
due = append(due, sub)
|
||||
}
|
||||
}
|
||||
s.mu.Unlock()
|
||||
for _, sub := range due {
|
||||
samples, err := s.registry.Sample(ctx, sub.request)
|
||||
if err != nil {
|
||||
_ = s.sendError(sub.id, "LIVE_SAMPLE_UNAVAILABLE", "Live sample is tijdelijk niet beschikbaar.")
|
||||
continue
|
||||
}
|
||||
s.mu.Lock()
|
||||
current, exists := s.subscriptions[sub.id]
|
||||
if !exists {
|
||||
s.mu.Unlock()
|
||||
continue
|
||||
}
|
||||
current.sequence++
|
||||
sequence := current.sequence
|
||||
s.mu.Unlock()
|
||||
if len(samples) > 10000 {
|
||||
_ = s.sendError(sub.id, "LIVE_SAMPLE_LIMIT", "Live sample bevat te veel punten.")
|
||||
continue
|
||||
}
|
||||
if err := s.send(SamplesMessage{SchemaVersion: schemaVersion, Type: "samples", SubscriptionID: sub.id, Sequence: sequence, ServerTime: now.UTC(), Samples: samples}); err != nil {
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user