diff --git a/pkg/rpc/client/wsclient.go b/pkg/rpc/client/wsclient.go index bdac24816..6774dd421 100644 --- a/pkg/rpc/client/wsclient.go +++ b/pkg/rpc/client/wsclient.go @@ -7,8 +7,11 @@ import ( "time" "github.com/gorilla/websocket" + "github.com/nspcc-dev/neo-go/pkg/core/block" + "github.com/nspcc-dev/neo-go/pkg/core/transaction" "github.com/nspcc-dev/neo-go/pkg/rpc/request" "github.com/nspcc-dev/neo-go/pkg/rpc/response" + "github.com/nspcc-dev/neo-go/pkg/rpc/response/result" ) // WSClient is a websocket-enabled RPC client that can be used with appropriate @@ -17,12 +20,28 @@ import ( // that is only provided via websockets (like event subscription mechanism). type WSClient struct { Client + // Notifications is a channel that is used to send events received from + // server. Client's code is supposed to be reading from this channel if + // it wants to use subscription mechanism, failing to do so will cause + // WSClient to block even regular requests. This channel is not buffered. + // In case of protocol error or upon connection closure this channel will + // be closed, so make sure to handle this. + Notifications chan Notification + ws *websocket.Conn done chan struct{} - notifications chan *request.In responses chan *response.Raw requests chan *request.Raw shutdown chan struct{} + subscriptions map[string]bool +} + +// Notification represents server-generated notification for client subscriptions. +// Value can be one of block.Block, result.ApplicationLog, result.NotificationEvent +// or transaction.Transaction based on Type. +type Notification struct { + Type response.EventID + Value interface{} } // requestResponse is a combined type for request and response since we can get @@ -59,12 +78,15 @@ func NewWS(ctx context.Context, endpoint string, opts Options) (*WSClient, error return nil, err } wsc := &WSClient{ - Client: *cl, - ws: ws, - shutdown: make(chan struct{}), - done: make(chan struct{}), - responses: make(chan *response.Raw), - requests: make(chan *request.Raw), + Client: *cl, + Notifications: make(chan Notification), + + ws: ws, + shutdown: make(chan struct{}), + done: make(chan struct{}), + responses: make(chan *response.Raw), + requests: make(chan *request.Raw), + subscriptions: make(map[string]bool), } go wsc.wsReader() go wsc.wsWriter() @@ -86,6 +108,7 @@ func (c *WSClient) Close() { func (c *WSClient) wsReader() { c.ws.SetReadLimit(wsReadLimit) c.ws.SetPongHandler(func(string) error { c.ws.SetReadDeadline(time.Now().Add(wsPongLimit)); return nil }) +readloop: for { rr := new(requestResponse) c.ws.SetReadDeadline(time.Now().Add(wsPongLimit)) @@ -95,9 +118,37 @@ func (c *WSClient) wsReader() { break } if rr.RawID == nil && rr.Method != "" { - if c.notifications != nil { - c.notifications <- &rr.In + event, err := response.GetEventIDFromString(rr.Method) + if err != nil { + // Bad event received. + break } + var slice []json.RawMessage + err = json.Unmarshal(rr.RawParams, &slice) + if err != nil || len(slice) != 1 { + // Bad event received. + break + } + var val interface{} + switch event { + case response.BlockEventID: + val = new(block.Block) + case response.TransactionEventID: + val = new(transaction.Transaction) + case response.NotificationEventID: + val = new(result.NotificationEvent) + case response.ExecutionEventID: + val = new(result.ApplicationLog) + default: + // Bad event received. + break readloop + } + err = json.Unmarshal(slice[0], val) + if err != nil || len(slice) != 1 { + // Bad event received. + break + } + c.Notifications <- Notification{event, val} } else if rr.RawID != nil && (rr.Error != nil || rr.Result != nil) { resp := new(response.Raw) resp.ID = rr.RawID @@ -112,9 +163,7 @@ func (c *WSClient) wsReader() { } close(c.done) close(c.responses) - if c.notifications != nil { - close(c.notifications) - } + close(c.Notifications) } func (c *WSClient) wsWriter() { @@ -158,3 +207,73 @@ func (c *WSClient) makeWsRequest(r *request.Raw) (*response.Raw, error) { return resp, nil } } + +func (c *WSClient) performSubscription(params request.RawParams) (string, error) { + var resp string + + if err := c.performRequest("subscribe", params, &resp); err != nil { + return "", err + } + c.subscriptions[resp] = true + return resp, nil +} + +func (c *WSClient) performUnsubscription(id string) error { + var resp bool + + if !c.subscriptions[id] { + return errors.New("no subscription with this ID") + } + if err := c.performRequest("unsubscribe", request.NewRawParams(id), &resp); err != nil { + return err + } + if !resp { + return errors.New("unsubscribe method returned false result") + } + delete(c.subscriptions, id) + return nil +} + +// SubscribeForNewBlocks adds subscription for new block events to this instance +// of client. +func (c *WSClient) SubscribeForNewBlocks() (string, error) { + params := request.NewRawParams("block_added") + return c.performSubscription(params) +} + +// SubscribeForNewTransactions adds subscription for new transaction events to +// this instance of client. +func (c *WSClient) SubscribeForNewTransactions() (string, error) { + params := request.NewRawParams("transaction_added") + return c.performSubscription(params) +} + +// SubscribeForExecutionNotifications adds subscription for notifications +// generated during transaction execution to this instance of client. +func (c *WSClient) SubscribeForExecutionNotifications() (string, error) { + params := request.NewRawParams("notification_from_execution") + return c.performSubscription(params) +} + +// SubscribeForTransactionExecutions adds subscription for application execution +// results generated during transaction execution to this instance of client. +func (c *WSClient) SubscribeForTransactionExecutions() (string, error) { + params := request.NewRawParams("transaction_executed") + return c.performSubscription(params) +} + +// Unsubscribe removes subscription for given event stream. +func (c *WSClient) Unsubscribe(id string) error { + return c.performUnsubscription(id) +} + +// UnsubscribeAll removes all active subscriptions of current client. +func (c *WSClient) UnsubscribeAll() error { + for id := range c.subscriptions { + err := c.performUnsubscription(id) + if err != nil { + return err + } + } + return nil +} diff --git a/pkg/rpc/client/wsclient_test.go b/pkg/rpc/client/wsclient_test.go index 2a996999a..f78481f05 100644 --- a/pkg/rpc/client/wsclient_test.go +++ b/pkg/rpc/client/wsclient_test.go @@ -2,8 +2,12 @@ package client import ( "context" + "net/http" + "net/http/httptest" "testing" + "time" + "github.com/gorilla/websocket" "github.com/stretchr/testify/require" ) @@ -14,3 +18,129 @@ func TestWSClientClose(t *testing.T) { require.NoError(t, err) wsc.Close() } + +func TestWSClientSubscription(t *testing.T) { + var cases = map[string]func(*WSClient) (string, error){ + "blocks": (*WSClient).SubscribeForNewBlocks, + "transactions": (*WSClient).SubscribeForNewTransactions, + "notifications": (*WSClient).SubscribeForExecutionNotifications, + "executions": (*WSClient).SubscribeForTransactionExecutions, + } + t.Run("good", func(t *testing.T) { + for name, f := range cases { + t.Run(name, func(t *testing.T) { + srv := initTestServer(t, `{"jsonrpc": "2.0", "id": 1, "result": "55aaff00"}`) + defer srv.Close() + wsc, err := NewWS(context.TODO(), httpURLtoWS(srv.URL), Options{}) + require.NoError(t, err) + id, err := f(wsc) + require.NoError(t, err) + require.Equal(t, "55aaff00", id) + }) + } + }) + t.Run("bad", func(t *testing.T) { + for name, f := range cases { + t.Run(name, func(t *testing.T) { + srv := initTestServer(t, `{"jsonrpc": "2.0", "id": 1, "error":{"code":-32602,"message":"Invalid Params"}}`) + defer srv.Close() + wsc, err := NewWS(context.TODO(), httpURLtoWS(srv.URL), Options{}) + require.NoError(t, err) + _, err = f(wsc) + require.Error(t, err) + }) + } + }) +} + +func TestWSClientUnsubscription(t *testing.T) { + type responseCheck struct { + response string + code func(*testing.T, *WSClient) + } + var cases = map[string]responseCheck{ + "good": {`{"jsonrpc": "2.0", "id": 1, "result": true}`, func(t *testing.T, wsc *WSClient) { + // We can't really subscribe using this stub server, so set up wsc internals. + wsc.subscriptions["0"] = true + err := wsc.Unsubscribe("0") + require.NoError(t, err) + }}, + "all": {`{"jsonrpc": "2.0", "id": 1, "result": true}`, func(t *testing.T, wsc *WSClient) { + // We can't really subscribe using this stub server, so set up wsc internals. + wsc.subscriptions["0"] = true + err := wsc.UnsubscribeAll() + require.NoError(t, err) + require.Equal(t, 0, len(wsc.subscriptions)) + }}, + "not subscribed": {`{"jsonrpc": "2.0", "id": 1, "result": true}`, func(t *testing.T, wsc *WSClient) { + err := wsc.Unsubscribe("0") + require.Error(t, err) + }}, + "error returned": {`{"jsonrpc": "2.0", "id": 1, "error":{"code":-32602,"message":"Invalid Params"}}`, func(t *testing.T, wsc *WSClient) { + // We can't really subscribe using this stub server, so set up wsc internals. + wsc.subscriptions["0"] = true + err := wsc.Unsubscribe("0") + require.Error(t, err) + }}, + "false returned": {`{"jsonrpc": "2.0", "id": 1, "result": false}`, func(t *testing.T, wsc *WSClient) { + // We can't really subscribe using this stub server, so set up wsc internals. + wsc.subscriptions["0"] = true + err := wsc.Unsubscribe("0") + require.Error(t, err) + }}, + } + for name, rc := range cases { + t.Run(name, func(t *testing.T) { + srv := initTestServer(t, rc.response) + defer srv.Close() + wsc, err := NewWS(context.TODO(), httpURLtoWS(srv.URL), Options{}) + require.NoError(t, err) + rc.code(t, wsc) + }) + } +} + +func TestWSClientEvents(t *testing.T) { + var ok bool + // Events from RPC server test chain. + var events = []string{ + `{"jsonrpc":"2.0","method":"transaction_executed","params":[{"txid":"0xe1cd5e57e721d2a2e05fb1f08721b12057b25ab1dd7fd0f33ee1639932fdfad7","executions":[{"trigger":"Application","contract":"0x0000000000000000000000000000000000000000","vmstate":"HALT","gas_consumed":"2.291","stack":[],"notifications":[{"contract":"0x1b4357bff5a01bdf2a6581247cf9ed1e24629176","state":{"type":"Array","value":[{"type":"ByteArray","value":"636f6e74726163742063616c6c"},{"type":"ByteArray","value":"7472616e73666572"},{"type":"Array","value":[{"type":"ByteArray","value":"769162241eedf97c2481652adf1ba0f5bf57431b"},{"type":"ByteArray","value":"316e851039019d39dfc2c37d6c3fee19fd580987"},{"type":"Integer","value":"1000"}]}]}},{"contract":"0x1b4357bff5a01bdf2a6581247cf9ed1e24629176","state":{"type":"Array","value":[{"type":"ByteArray","value":"7472616e73666572"},{"type":"ByteArray","value":"769162241eedf97c2481652adf1ba0f5bf57431b"},{"type":"ByteArray","value":"316e851039019d39dfc2c37d6c3fee19fd580987"},{"type":"Integer","value":"1000"}]}}]}]}]}`, + `{"jsonrpc":"2.0","method":"notification_from_execution","params":[{"contract":"0x1b4357bff5a01bdf2a6581247cf9ed1e24629176","state":{"type":"Array","value":[{"type":"ByteArray","value":"636f6e74726163742063616c6c"},{"type":"ByteArray","value":"7472616e73666572"},{"type":"Array","value":[{"type":"ByteArray","value":"769162241eedf97c2481652adf1ba0f5bf57431b"},{"type":"ByteArray","value":"316e851039019d39dfc2c37d6c3fee19fd580987"},{"type":"Integer","value":"1000"}]}]}}]}`, + `{"jsonrpc":"2.0","method":"transaction_added","params":[{"txid":"0xe1cd5e57e721d2a2e05fb1f08721b12057b25ab1dd7fd0f33ee1639932fdfad7","size":277,"type":"InvocationTransaction","version":1,"nonce":9,"sender":"ALHF9wsXZVEuCGgmDA6ZNsCLtrb4A1g4yG","sys_fee":"0","net_fee":"0.0037721","valid_until_block":1200,"attributes":[],"cosigners":[{"account":"0x870958fd19ee3f6c7dc3c2df399d013910856e31","scopes":1}],"vin":[],"vout":[],"scripts":[{"invocation":"0c4027727296b84853c5d9e07fb8a40e885246ae25641383b16eefbe92027ecb1635b794aacf6bbfc3e828c73829b14791c483d19eb758b57638e3191393dbf2d288","verification":"0c2102b3622bf4017bdfe317c58aed5f4c753f206b7db896046fa7d774bbc4bf7f8dc20b410a906ad4"}],"script":"01e8030c14316e851039019d39dfc2c37d6c3fee19fd5809870c14769162241eedf97c2481652adf1ba0f5bf57431b13c00c087472616e736665720c14769162241eedf97c2481652adf1ba0f5bf57431b41627d5b5238"}]}`, + `{"jsonrpc":"2.0","method":"block_added","params":[{"version":0,"previousblockhash":"0x04f7580b111ec75f0ce68d3a9fd70a0544b4521b4a98541694d8575c548b759e","merkleroot":"0xb2c7230ebee4cb83bc03afadbba413e6bca8fcdeaf9c077bea060918da0e52a1","time":1590006200,"height":207,"next_consensus":"0x4138145d67638db97b402ee0b6751ef16253ecab","script":{"invocation":"0c4063429fca5ff75c964d9e38179c75978e33f8174d91a780c2e825265cf2447281594afdd5f3e216dcaf5ff0693aec83f415996cf224454495495f6bd0a4c5d08f0c4099680903a954278580d8533121c2cd3e53a089817b6a784901ec06178a60b5f1da6e70422bdcadc89029767e08d66ce4180b99334cb2d42f42e4216394af15920c4067d5e362189e48839a24e187c59d46f5d9db862c8a029777f1548b19632bfdc73ad373827ed02369f925e89c2303b64e6b9838dca229949b9b9d3bd4c0c3ed8f0c4021d4c00d4522805883f1db929554441bcbbee127c48f6b7feeeb69a72a78c7f0a75011663e239c0820ef903f36168f42936de10f0ef20681cb735a4b53d0390f","verification":"130c2102103a7f7dd016558597f7960d27c516a4394fd968b9e65155eb4b013e4040406e0c2102a7bc55fe8684e0119768d104ba30795bdcc86619e864add26156723ed185cd620c2102b3622bf4017bdfe317c58aed5f4c753f206b7db896046fa7d774bbc4bf7f8dc20c2103d90c07df63e690ce77912e10ab51acc944b66860237b608c4f8f8309e71ee699140b413073b3bb"},"consensus_data":{"primary":0,"nonce":1111},"tx":[{"txid":"0xf736cd91ab84062a21a09b424346b241987f6245ffe8c2b2db39d595c3c222f7","size":204,"type":"InvocationTransaction","version":1,"nonce":8,"sender":"ALHF9wsXZVEuCGgmDA6ZNsCLtrb4A1g4yG","sys_fee":"0","net_fee":"0.0030421","valid_until_block":1200,"attributes":[],"cosigners":[],"vin":[],"vout":[],"scripts":[{"invocation":"0c4016e7a112742409cdfaad89dcdbcb52c94c5c1a69dfe5d8b999649eaaa787e31ca496d1734d6ea606c749ad36e9a88892240ae59e0efa7f544e0692124898d512","verification":"0c2102b3622bf4017bdfe317c58aed5f4c753f206b7db896046fa7d774bbc4bf7f8dc20b410a906ad4"}],"script":"10c00c04696e69740c14769162241eedf97c2481652adf1ba0f5bf57431b41627d5b52"},{"txid":"0xe1cd5e57e721d2a2e05fb1f08721b12057b25ab1dd7fd0f33ee1639932fdfad7","size":277,"type":"InvocationTransaction","version":1,"nonce":9,"sender":"ALHF9wsXZVEuCGgmDA6ZNsCLtrb4A1g4yG","sys_fee":"0","net_fee":"0.0037721","valid_until_block":1200,"attributes":[],"cosigners":[{"account":"0x870958fd19ee3f6c7dc3c2df399d013910856e31","scopes":1}],"vin":[],"vout":[],"scripts":[{"invocation":"0c4027727296b84853c5d9e07fb8a40e885246ae25641383b16eefbe92027ecb1635b794aacf6bbfc3e828c73829b14791c483d19eb758b57638e3191393dbf2d288","verification":"0c2102b3622bf4017bdfe317c58aed5f4c753f206b7db896046fa7d774bbc4bf7f8dc20b410a906ad4"}],"script":"01e8030c14316e851039019d39dfc2c37d6c3fee19fd5809870c14769162241eedf97c2481652adf1ba0f5bf57431b13c00c087472616e736665720c14769162241eedf97c2481652adf1ba0f5bf57431b41627d5b5238"}]}]}`, + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + if req.URL.Path == "/ws" && req.Method == "GET" { + var upgrader = websocket.Upgrader{} + ws, err := upgrader.Upgrade(w, req, nil) + require.NoError(t, err) + for _, event := range events { + ws.SetWriteDeadline(time.Now().Add(2 * time.Second)) + err = ws.WriteMessage(1, []byte(event)) + if err != nil { + break + } + } + ws.Close() + return + } + })) + + wsc, err := NewWS(context.TODO(), httpURLtoWS(srv.URL), Options{}) + require.NoError(t, err) + for range events { + select { + case _, ok = <-wsc.Notifications: + case <-time.After(time.Second): + t.Fatal("timeout waiting for event") + } + require.True(t, ok) + } + select { + case _, ok = <-wsc.Notifications: + case <-time.After(time.Second): + t.Fatal("timeout waiting for event") + } + // Connection closed by server. + require.False(t, ok) +}