feat(ws): release generic WebSocket client and server v0.31.0-preview

ws/rpc/subscriptions.go

@@ -0,0 +1,159 @@
+package rpc
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"sort"
+	"strings"
+	"unicode/utf8"
+)
+
+// Subscription is the generic request sent to the configured subscribe method.
+// The server must authenticate and authorize Key and Params itself.
+type Subscription struct {
+	Key    string          `json:"key"`
+	Params json.RawMessage `json:"params"`
+}
+
+type SubscriptionConfig struct {
+	MaxSubscriptions  int    `json:"maxSubscriptions"`
+	MaxParamsBytes    int    `json:"maxParamsBytes"`
+	SubscribeMethod   string `json:"subscribeMethod"`
+	UnsubscribeMethod string `json:"unsubscribeMethod"`
+}
+
+// Subscriptions stores desired subscriptions, never received messages. Reuse it
+// across sessions and explicitly Restore after each new Peer is bound. Server
+// subscribe handlers must be idempotent for a key within a session; authorization
+// is checked anew on every restore. No missed-event replay is implied.
+type Subscriptions struct {
+	cfg     SubscriptionConfig
+	gate    chan struct{}
+	desired map[string]json.RawMessage
+}
+
+func validSubscriptionKey(key string) bool {
+	return key != "" && len(key) <= 256 && utf8.ValidString(key) && !strings.ContainsAny(key, "\r\n\x00")
+}
+
+func NewSubscriptions(cfg SubscriptionConfig) (*Subscriptions, error) {
+	if cfg.MaxSubscriptions < 0 || cfg.MaxSubscriptions > 65536 || cfg.MaxParamsBytes < 0 || cfg.MaxParamsBytes > 1<<20 {
+		return nil, errors.New("rpc: invalid subscription limits")
+	}
+	if cfg.MaxSubscriptions == 0 {
+		cfg.MaxSubscriptions = 256
+	}
+	if cfg.MaxParamsBytes == 0 {
+		cfg.MaxParamsBytes = 16 << 10
+	}
+	if cfg.SubscribeMethod == "" {
+		cfg.SubscribeMethod = "rpc.subscribe"
+	}
+	if cfg.UnsubscribeMethod == "" {
+		cfg.UnsubscribeMethod = "rpc.unsubscribe"
+	}
+	if !validName(cfg.SubscribeMethod, 128) || !validName(cfg.UnsubscribeMethod, 128) || cfg.SubscribeMethod == cfg.UnsubscribeMethod {
+		return nil, errors.New("rpc: distinct subscription methods required")
+	}
+	return &Subscriptions{cfg: cfg, gate: make(chan struct{}, 1), desired: make(map[string]json.RawMessage)}, nil
+}
+func (s *Subscriptions) lock(ctx context.Context) error {
+	select {
+	case s.gate <- struct{}{}:
+		if err := ctx.Err(); err != nil {
+			<-s.gate
+			return err
+		}
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+func (s *Subscriptions) unlock() { <-s.gate }
+
+// Subscribe records the desired state before attempting the call. On failure,
+// it remains desired for a later Restore; the returned error must be inspected.
+// A nil peer stores the subscription for the next explicit Restore.
+func (s *Subscriptions) Subscribe(ctx context.Context, p *Peer, key string, params any) error {
+	if !validSubscriptionKey(key) {
+		return errors.New("rpc: invalid subscription key")
+	}
+	raw, err := json.Marshal(params)
+	if err != nil {
+		return err
+	}
+	if len(raw) > s.cfg.MaxParamsBytes {
+		return errors.New("rpc: subscription params exceed limit")
+	}
+	if err = s.lock(ctx); err != nil {
+		return err
+	}
+	defer s.unlock()
+	if _, exists := s.desired[key]; !exists && len(s.desired) >= s.cfg.MaxSubscriptions {
+		return ErrBusy
+	}
+	s.desired[key] = raw
+	if p == nil {
+		return nil
+	}
+	return p.Call(ctx, s.cfg.SubscribeMethod, Subscription{key, raw}, nil)
+}
+
+// Unsubscribe removes desired state even if the remote acknowledgement fails.
+// The current connection may still receive events until remote removal or close.
+func (s *Subscriptions) Unsubscribe(ctx context.Context, p *Peer, key string) error {
+	if !validSubscriptionKey(key) {
+		return errors.New("rpc: invalid subscription key")
+	}
+	if err := s.lock(ctx); err != nil {
+		return err
+	}
+	defer s.unlock()
+	delete(s.desired, key)
+	if p == nil {
+		return nil
+	}
+	return p.Call(ctx, s.cfg.UnsubscribeMethod, struct {
+		Key string `json:"key"`
+	}{key}, nil)
+}
+
+// Restore reauthorizes all desired subscriptions in deterministic order. It
+// stops at the first error and reports how many were acknowledged. A caller can
+// retry on the same peer only if its subscription handler is idempotent.
+func (s *Subscriptions) Restore(ctx context.Context, p *Peer) (int, error) {
+	if p == nil {
+		return 0, ErrNotConnected
+	}
+	if err := s.lock(ctx); err != nil {
+		return 0, err
+	}
+	defer s.unlock()
+	keys := make([]string, 0, len(s.desired))
+	for key := range s.desired {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+	for i, key := range keys {
+		if err := p.Call(ctx, s.cfg.SubscribeMethod, Subscription{key, s.desired[key]}, nil); err != nil {
+			return i, err
+		}
+	}
+	return len(keys), nil
+}
+
+// Keys returns a sorted snapshot. In-flight restoration serializes desired
+// changes; waiting for that operation always respects ctx.
+func (s *Subscriptions) Keys(ctx context.Context) ([]string, error) {
+	if err := s.lock(ctx); err != nil {
+		return nil, err
+	}
+	defer s.unlock()
+	keys := make([]string, 0, len(s.desired))
+	for key := range s.desired {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+	return keys, nil
+}