mirror of
https://github.com/ethereum/go-ethereum.git
synced 2026-07-27 23:26:44 +00:00
p2p/simulations: Use WebSocket RPC
Signed-off-by: Lewis Marshall <lewis@lmars.net>
This commit is contained in:
parent
1a4b945bff
commit
69c2e537c6
6 changed files with 106 additions and 278 deletions
|
|
@ -203,7 +203,7 @@ func (api *PrivateAdminAPI) StartWS(host *string, port *int, allowedOrigins *str
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := api.node.startWS(fmt.Sprintf("%s:%d", *host, *port), api.node.rpcAPIs, modules, origins); err != nil {
|
if err := api.node.startWS(fmt.Sprintf("%s:%d", *host, *port), api.node.rpcAPIs, modules, origins, api.node.config.WSExposeAll); err != nil {
|
||||||
return false, err
|
return false, err
|
||||||
}
|
}
|
||||||
return true, nil
|
return true, nil
|
||||||
|
|
|
||||||
|
|
@ -128,6 +128,13 @@ type Config struct {
|
||||||
// If the module list is empty, all RPC API endpoints designated public will be
|
// If the module list is empty, all RPC API endpoints designated public will be
|
||||||
// exposed.
|
// exposed.
|
||||||
WSModules []string `toml:",omitempty"`
|
WSModules []string `toml:",omitempty"`
|
||||||
|
|
||||||
|
// WSExposeAll exposes all API modules via the WebSocket RPC interface rather
|
||||||
|
// than just the public ones.
|
||||||
|
//
|
||||||
|
// *WARNING* Only set this if the node is running in a trusted network, exposing
|
||||||
|
// private APIs to untrusted users is a major security risk.
|
||||||
|
WSExposeAll bool `toml:",omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// IPCEndpoint resolves an IPC endpoint based on a configured value, taking into
|
// IPCEndpoint resolves an IPC endpoint based on a configured value, taking into
|
||||||
|
|
|
||||||
|
|
@ -275,7 +275,7 @@ func (n *Node) startRPC(services map[reflect.Type]Service) error {
|
||||||
n.stopInProc()
|
n.stopInProc()
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if err := n.startWS(n.wsEndpoint, apis, n.config.WSModules, n.config.WSOrigins); err != nil {
|
if err := n.startWS(n.wsEndpoint, apis, n.config.WSModules, n.config.WSOrigins, n.config.WSExposeAll); err != nil {
|
||||||
n.stopHTTP()
|
n.stopHTTP()
|
||||||
n.stopIPC()
|
n.stopIPC()
|
||||||
n.stopInProc()
|
n.stopInProc()
|
||||||
|
|
@ -426,7 +426,7 @@ func (n *Node) stopHTTP() {
|
||||||
}
|
}
|
||||||
|
|
||||||
// startWS initializes and starts the websocket RPC endpoint.
|
// startWS initializes and starts the websocket RPC endpoint.
|
||||||
func (n *Node) startWS(endpoint string, apis []rpc.API, modules []string, wsOrigins []string) error {
|
func (n *Node) startWS(endpoint string, apis []rpc.API, modules []string, wsOrigins []string, exposeAll bool) error {
|
||||||
// Short circuit if the WS endpoint isn't being exposed
|
// Short circuit if the WS endpoint isn't being exposed
|
||||||
if endpoint == "" {
|
if endpoint == "" {
|
||||||
return nil
|
return nil
|
||||||
|
|
@ -439,7 +439,7 @@ func (n *Node) startWS(endpoint string, apis []rpc.API, modules []string, wsOrig
|
||||||
// Register all the APIs exposed by the services
|
// Register all the APIs exposed by the services
|
||||||
handler := rpc.NewServer()
|
handler := rpc.NewServer()
|
||||||
for _, api := range apis {
|
for _, api := range apis {
|
||||||
if whitelist[api.Namespace] || (len(whitelist) == 0 && api.Public) {
|
if exposeAll || whitelist[api.Namespace] || (len(whitelist) == 0 && api.Public) {
|
||||||
if err := handler.RegisterName(api.Namespace, api.Service); err != nil {
|
if err := handler.RegisterName(api.Namespace, api.Service); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
@ -455,7 +455,7 @@ func (n *Node) startWS(endpoint string, apis []rpc.API, modules []string, wsOrig
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
go rpc.NewWSServer(wsOrigins, handler).Serve(listener)
|
go rpc.NewWSServer(wsOrigins, handler).Serve(listener)
|
||||||
log.Info(fmt.Sprintf("WebSocket endpoint opened: ws://%s", endpoint))
|
log.Info(fmt.Sprintf("WebSocket endpoint opened: ws://%s", listener.Addr()))
|
||||||
|
|
||||||
// All listeners booted successfully
|
// All listeners booted successfully
|
||||||
n.wsEndpoint = endpoint
|
n.wsEndpoint = endpoint
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,7 @@
|
||||||
package adapters
|
package adapters
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bufio"
|
||||||
"context"
|
"context"
|
||||||
"crypto/ecdsa"
|
"crypto/ecdsa"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
|
@ -12,7 +13,9 @@ import (
|
||||||
"os/exec"
|
"os/exec"
|
||||||
"os/signal"
|
"os/signal"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"regexp"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"syscall"
|
"syscall"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
|
@ -22,6 +25,7 @@ import (
|
||||||
"github.com/ethereum/go-ethereum/p2p"
|
"github.com/ethereum/go-ethereum/p2p"
|
||||||
"github.com/ethereum/go-ethereum/p2p/discover"
|
"github.com/ethereum/go-ethereum/p2p/discover"
|
||||||
"github.com/ethereum/go-ethereum/rpc"
|
"github.com/ethereum/go-ethereum/rpc"
|
||||||
|
"golang.org/x/net/websocket"
|
||||||
)
|
)
|
||||||
|
|
||||||
// ExecAdapter is a NodeAdapter which runs nodes by executing the current
|
// ExecAdapter is a NodeAdapter which runs nodes by executing the current
|
||||||
|
|
@ -69,6 +73,10 @@ func (e *ExecAdapter) NewNode(config *NodeConfig) (Node, error) {
|
||||||
Node: config,
|
Node: config,
|
||||||
}
|
}
|
||||||
conf.Stack.DataDir = filepath.Join(dir, "data")
|
conf.Stack.DataDir = filepath.Join(dir, "data")
|
||||||
|
conf.Stack.WSHost = "127.0.0.1"
|
||||||
|
conf.Stack.WSPort = 0
|
||||||
|
conf.Stack.WSOrigins = []string{"*"}
|
||||||
|
conf.Stack.WSExposeAll = true
|
||||||
conf.Stack.P2P.EnableMsgEvents = false
|
conf.Stack.P2P.EnableMsgEvents = false
|
||||||
conf.Stack.P2P.NoDiscovery = true
|
conf.Stack.P2P.NoDiscovery = true
|
||||||
conf.Stack.P2P.NAT = nil
|
conf.Stack.P2P.NAT = nil
|
||||||
|
|
@ -101,7 +109,7 @@ type ExecNode struct {
|
||||||
Info *p2p.NodeInfo
|
Info *p2p.NodeInfo
|
||||||
|
|
||||||
client *rpc.Client
|
client *rpc.Client
|
||||||
rpcMux *rpcMux
|
wsAddr string
|
||||||
newCmd func() *exec.Cmd
|
newCmd func() *exec.Cmd
|
||||||
key *ecdsa.PrivateKey
|
key *ecdsa.PrivateKey
|
||||||
}
|
}
|
||||||
|
|
@ -120,6 +128,10 @@ func (n *ExecNode) Client() (*rpc.Client, error) {
|
||||||
return n.client, nil
|
return n.client, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// wsAddrPattern is a regex used to read the WebSocket address from the node's
|
||||||
|
// log
|
||||||
|
var wsAddrPattern = regexp.MustCompile(`ws://[\d.:]+`)
|
||||||
|
|
||||||
// Start exec's the node passing the ID and service as command line arguments
|
// Start exec's the node passing the ID and service as command line arguments
|
||||||
// and the node config encoded as JSON in the _P2P_NODE_CONFIG environment
|
// and the node config encoded as JSON in the _P2P_NODE_CONFIG environment
|
||||||
// variable
|
// variable
|
||||||
|
|
@ -142,29 +154,61 @@ func (n *ExecNode) Start(snapshots map[string][]byte) (err error) {
|
||||||
return fmt.Errorf("error generating node config: %s", err)
|
return fmt.Errorf("error generating node config: %s", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// create a net.Pipe for RPC communication over stdin / stdout
|
// use a pipe for stderr so we can both copy the node's stderr to
|
||||||
pipe1, pipe2 := net.Pipe()
|
// os.Stderr and read the WebSocket address from the logs
|
||||||
|
stderrR, stderrW := io.Pipe()
|
||||||
|
stderr := io.MultiWriter(os.Stderr, stderrW)
|
||||||
|
|
||||||
// start the node
|
// start the node
|
||||||
cmd := n.newCmd()
|
cmd := n.newCmd()
|
||||||
cmd.Stdin = pipe1
|
cmd.Stdout = os.Stdout
|
||||||
cmd.Stdout = pipe1
|
cmd.Stderr = stderr
|
||||||
cmd.Stderr = os.Stderr
|
|
||||||
cmd.Env = append(os.Environ(), fmt.Sprintf("_P2P_NODE_CONFIG=%s", confData))
|
cmd.Env = append(os.Environ(), fmt.Sprintf("_P2P_NODE_CONFIG=%s", confData))
|
||||||
if err := cmd.Start(); err != nil {
|
if err := cmd.Start(); err != nil {
|
||||||
return fmt.Errorf("error starting node: %s", err)
|
return fmt.Errorf("error starting node: %s", err)
|
||||||
}
|
}
|
||||||
n.Cmd = cmd
|
n.Cmd = cmd
|
||||||
|
|
||||||
|
// read the WebSocket address from the stderr logs
|
||||||
|
var wsAddr string
|
||||||
|
errC := make(chan error)
|
||||||
|
go func() {
|
||||||
|
s := bufio.NewScanner(stderrR)
|
||||||
|
for s.Scan() {
|
||||||
|
if strings.Contains(s.Text(), "WebSocket endpoint opened:") {
|
||||||
|
wsAddr = wsAddrPattern.FindString(s.Text())
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case errC <- s.Err():
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
select {
|
||||||
|
case err := <-errC:
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("error reading WebSocket address from stderr: %s", err)
|
||||||
|
} else if wsAddr == "" {
|
||||||
|
return errors.New("failed to read WebSocket address from stderr")
|
||||||
|
}
|
||||||
|
case <-time.After(10 * time.Second):
|
||||||
|
return errors.New("timed out waiting for WebSocket address on stderr")
|
||||||
|
}
|
||||||
|
|
||||||
// create the RPC client and load the node info
|
// create the RPC client and load the node info
|
||||||
n.rpcMux = newRPCMux(pipe2)
|
|
||||||
n.client = n.rpcMux.Client()
|
|
||||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
client, err := rpc.DialWebsocket(ctx, wsAddr, "")
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("error dialing rpc websocket: %s", err)
|
||||||
|
}
|
||||||
var info p2p.NodeInfo
|
var info p2p.NodeInfo
|
||||||
if err := n.client.CallContext(ctx, &info, "admin_nodeInfo"); err != nil {
|
if err := client.CallContext(ctx, &info, "admin_nodeInfo"); err != nil {
|
||||||
return fmt.Errorf("error getting node info: %s", err)
|
return fmt.Errorf("error getting node info: %s", err)
|
||||||
}
|
}
|
||||||
|
n.client = client
|
||||||
|
n.wsAddr = wsAddr
|
||||||
n.Info = &info
|
n.Info = &info
|
||||||
|
|
||||||
return nil
|
return nil
|
||||||
|
|
@ -193,6 +237,7 @@ func (n *ExecNode) Stop() error {
|
||||||
if n.client != nil {
|
if n.client != nil {
|
||||||
n.client.Close()
|
n.client.Close()
|
||||||
n.client = nil
|
n.client = nil
|
||||||
|
n.wsAddr = ""
|
||||||
n.Info = nil
|
n.Info = nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -222,13 +267,30 @@ func (n *ExecNode) NodeInfo() *p2p.NodeInfo {
|
||||||
return info
|
return info
|
||||||
}
|
}
|
||||||
|
|
||||||
// ServeRPC serves RPC requests over the given connection using the node's
|
// ServeRPC serves RPC requests over the given connection by dialling the
|
||||||
// RPC multiplexer
|
// node's WebSocket address and joining the two connections
|
||||||
func (n *ExecNode) ServeRPC(conn net.Conn) error {
|
func (n *ExecNode) ServeRPC(clientConn net.Conn) error {
|
||||||
if n.rpcMux == nil {
|
conn, err := websocket.Dial(n.wsAddr, "", "http://localhost")
|
||||||
return errors.New("RPC not started")
|
if err != nil {
|
||||||
|
return err
|
||||||
}
|
}
|
||||||
n.rpcMux.Serve(conn)
|
var wg sync.WaitGroup
|
||||||
|
wg.Add(2)
|
||||||
|
join := func(src, dst net.Conn) {
|
||||||
|
defer wg.Done()
|
||||||
|
io.Copy(dst, src)
|
||||||
|
// close the write end of the destination connection
|
||||||
|
if cw, ok := dst.(interface {
|
||||||
|
CloseWrite() error
|
||||||
|
}); ok {
|
||||||
|
cw.CloseWrite()
|
||||||
|
} else {
|
||||||
|
dst.Close()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
go join(conn, clientConn)
|
||||||
|
go join(clientConn, conn)
|
||||||
|
wg.Wait()
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -313,15 +375,6 @@ func execP2PNode() {
|
||||||
log.Crit("error starting p2p node", "err", err)
|
log.Crit("error starting p2p node", "err", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// use stdin / stdout for RPC to avoid the parent needing to access
|
|
||||||
// either the local filesystem or TCP stack (useful when running in
|
|
||||||
// Docker)
|
|
||||||
handler, err := stack.RPCHandler()
|
|
||||||
if err != nil {
|
|
||||||
log.Crit("error getting RPC server", "err", err)
|
|
||||||
}
|
|
||||||
go handler.ServeCodec(rpc.NewJSONCodec(&stdioConn{os.Stdin, os.Stdout}), rpc.OptionMethodInvocation|rpc.OptionSubscriptions)
|
|
||||||
|
|
||||||
go func() {
|
go func() {
|
||||||
sigc := make(chan os.Signal, 1)
|
sigc := make(chan os.Signal, 1)
|
||||||
signal.Notify(sigc, syscall.SIGTERM)
|
signal.Notify(sigc, syscall.SIGTERM)
|
||||||
|
|
@ -404,14 +457,3 @@ func (api SnapshotAPI) Snapshot() (map[string][]byte, error) {
|
||||||
}
|
}
|
||||||
return snapshots, nil
|
return snapshots, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// stdioConn wraps os.Stdin / os.Stdout with a no-op Close method so we can
|
|
||||||
// use stdio for RPC messages
|
|
||||||
type stdioConn struct {
|
|
||||||
io.Reader
|
|
||||||
io.Writer
|
|
||||||
}
|
|
||||||
|
|
||||||
func (r *stdioConn) Close() error {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
|
||||||
|
|
@ -119,7 +119,6 @@ type SimNode struct {
|
||||||
node *node.Node
|
node *node.Node
|
||||||
running []node.Service
|
running []node.Service
|
||||||
client *rpc.Client
|
client *rpc.Client
|
||||||
rpcMux *rpcMux
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Addr returns the node's discovery address
|
// Addr returns the node's discovery address
|
||||||
|
|
@ -135,24 +134,28 @@ func (self *SimNode) Node() *discover.Node {
|
||||||
// Client returns an rpc.Client which can be used to communicate with the
|
// Client returns an rpc.Client which can be used to communicate with the
|
||||||
// underlying service (it is set once the node has started)
|
// underlying service (it is set once the node has started)
|
||||||
func (self *SimNode) Client() (*rpc.Client, error) {
|
func (self *SimNode) Client() (*rpc.Client, error) {
|
||||||
self.lock.Lock()
|
self.lock.RLock()
|
||||||
defer self.lock.Unlock()
|
defer self.lock.RUnlock()
|
||||||
if self.client == nil {
|
if self.client == nil {
|
||||||
return nil, errors.New("RPC not started")
|
return nil, errors.New("node not started")
|
||||||
}
|
}
|
||||||
return self.client, nil
|
return self.client, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// ServeRPC serves RPC requests over the given connection using the node's
|
// ServeRPC serves RPC requests over the given connection by creating an
|
||||||
// RPC multiplexer
|
// in-process client to the node's RPC server
|
||||||
func (self *SimNode) ServeRPC(conn net.Conn) error {
|
func (self *SimNode) ServeRPC(conn net.Conn) error {
|
||||||
self.lock.Lock()
|
self.lock.RLock()
|
||||||
mux := self.rpcMux
|
node := self.node
|
||||||
self.lock.Unlock()
|
self.lock.RUnlock()
|
||||||
if mux == nil {
|
if node == nil {
|
||||||
return errors.New("RPC not started")
|
return errors.New("node not started")
|
||||||
}
|
}
|
||||||
mux.Serve(conn)
|
handler, err := node.RPCHandler()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
handler.ServeCodec(rpc.NewJSONCodec(conn), rpc.OptionMethodInvocation|rpc.OptionSubscriptions)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -213,18 +216,12 @@ func (self *SimNode) Start(snapshots map[string][]byte) error {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// create an in-process RPC client
|
||||||
handler, err := node.RPCHandler()
|
handler, err := node.RPCHandler()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
self.client = rpc.DialInProc(handler)
|
||||||
// create an in-process RPC multiplexer
|
|
||||||
pipe1, pipe2 := net.Pipe()
|
|
||||||
go handler.ServeCodec(rpc.NewJSONCodec(pipe1), rpc.OptionMethodInvocation|rpc.OptionSubscriptions)
|
|
||||||
self.rpcMux = newRPCMux(pipe2)
|
|
||||||
|
|
||||||
// create an in-process RPC client
|
|
||||||
self.client = self.rpcMux.Client()
|
|
||||||
|
|
||||||
self.node = node
|
self.node = node
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,218 +0,0 @@
|
||||||
// Copyright 2016 The go-ethereum Authors
|
|
||||||
// This file is part of the go-ethereum library.
|
|
||||||
//
|
|
||||||
// The go-ethereum library is free software: you can redistribute it and/or modify
|
|
||||||
// it under the terms of the GNU Lesser General Public License as published by
|
|
||||||
// the Free Software Foundation, either version 3 of the License, or
|
|
||||||
// (at your option) any later version.
|
|
||||||
//
|
|
||||||
// The go-ethereum library is distributed in the hope that it will be useful,
|
|
||||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
||||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
||||||
// GNU Lesser General Public License for more details.
|
|
||||||
//
|
|
||||||
// You should have received a copy of the GNU Lesser General Public License
|
|
||||||
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
|
|
||||||
|
|
||||||
package adapters
|
|
||||||
|
|
||||||
import (
|
|
||||||
"encoding/json"
|
|
||||||
"net"
|
|
||||||
"strconv"
|
|
||||||
"strings"
|
|
||||||
"sync"
|
|
||||||
|
|
||||||
"github.com/ethereum/go-ethereum/rpc"
|
|
||||||
)
|
|
||||||
|
|
||||||
// rpcMux is an RPC multiplexer which allows many clients to make RPC requests
|
|
||||||
// over a single connection by changing each request's ID to a unique value.
|
|
||||||
//
|
|
||||||
// This is used by node adapters so that simulations can create many RPC
|
|
||||||
// clients all sending requests over the underlying node's stdin / stdout.
|
|
||||||
type rpcMux struct {
|
|
||||||
conn net.Conn
|
|
||||||
|
|
||||||
mtx sync.Mutex
|
|
||||||
idCounter uint64
|
|
||||||
msgMap map[uint64]*rpcMsg
|
|
||||||
subMap map[string]*rpcReply
|
|
||||||
send chan *rpcMsg
|
|
||||||
}
|
|
||||||
|
|
||||||
type rpcMsg struct {
|
|
||||||
Method string `json:"method,omitempty"`
|
|
||||||
Version string `json:"jsonrpc,omitempty"`
|
|
||||||
ID json.RawMessage `json:"id,omitempty"`
|
|
||||||
Payload json.RawMessage `json:"params,omitempty"`
|
|
||||||
Result json.RawMessage `json:"result,omitempty"`
|
|
||||||
Error json.RawMessage `json:"error,omitempty"`
|
|
||||||
|
|
||||||
id uint64
|
|
||||||
reply *rpcReply
|
|
||||||
}
|
|
||||||
|
|
||||||
// rpcSub is the payload or result of a subscription RPC message
|
|
||||||
type rpcSub struct {
|
|
||||||
Subscription string `json:"subscription"`
|
|
||||||
Result json.RawMessage `json:"result,omitempty"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// rpcReply receives replies to RPC messages for a particular client
|
|
||||||
type rpcReply struct {
|
|
||||||
ch chan *rpcMsg
|
|
||||||
closeOnce sync.Once
|
|
||||||
}
|
|
||||||
|
|
||||||
func (r *rpcReply) close() {
|
|
||||||
r.closeOnce.Do(func() { close(r.ch) })
|
|
||||||
}
|
|
||||||
|
|
||||||
func newRPCMux(conn net.Conn) *rpcMux {
|
|
||||||
mux := &rpcMux{
|
|
||||||
msgMap: make(map[uint64]*rpcMsg),
|
|
||||||
subMap: make(map[string]*rpcReply),
|
|
||||||
send: make(chan *rpcMsg),
|
|
||||||
}
|
|
||||||
go mux.sendLoop(conn)
|
|
||||||
go mux.recvLoop(conn)
|
|
||||||
return mux
|
|
||||||
}
|
|
||||||
|
|
||||||
// Client creates a new RPC client which sends messages through the multiplexer
|
|
||||||
func (mux *rpcMux) Client() *rpc.Client {
|
|
||||||
pipe1, pipe2 := net.Pipe()
|
|
||||||
go mux.Serve(pipe1)
|
|
||||||
return rpc.NewClientWithConn(pipe2)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Serve reads RPC messages from the given connection, forwards them to the
|
|
||||||
// multiplexed connnection and writes replies back to the given connection
|
|
||||||
func (mux *rpcMux) Serve(conn net.Conn) {
|
|
||||||
// reply will receive replies to any messages we send
|
|
||||||
reply := &rpcReply{ch: make(chan *rpcMsg)}
|
|
||||||
defer func() {
|
|
||||||
// drain the channel to prevent blocking the recvLoop
|
|
||||||
for range reply.ch {
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
|
|
||||||
// start a goroutine to read RPC messages from the connection and
|
|
||||||
// forward them to the sendLoop
|
|
||||||
done := make(chan struct{})
|
|
||||||
go func() {
|
|
||||||
defer close(done)
|
|
||||||
dec := json.NewDecoder(conn)
|
|
||||||
for {
|
|
||||||
msg := &rpcMsg{}
|
|
||||||
if err := dec.Decode(msg); err != nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
msg.reply = reply
|
|
||||||
mux.send <- msg
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
|
|
||||||
// write message replies to the connection
|
|
||||||
enc := json.NewEncoder(conn)
|
|
||||||
for {
|
|
||||||
select {
|
|
||||||
case msg, ok := <-reply.ch:
|
|
||||||
if !ok {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
if err := enc.Encode(msg); err != nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
case <-done:
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// sendLoop receives messages from the send channel, changes their ID and
|
|
||||||
// writes them to the given connection
|
|
||||||
func (mux *rpcMux) sendLoop(conn net.Conn) {
|
|
||||||
enc := json.NewEncoder(conn)
|
|
||||||
for msg := range mux.send {
|
|
||||||
if err := enc.Encode(mux.newMsg(msg)); err != nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// recvLoop reads messages from the given connection, changes their ID back
|
|
||||||
// to the oringal value and sends them to the message's reply channel
|
|
||||||
func (mux *rpcMux) recvLoop(conn net.Conn) {
|
|
||||||
// close all reply channels if we get an error
|
|
||||||
defer func() {
|
|
||||||
mux.mtx.Lock()
|
|
||||||
defer mux.mtx.Unlock()
|
|
||||||
for _, msg := range mux.msgMap {
|
|
||||||
msg.reply.close()
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
|
|
||||||
dec := json.NewDecoder(conn)
|
|
||||||
for {
|
|
||||||
msg := &rpcMsg{}
|
|
||||||
if err := dec.Decode(msg); err != nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
if reply := mux.lookup(msg); reply != nil {
|
|
||||||
reply.ch <- msg
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// newMsg copies the given message and changes it's ID to a unique value
|
|
||||||
func (mux *rpcMux) newMsg(msg *rpcMsg) *rpcMsg {
|
|
||||||
mux.mtx.Lock()
|
|
||||||
defer mux.mtx.Unlock()
|
|
||||||
id := mux.idCounter
|
|
||||||
mux.idCounter++
|
|
||||||
mux.msgMap[id] = msg
|
|
||||||
newMsg := *msg
|
|
||||||
newMsg.ID = json.RawMessage(strconv.FormatUint(id, 10))
|
|
||||||
return &newMsg
|
|
||||||
}
|
|
||||||
|
|
||||||
// lookup looks up the original message for which the given message is a reply
|
|
||||||
func (mux *rpcMux) lookup(msg *rpcMsg) *rpcReply {
|
|
||||||
mux.mtx.Lock()
|
|
||||||
defer mux.mtx.Unlock()
|
|
||||||
|
|
||||||
// if the message has no ID, it is a subscription notification so
|
|
||||||
// lookup the original subscribe message
|
|
||||||
if msg.ID == nil {
|
|
||||||
sub := &rpcSub{}
|
|
||||||
if err := json.Unmarshal(msg.Payload, sub); err != nil {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
return mux.subMap[sub.Subscription]
|
|
||||||
}
|
|
||||||
|
|
||||||
// lookup the original message and restore the ID
|
|
||||||
id, err := strconv.ParseUint(string(msg.ID), 10, 64)
|
|
||||||
if err != nil {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
origMsg, ok := mux.msgMap[id]
|
|
||||||
if !ok {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
delete(mux.msgMap, id)
|
|
||||||
msg.ID = origMsg.ID
|
|
||||||
|
|
||||||
// if the original message was a subscription, store the subscription
|
|
||||||
// ID so we can detect notifications
|
|
||||||
if strings.HasSuffix(string(origMsg.Method), "_subscribe") {
|
|
||||||
var result string
|
|
||||||
if err := json.Unmarshal(msg.Result, &result); err == nil {
|
|
||||||
mux.subMap[result] = origMsg.reply
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return origMsg.reply
|
|
||||||
}
|
|
||||||
Loading…
Reference in a new issue