Feat/headers adds header functionality to http and websockets (umbrac…

…le#140) * added in header functionality * draft implementation for websockets * minor formatting changes Co-authored-by: Ben Wilson <[email protected]>
cwarden · Jan 12, 2022 · a8fbcb8 · a8fbcb8
1 parent 013846a
commit a8fbcb8
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 9 deletions.
diff --git a/jsonrpc/client.go b/jsonrpc/client.go
@@ -16,24 +16,41 @@ type endpoints struct {
 	n *Net
 	d *Debug
 }
+type Config struct {
+	headers map[string]string
+}
+
+type ConfigOption func(*Config)
+
+func WithHeaders(headers map[string]string) ConfigOption {
+	return func(c *Config) {
+		for k, v := range headers {
+			c.headers[k] = v
+		}
+	}
+}
+
+func NewClient(addr string, opts ...ConfigOption) (*Client, error) {
+	config := &Config{headers: map[string]string{}}
+	for _, opt := range opts {
+		opt(config)
+	}
 
-// NewClient creates a new client
-func NewClient(addr string) (*Client, error) {
 	c := &Client{}
 	c.endpoints.w = &Web3{c}
 	c.endpoints.e = &Eth{c}
 	c.endpoints.n = &Net{c}
 	c.endpoints.d = &Debug{c}
 
-	t, err := transport.NewTransport(addr)
+	t, err := transport.NewTransport(addr, config.headers)
 	if err != nil {
 		return nil, err
 	}
 	c.transport = t
 	return c, nil
 }
 
-// Close closes the tranport
+// Close closes the transport
 func (c *Client) Close() error {
 	return c.transport.Close()
 }

diff --git a/jsonrpc/transport/http.go b/jsonrpc/transport/http.go
@@ -11,12 +11,14 @@ import (
 type HTTP struct {
 	addr   string
 	client *fasthttp.Client
+	headers map[string]string
 }
 
-func newHTTP(addr string) *HTTP {
+func newHTTP(addr string, headers map[string]string) *HTTP {
 	return &HTTP{
 		addr:   addr,
 		client: &fasthttp.Client{},
+		headers: headers,
 	}
 }
 
@@ -53,6 +55,9 @@ func (h *HTTP) Call(method string, out interface{}, params ...interface{}) error
 	req.SetRequestURI(h.addr)
 	req.Header.SetMethod("POST")
 	req.Header.SetContentType("application/json")
+	for k, v := range h.headers {
+		req.Header.Add(k, v)
+	}
 	req.SetBody(raw)
 
 	if err := h.client.Do(req, res); err != nil {

diff --git a/jsonrpc/transport/transport.go b/jsonrpc/transport/transport.go
@@ -29,9 +29,9 @@ const (
 )
 
 // NewTransport creates a new transport object
-func NewTransport(url string) (Transport, error) {
+func NewTransport(url string, headers map[string]string) (Transport, error) {
 	if strings.HasPrefix(url, wsPrefix) || strings.HasPrefix(url, wssPrefix) {
-		t, err := newWebsocket(url)
+		t, err := newWebsocket(url, headers)
 		if err != nil {
 			return nil, err
 		}
@@ -45,5 +45,5 @@ func NewTransport(url string) (Transport, error) {
 		}
 		return t, nil
 	}
-	return newHTTP(url), nil
+	return newHTTP(url, headers), nil
 }
diff --git a/jsonrpc/transport/websocket.go b/jsonrpc/transport/websocket.go
@@ -12,7 +12,11 @@ import (
 	"github.com/umbracle/go-web3/jsonrpc/codec"
 )
 
-func newWebsocket(url string) (Transport, error) {
+func newWebsocket(url string, headers map[string]string) (Transport, error) {
+	wsHeaders := http.Header{}
+	for k, v := range headers {
+		wsHeaders.Add(k, v)
+	}
 	wsConn, _, err := websocket.DefaultDialer.Dial(url, http.Header{})
 	if err != nil {
 		return nil, err