diff --git a/mailbox/rpc/mux.go b/mailbox/rpc/mux.go new file mode 100644 index 000000000..1097b02e0 --- /dev/null +++ b/mailbox/rpc/mux.go @@ -0,0 +1,110 @@ +package mailboxrpc + +import ( + "context" + "errors" + "fmt" + "sync" + + "google.golang.org/protobuf/proto" +) + +// ErrNoHandler is returned when a (service, method) pair has no registered +// handler. +var ErrNoHandler = errors.New("no handler registered") + +// ServeMux is an in-process router that maps (service, method) pairs to typed +// handlers. +// +// ServeMux is intended as a small, dependency-free building block for mailbox +// RPC servers. It does not implement transport concerns such as authentication, +// retries, persistence, or acking. +type ServeMux struct { + mu sync.RWMutex + handlers map[routeKey]handlerEntry +} + +type routeKey struct { + service string + method string +} + +type handlerEntry struct { + newReq func() proto.Message + fn HandlerFunc +} + +// NewServeMux creates an empty mux. +func NewServeMux() *ServeMux { + return &ServeMux{ + handlers: make(map[routeKey]handlerEntry), + } +} + +// Handle registers a typed handler for a single RPC method. +func (m *ServeMux) Handle(service string, method string, + newReq func() proto.Message, fn HandlerFunc) { + + if service == "" { + panic("mailboxrpc: empty service name") + } + if method == "" { + panic("mailboxrpc: empty method name") + } + if newReq == nil { + panic("mailboxrpc: nil request constructor") + } + if fn == nil { + panic("mailboxrpc: nil handler function") + } + + m.mu.Lock() + defer m.mu.Unlock() + + m.handlers[routeKey{ + service: service, + method: method, + }] = handlerEntry{ + newReq: newReq, + fn: fn, + } +} + +// ServeRPC unmarshals reqBytes into the registered request type for +// (service, method) and invokes the handler. +func (m *ServeMux) ServeRPC(ctx context.Context, service string, + method string, reqBytes []byte) (proto.Message, error) { + + entry, ok := m.lookup(service, method) + if !ok { + return nil, fmt.Errorf("%w: %s/%s", ErrNoHandler, + service, method) + } + + req := entry.newReq() + if req == nil { + return nil, fmt.Errorf("nil request prototype for %s/%s", + service, method) + } + + if err := (proto.UnmarshalOptions{ + DiscardUnknown: true, + }).Unmarshal(reqBytes, req); err != nil { + return nil, err + } + + return entry.fn(ctx, req) +} + +// lookup returns the handler entry for (service, method) if present. +func (m *ServeMux) lookup(service string, method string) (handlerEntry, bool) { + m.mu.RLock() + defer m.mu.RUnlock() + + entry, ok := m.handlers[routeKey{ + service: service, + method: method, + }] + + return entry, ok +} diff --git a/mailbox/rpc/mux_test.go b/mailbox/rpc/mux_test.go new file mode 100644 index 000000000..45b4894ab --- /dev/null +++ b/mailbox/rpc/mux_test.go @@ -0,0 +1,79 @@ +package mailboxrpc_test + +import ( + "context" + "errors" + "testing" + + "github.com/lightninglabs/darepo-client/arkrpc" + mailboxrpc "github.com/lightninglabs/darepo-client/mailbox/rpc" + "google.golang.org/protobuf/proto" +) + +// TestServeMuxDispatch verifies a handler is invoked after unmarshalling. +func TestServeMuxDispatch(t *testing.T) { + t.Parallel() + + mux := mailboxrpc.NewServeMux() + mux.Handle(arkrpcServiceName, "GetInfo", func() proto.Message { + return new(arkrpc.GetInfoRequest) + }, func(_ context.Context, msg proto.Message) (proto.Message, error) { + _, ok := msg.(*arkrpc.GetInfoRequest) + if !ok { + t.Fatalf("unexpected request type: %T", msg) + } + + return &arkrpc.GetInfoResponse{ + Version: "test", + Pubkey: []byte{1, 2, 3}, + Network: "regtest", + BlockHeight: 7, + }, nil + }) + + reqBytes, err := proto.Marshal(&arkrpc.GetInfoRequest{}) + if err != nil { + t.Fatalf("marshal request: %v", err) + } + + respMsg, err := mux.ServeRPC( + t.Context(), + arkrpcServiceName, + "GetInfo", + reqBytes, + ) + if err != nil { + t.Fatalf("ServeRPC: %v", err) + } + + resp, ok := respMsg.(*arkrpc.GetInfoResponse) + if !ok { + t.Fatalf("unexpected response type: %T", respMsg) + } + + if resp.GetVersion() != "test" { + t.Fatalf("unexpected version: %q", resp.GetVersion()) + } +} + +// TestServeMuxMissingHandler verifies unknown routes fail with ErrNoHandler. +func TestServeMuxMissingHandler(t *testing.T) { + t.Parallel() + + mux := mailboxrpc.NewServeMux() + + _, err := mux.ServeRPC( + t.Context(), + arkrpcServiceName, + "Missing", + nil, + ) + if err == nil { + t.Fatalf("expected error") + } + if !errors.Is(err, mailboxrpc.ErrNoHandler) { + t.Fatalf("expected ErrNoHandler, got %v", err) + } +} + +const arkrpcServiceName = "arkrpc.ArkService"