go-omni/examples/websocketRPC/main.go

// websocketRPC demonstrates bidirectional calls and explicit subscription
// restoration over two successive real TCP connections, without external services.
package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/ws"
	"git.yuntuo.hk/k2safe/go-omni/ws/hub"
	"git.yuntuo.hk/k2safe/go-omni/ws/rpc"
)

func main() {
	if err := run(); err != nil {
		panic(err)
	}
}

func run() (runErr error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	groups, err := hub.New(hub.Config{})
	if err != nil {
		return err
	}
	defer groups.Close()
	server, err := ws.NewServer(ws.ServerConfig{
		Context: ctx, Subprotocols: []string{rpc.Subprotocol}, RequireSubprotocol: true,
		Prepare: func(context.Context, *http.Request) (ws.Handlers, error) {
			peer, err := rpc.New(rpc.Config{})
			if err != nil {
				return ws.Handlers{}, err
			}
			var session *ws.Session
			if err = peer.Handle("double", func(_ context.Context, raw json.RawMessage) (any, error) {
				var n int
				if json.Unmarshal(raw, &n) != nil {
					return nil, &rpc.Error{Code: "INVALID_INPUT", Message: "integer required"}
				}
				return n * 2, nil
			}); err != nil {
				return ws.Handlers{}, err
			}
			if err = peer.Handle("identify", func(ctx context.Context, _ json.RawMessage) (any, error) {
				var name string
				err := peer.Call(ctx, "client.name", nil, &name)
				return name, err
			}); err != nil {
				return ws.Handlers{}, err
			}
			if err = peer.Handle("rpc.subscribe", func(_ context.Context, raw json.RawMessage) (any, error) {
				var sub rpc.Subscription
				if json.Unmarshal(raw, &sub) != nil {
					return nil, &rpc.Error{Code: "INVALID_INPUT", Message: "invalid subscription"}
				}
				// Demo policy. A real project authenticates in Prepare, then authorizes the
				// requested topic against that verified identity on every subscription.
				if sub.Key != "updates" {
					return nil, &rpc.Error{Code: "DENIED", Message: "topic not allowed"}
				}
				return nil, groups.Join(session.ID(), sub.Key)
			}); err != nil {
				return ws.Handlers{}, err
			}
			if err = peer.Handle("rpc.unsubscribe", func(_ context.Context, raw json.RawMessage) (any, error) {
				var sub struct {
					Key string `json:"key"`
				}
				if json.Unmarshal(raw, &sub) != nil {
					return nil, &rpc.Error{Code: "INVALID_INPUT", Message: "invalid subscription"}
				}
				groups.Leave(session.ID(), sub.Key)
				return nil, nil
			}); err != nil {
				return ws.Handlers{}, err
			}
			handlers := peer.Handlers()
			open := handlers.OnOpen
			handlers.OnOpen = func(ctx context.Context, s *ws.Session) error {
				session = s
				if err := open(ctx, s); err != nil {
					return err
				}
				return groups.Add(s)
			}
			closeHandler := handlers.OnClose
			handlers.OnClose = func(s *ws.Session, err error) { closeHandler(s, err); groups.Remove(s.ID()) }
			return handlers, nil
		},
	})
	if err != nil {
		return err
	}
	httpServer := httptest.NewServer(server.Handler(ws.Handlers{}))
	defer httpServer.Close()
	defer func() {
		stop, done := context.WithTimeout(context.Background(), 2*time.Second)
		defer done()
		runErr = errors.Join(runErr, server.Shutdown(stop))
	}()
	client, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct", Subprotocols: []string{rpc.Subprotocol}, RequireSubprotocol: true})
	if err != nil {
		return err
	}
	defer func() {
		stop, done := context.WithTimeout(context.Background(), 2*time.Second)
		defer done()
		runErr = errors.Join(runErr, client.Shutdown(stop))
	}()
	subscriptions, err := rpc.NewSubscriptions(rpc.SubscriptionConfig{})
	if err != nil {
		return err
	}
	if err = subscriptions.Subscribe(ctx, nil, "updates", nil); err != nil {
		return err
	}
	endpoint := "ws" + strings.TrimPrefix(httpServer.URL, "http")
	for generation := 1; generation <= 2; generation++ {
		peer, err := rpc.New(rpc.Config{})
		if err != nil {
			return err
		}
		events := make(chan int, 1)
		if err = peer.Handle("client.name", func(context.Context, json.RawMessage) (any, error) { return "generic Go client", nil }); err != nil {
			return err
		}
		if err = peer.Handle("updates.changed", func(ctx context.Context, raw json.RawMessage) (any, error) {
			var v int
			if err := json.Unmarshal(raw, &v); err != nil {
				return nil, err
			}
			select {
			case events <- v:
				return nil, nil
			case <-ctx.Done():
				return nil, ctx.Err()
			}
		}); err != nil {
			return err
		}
		session, err := client.Dial(ctx, endpoint, peer.Handlers())
		if err != nil {
			return err
		}
		if err = peer.WaitReady(ctx); err != nil {
			return err
		}
		var doubled int
		if err = peer.Call(ctx, "double", 21, &doubled); err != nil {
			return err
		}
		var name string
		if err = peer.Call(ctx, "identify", nil, &name); err != nil {
			return err
		}
		restored, err := subscriptions.Restore(ctx, peer)
		if err != nil {
			return err
		}
		// The documented envelope also permits native browser clients and Hub fanout.
		payload, err := json.Marshal(map[string]any{"v": 1, "type": "event", "method": "updates.changed", "params": generation})
		if err != nil {
			return err
		}
		sent, err := groups.SendGroup("updates", ws.Message{Type: ws.Text, Data: payload})
		if err != nil {
			return err
		}
		if sent.Queued != 1 || len(sent.Errors) > 0 {
			return fmt.Errorf("broadcast failed: queued %d of %d", sent.Queued, sent.Attempted)
		}
		select {
		case received := <-events:
			fmt.Printf("connection %d: double=%d, remote callback=%q, restored=%d, event=%d\n", generation, doubled, name, restored, received)
		case <-ctx.Done():
			return ctx.Err()
		}
		if err = subscriptions.Unsubscribe(ctx, peer, "updates"); err != nil {
			return err
		}
		// Keep desired state separately for the next physical connection.
		if err = subscriptions.Subscribe(ctx, nil, "updates", nil); err != nil {
			return err
		}
		if err = session.Close(1000, ""); err != nil {
			return err
		}
		if err = session.Wait(ctx); err != nil {
			return err
		}
	}
	return nil
}