From 555f42cfd836dabe3151becffa63976c91e5344b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Szil=C3=A1gyi?= Date: Tue, 13 Mar 2018 13:23:44 +0200 Subject: [PATCH] rpc: enforce the 128KB request limits on websockets too --- rpc/http.go | 10 +++++----- rpc/http_test.go | 2 +- rpc/json.go | 46 +++++++++++++++++++++++++++++++--------------- rpc/websocket.go | 30 +++++++++++++++++++++++++++++- 4 files changed, 66 insertions(+), 22 deletions(-) diff --git a/rpc/http.go b/rpc/http.go index a46d8c2b3..9805d69b6 100644 --- a/rpc/http.go +++ b/rpc/http.go @@ -27,16 +27,16 @@ import ( "mime" "net" "net/http" + "strings" "sync" "time" "github.com/rs/cors" - "strings" ) const ( - contentType = "application/json" - maxHTTPRequestContentLength = 1024 * 128 + contentType = "application/json" + maxRequestContentLength = 1024 * 128 ) var nullAddr, _ = net.ResolveTCPAddr("tcp", "127.0.0.1:0") @@ -182,8 +182,8 @@ func validateRequest(r *http.Request) (int, error) { if r.Method == http.MethodPut || r.Method == http.MethodDelete { return http.StatusMethodNotAllowed, errors.New("method not allowed") } - if r.ContentLength > maxHTTPRequestContentLength { - err := fmt.Errorf("content length too large (%d>%d)", r.ContentLength, maxHTTPRequestContentLength) + if r.ContentLength > maxRequestContentLength { + err := fmt.Errorf("content length too large (%d>%d)", r.ContentLength, maxRequestContentLength) return http.StatusRequestEntityTooLarge, err } mt, _, err := mime.ParseMediaType(r.Header.Get("content-type")) diff --git a/rpc/http_test.go b/rpc/http_test.go index aed84f683..b3f694d8a 100644 --- a/rpc/http_test.go +++ b/rpc/http_test.go @@ -32,7 +32,7 @@ func TestHTTPErrorResponseWithPut(t *testing.T) { } func TestHTTPErrorResponseWithMaxContentLength(t *testing.T) { - body := make([]rune, maxHTTPRequestContentLength+1) + body := make([]rune, maxRequestContentLength+1) testHTTPErrorResponse(t, http.MethodPost, contentType, string(body), http.StatusRequestEntityTooLarge) } diff --git a/rpc/json.go b/rpc/json.go index 2e7fd599e..837011f51 100644 --- a/rpc/json.go +++ b/rpc/json.go @@ -76,13 +76,13 @@ type jsonNotification struct { // jsonCodec reads and writes JSON-RPC messages to the underlying connection. It // also has support for parsing arguments and serializing (result) objects. type jsonCodec struct { - closer sync.Once // close closed channel once - closed chan interface{} // closed on Close - decMu sync.Mutex // guards d - d *json.Decoder // decodes incoming requests - encMu sync.Mutex // guards e - e *json.Encoder // encodes responses - rw io.ReadWriteCloser // connection + closer sync.Once // close closed channel once + closed chan interface{} // closed on Close + decMu sync.Mutex // guards the decoder + decode func(v interface{}) error // decoder to allow multiple transports + encMu sync.Mutex // guards the encoder + encode func(v interface{}) error // encoder to allow multiple transports + rw io.ReadWriteCloser // connection } func (err *jsonError) Error() string { @@ -96,11 +96,29 @@ func (err *jsonError) ErrorCode() int { return err.Code } -// NewJSONCodec creates a new RPC server codec with support for JSON-RPC 2.0 +// NewCodec creates a new RPC server codec with support for JSON-RPC 2.0 based +// on explicitly given encoding and decoding methods. +func NewCodec(rwc io.ReadWriteCloser, encode, decode func(v interface{}) error) ServerCodec { + return &jsonCodec{ + closed: make(chan interface{}), + encode: encode, + decode: decode, + rw: rwc, + } +} + +// NewJSONCodec creates a new RPC server codec with support for JSON-RPC 2.0. func NewJSONCodec(rwc io.ReadWriteCloser) ServerCodec { - d := json.NewDecoder(rwc) - d.UseNumber() - return &jsonCodec{closed: make(chan interface{}), d: d, e: json.NewEncoder(rwc), rw: rwc} + enc := json.NewEncoder(rwc) + dec := json.NewDecoder(rwc) + dec.UseNumber() + + return &jsonCodec{ + closed: make(chan interface{}), + encode: enc.Encode, + decode: dec.Decode, + rw: rwc, + } } // isBatch returns true when the first non-whitespace characters is '[' @@ -123,14 +141,12 @@ func (c *jsonCodec) ReadRequestHeaders() ([]rpcRequest, bool, Error) { defer c.decMu.Unlock() var incomingMsg json.RawMessage - if err := c.d.Decode(&incomingMsg); err != nil { + if err := c.decode(&incomingMsg); err != nil { return nil, false, &invalidRequestError{err.Error()} } - if isBatch(incomingMsg) { return parseBatchRequest(incomingMsg) } - return parseRequest(incomingMsg) } @@ -338,7 +354,7 @@ func (c *jsonCodec) Write(res interface{}) error { c.encMu.Lock() defer c.encMu.Unlock() - return c.e.Encode(res) + return c.encode(res) } // Close the underlying connection diff --git a/rpc/websocket.go b/rpc/websocket.go index 4214fc86a..a6e1cec28 100644 --- a/rpc/websocket.go +++ b/rpc/websocket.go @@ -17,8 +17,10 @@ package rpc import ( + "bytes" "context" "crypto/tls" + "encoding/json" "fmt" "net" "net/http" @@ -32,6 +34,23 @@ import ( "gopkg.in/fatih/set.v0" ) +// websocketJSONCodec is a custom JSON codec with payload size enforcement and +// special number parsing. +var websocketJSONCodec = websocket.Codec{ + // Marshal is the stock JSON marshaller used by the websocket library too. + Marshal: func(v interface{}) ([]byte, byte, error) { + msg, err := json.Marshal(v) + return msg, websocket.TextFrame, err + }, + // Unmarshal is a specialized unmarshaller to properly convert numbers. + Unmarshal: func(msg []byte, payloadType byte, v interface{}) error { + dec := json.NewDecoder(bytes.NewReader(msg)) + dec.UseNumber() + + return dec.Decode(v) + }, +} + // WebsocketHandler returns a handler that serves JSON-RPC to WebSocket connections. // // allowedOrigins should be a comma-separated list of allowed origin URLs. @@ -40,7 +59,16 @@ func (srv *Server) WebsocketHandler(allowedOrigins []string) http.Handler { return websocket.Server{ Handshake: wsHandshakeValidator(allowedOrigins), Handler: func(conn *websocket.Conn) { - srv.ServeCodec(NewJSONCodec(conn), OptionMethodInvocation|OptionSubscriptions) + // Create a custom encode/decode pair to enforce payload size and number encoding + conn.MaxPayloadBytes = maxRequestContentLength + + encoder := func(v interface{}) error { + return websocketJSONCodec.Send(conn, v) + } + decoder := func(v interface{}) error { + return websocketJSONCodec.Receive(conn, v) + } + srv.ServeCodec(NewCodec(conn, encoder, decoder), OptionMethodInvocation|OptionSubscriptions) }, } }