Merge pull request checksum0#2 from keep-network/websocket

Support WebSocket protocol

electrum/network.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"net/url"
 	"sync"
 	"sync/atomic"
 )
@@ -70,6 +71,28 @@ type Client struct {
 	nextID uint64
 }
 
+// NewClient initializes a new client for remote server and connects to it using
+// a transport protocol resolved from the URL's protocol scheme.
+// A remote server URL should be provided in the `scheme://hostname:port` format
+// (e.g. `tcp://electrum.io:50001`).
+func NewClient(ctx context.Context, urlStr string, tlsConfig *tls.Config) (*Client, error) {
+	u, err := url.Parse(urlStr)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse url [%s]: [%w]", urlStr, err)
+	}
+
+	switch u.Scheme {
+	case "tcp":
+		return NewClientTCP(ctx, u.Host)
+	case "ssl":
+		return NewClientSSL(ctx, u.Host, tlsConfig)
+	case "ws", "wss":
+		return NewClientWebSocket(ctx, u.String(), tlsConfig)
+	}
+
+	return nil, fmt.Errorf("unsupported protocol scheme: [%s]", u.Scheme)
+}
+
 // NewClientTCP initialize a new client for remote server and connects to the remote server using TCP
 func NewClientTCP(ctx context.Context, addr string) (*Client, error) {
 	transport, err := NewTCPTransport(ctx, addr)
@@ -112,6 +135,28 @@ func NewClientSSL(ctx context.Context, addr string, config *tls.Config) (*Client
 	return c, nil
 }
 
+// NewClientWebSocket initialize a new client for remote server and connects to
+// the remote server using WebSocket.
+func NewClientWebSocket(ctx context.Context, url string, config *tls.Config) (*Client, error) {
+	transport, err := NewWebSocketTransport(ctx, url, config)
+	if err != nil {
+		return nil, err
+	}
+
+	c := &Client{
+		handlers:     make(map[uint64]chan *container),
+		pushHandlers: make(map[string][]chan *container),
+
+		Error: make(chan error),
+		quit:  make(chan struct{}),
+	}
+
+	c.transport = transport
+	go c.listen()
+
+	return c, nil
+}
+
 // JSON-RPC 2.0 Error Object
 // See: https://www.jsonrpc.org/specificationJSON#error_object
 type apiErr struct {

electrum/transport_ws.go

@@ -0,0 +1,120 @@
+package electrum
+
+import (
+	"context"
+	"crypto/tls"
+	"log"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+type WebSocketTransport struct {
+	conn      *websocket.Conn
+	responses chan []byte
+	errors    chan error
+	// close is a channel used for graceful connection closure
+	close chan struct{}
+}
+
+const webSocketClosingTimeout = 2 * time.Second
+
+// NewWebSocketTransport initializes new WebSocket transport.
+func NewWebSocketTransport(
+	ctx context.Context,
+	url string,
+	tlsConfig *tls.Config,
+) (*WebSocketTransport, error) {
+	dialer := websocket.Dialer{
+		TLSClientConfig: tlsConfig,
+	}
+
+	conn, response, err := dialer.DialContext(ctx, url, nil)
+	if err != nil {
+		if DebugMode {
+			log.Printf(
+				"%s [debug] connect -> status: %v, error: %v",
+				time.Now().Format("2006-01-02 15:04:05"),
+				response.Status,
+				err,
+			)
+		}
+		return nil, err
+	}
+
+	ws := &WebSocketTransport{
+		conn:      conn,
+		responses: make(chan []byte),
+		errors:    make(chan error),
+		close:     make(chan struct{}),
+	}
+
+	go ws.listen()
+
+	return ws, nil
+}
+
+func (t *WebSocketTransport) listen() {
+	defer t.conn.Close()
+	defer close(t.close)
+
+	for {
+		_, msg, err := t.conn.ReadMessage()
+		if DebugMode {
+			log.Printf(
+				"%s [debug] %s -> msg: %s, err: %v",
+				time.Now().Format("2006-01-02 15:04:05"),
+				t.conn.RemoteAddr(),
+				msg,
+				err,
+			)
+		}
+		if err != nil {
+			isNormalClose := websocket.IsCloseError(err, websocket.CloseNormalClosure)
+			if !isNormalClose {
+				t.errors <- err
+			}
+
+			break
+		}
+
+		t.responses <- msg
+	}
+}
+
+// SendMessage sends a message to the remote server through the WebSocket transport.
+func (t *WebSocketTransport) SendMessage(body []byte) error {
+	if DebugMode {
+		log.Printf("%s [debug] %s <- %s", time.Now().Format("2006-01-02 15:04:05"), t.conn.RemoteAddr(), body)
+	}
+
+	return t.conn.WriteMessage(websocket.TextMessage, body)
+}
+
+// Responses returns chan to WebSocket transport responses.
+func (t *WebSocketTransport) Responses() <-chan []byte {
+	return t.responses
+}
+
+// Errors returns chan to WebSocket transport errors.
+func (t *WebSocketTransport) Errors() <-chan error {
+	return t.errors
+}
+
+// Close closes WebSocket transport.
+func (t *WebSocketTransport) Close() error {
+	// Cleanly close the connection by sending a close message and then
+	// waiting (with timeout) for the server to close the connection.
+	err := t.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+	if err != nil {
+		log.Printf("%s [error] %s -> close error: %s", time.Now().Format("2006-01-02 15:04:05"), t.conn.RemoteAddr(), err)
+	}
+
+	select {
+	case <-t.close:
+	case <-time.After(webSocketClosingTimeout):
+		return t.conn.Close()
+	}
+
+	return nil
+}

go.mod

@@ -5,6 +5,7 @@ go 1.18
 require (
 	github.com/btcsuite/btcd v0.23.1
 	github.com/btcsuite/btcd/btcutil v1.1.1
+	github.com/gorilla/websocket v1.5.0
 	github.com/stretchr/testify v1.7.0
 )
 

go.sum

@@ -43,6 +43,8 @@ github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw
 github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
+github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/jessevdk/go-flags v0.0.0-20141203071132-1679536dcc89/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=
 github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=