github.com/luckypickle/go-ethereum-vet@v1.14.2/rpc/websocket.go (about)

     1  // Copyright 2015 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package rpc
    18  
    19  import (
    20  	"bytes"
    21  	"context"
    22  	"crypto/tls"
    23  	"encoding/json"
    24  	"fmt"
    25  	"net"
    26  	"net/http"
    27  	"net/url"
    28  	"os"
    29  	"strings"
    30  	"time"
    31  
    32  	mapset "github.com/deckarep/golang-set"
    33  	"github.com/luckypickle/go-ethereum-vet/log"
    34  	"golang.org/x/net/websocket"
    35  )
    36  
    37  // websocketJSONCodec is a custom JSON codec with payload size enforcement and
    38  // special number parsing.
    39  var websocketJSONCodec = websocket.Codec{
    40  	// Marshal is the stock JSON marshaller used by the websocket library too.
    41  	Marshal: func(v interface{}) ([]byte, byte, error) {
    42  		msg, err := json.Marshal(v)
    43  		return msg, websocket.TextFrame, err
    44  	},
    45  	// Unmarshal is a specialized unmarshaller to properly convert numbers.
    46  	Unmarshal: func(msg []byte, payloadType byte, v interface{}) error {
    47  		dec := json.NewDecoder(bytes.NewReader(msg))
    48  		dec.UseNumber()
    49  
    50  		return dec.Decode(v)
    51  	},
    52  }
    53  
    54  // WebsocketHandler returns a handler that serves JSON-RPC to WebSocket connections.
    55  //
    56  // allowedOrigins should be a comma-separated list of allowed origin URLs.
    57  // To allow connections with any origin, pass "*".
    58  func (srv *Server) WebsocketHandler(allowedOrigins []string) http.Handler {
    59  	return websocket.Server{
    60  		Handshake: wsHandshakeValidator(allowedOrigins),
    61  		Handler: func(conn *websocket.Conn) {
    62  			// Create a custom encode/decode pair to enforce payload size and number encoding
    63  			conn.MaxPayloadBytes = maxRequestContentLength
    64  
    65  			encoder := func(v interface{}) error {
    66  				return websocketJSONCodec.Send(conn, v)
    67  			}
    68  			decoder := func(v interface{}) error {
    69  				return websocketJSONCodec.Receive(conn, v)
    70  			}
    71  			srv.ServeCodec(NewCodec(conn, encoder, decoder), OptionMethodInvocation|OptionSubscriptions)
    72  		},
    73  	}
    74  }
    75  
    76  // NewWSServer creates a new websocket RPC server around an API provider.
    77  //
    78  // Deprecated: use Server.WebsocketHandler
    79  func NewWSServer(allowedOrigins []string, srv *Server) *http.Server {
    80  	return &http.Server{Handler: srv.WebsocketHandler(allowedOrigins)}
    81  }
    82  
    83  // wsHandshakeValidator returns a handler that verifies the origin during the
    84  // websocket upgrade process. When a '*' is specified as an allowed origins all
    85  // connections are accepted.
    86  func wsHandshakeValidator(allowedOrigins []string) func(*websocket.Config, *http.Request) error {
    87  	origins := mapset.NewSet()
    88  	allowAllOrigins := false
    89  
    90  	for _, origin := range allowedOrigins {
    91  		if origin == "*" {
    92  			allowAllOrigins = true
    93  		}
    94  		if origin != "" {
    95  			origins.Add(strings.ToLower(origin))
    96  		}
    97  	}
    98  
    99  	// allow localhost if no allowedOrigins are specified.
   100  	if len(origins.ToSlice()) == 0 {
   101  		origins.Add("http://localhost")
   102  		if hostname, err := os.Hostname(); err == nil {
   103  			origins.Add("http://" + strings.ToLower(hostname))
   104  		}
   105  	}
   106  
   107  	log.Debug(fmt.Sprintf("Allowed origin(s) for WS RPC interface %v\n", origins.ToSlice()))
   108  
   109  	f := func(cfg *websocket.Config, req *http.Request) error {
   110  		origin := strings.ToLower(req.Header.Get("Origin"))
   111  		if allowAllOrigins || origins.Contains(origin) {
   112  			return nil
   113  		}
   114  		log.Warn(fmt.Sprintf("origin '%s' not allowed on WS-RPC interface\n", origin))
   115  		return fmt.Errorf("origin %s not allowed", origin)
   116  	}
   117  
   118  	return f
   119  }
   120  
   121  // DialWebsocket creates a new RPC client that communicates with a JSON-RPC server
   122  // that is listening on the given endpoint.
   123  //
   124  // The context is used for the initial connection establishment. It does not
   125  // affect subsequent interactions with the client.
   126  func DialWebsocket(ctx context.Context, endpoint, origin string) (*Client, error) {
   127  	if origin == "" {
   128  		var err error
   129  		if origin, err = os.Hostname(); err != nil {
   130  			return nil, err
   131  		}
   132  		if strings.HasPrefix(endpoint, "wss") {
   133  			origin = "https://" + strings.ToLower(origin)
   134  		} else {
   135  			origin = "http://" + strings.ToLower(origin)
   136  		}
   137  	}
   138  	config, err := websocket.NewConfig(endpoint, origin)
   139  	if err != nil {
   140  		return nil, err
   141  	}
   142  
   143  	return newClient(ctx, func(ctx context.Context) (net.Conn, error) {
   144  		return wsDialContext(ctx, config)
   145  	})
   146  }
   147  
   148  func wsDialContext(ctx context.Context, config *websocket.Config) (*websocket.Conn, error) {
   149  	var conn net.Conn
   150  	var err error
   151  	switch config.Location.Scheme {
   152  	case "ws":
   153  		conn, err = dialContext(ctx, "tcp", wsDialAddress(config.Location))
   154  	case "wss":
   155  		dialer := contextDialer(ctx)
   156  		conn, err = tls.DialWithDialer(dialer, "tcp", wsDialAddress(config.Location), config.TlsConfig)
   157  	default:
   158  		err = websocket.ErrBadScheme
   159  	}
   160  	if err != nil {
   161  		return nil, err
   162  	}
   163  	ws, err := websocket.NewClient(config, conn)
   164  	if err != nil {
   165  		conn.Close()
   166  		return nil, err
   167  	}
   168  	return ws, err
   169  }
   170  
   171  var wsPortMap = map[string]string{"ws": "80", "wss": "443"}
   172  
   173  func wsDialAddress(location *url.URL) string {
   174  	if _, ok := wsPortMap[location.Scheme]; ok {
   175  		if _, _, err := net.SplitHostPort(location.Host); err != nil {
   176  			return net.JoinHostPort(location.Host, wsPortMap[location.Scheme])
   177  		}
   178  	}
   179  	return location.Host
   180  }
   181  
   182  func dialContext(ctx context.Context, network, addr string) (net.Conn, error) {
   183  	d := &net.Dialer{KeepAlive: tcpKeepAliveInterval}
   184  	return d.DialContext(ctx, network, addr)
   185  }
   186  
   187  func contextDialer(ctx context.Context) *net.Dialer {
   188  	dialer := &net.Dialer{Cancel: ctx.Done(), KeepAlive: tcpKeepAliveInterval}
   189  	if deadline, ok := ctx.Deadline(); ok {
   190  		dialer.Deadline = deadline
   191  	} else {
   192  		dialer.Deadline = time.Now().Add(defaultDialTimeout)
   193  	}
   194  	return dialer
   195  }