go-omni/ws/rpc/subscriptions.go

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
}