1 files changed, 17 insertions, 2 deletions
diff --git a/rpc/websocket.go b/rpc/websocket.go
index f4271fda8..587010820 100644
--- a/rpc/websocket.go
+++ b/rpc/websocket.go
@@ -17,6 +17,7 @@
 package rpc
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"net"
@@ -24,10 +25,9 @@ import (
 	"net/url"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/ethereum/go-ethereum/log"
-
-	"golang.org/x/net/context"
 	"golang.org/x/net/websocket"
 	"gopkg.in/fatih/set.v0"
 )
@@ -150,3 +150,18 @@ func wsDialAddress(location *url.URL) string {
 	}
 	return location.Host
 }
+
+func dialContext(ctx context.Context, network, addr string) (net.Conn, error) {
+	d := &net.Dialer{KeepAlive: tcpKeepAliveInterval}
+	return d.DialContext(ctx, network, addr)
+}
+
+func contextDialer(ctx context.Context) *net.Dialer {
+	dialer := &net.Dialer{Cancel: ctx.Done(), KeepAlive: tcpKeepAliveInterval}
+	if deadline, ok := ctx.Deadline(); ok {
+		dialer.Deadline = deadline
+	} else {
+		dialer.Deadline = time.Now().Add(defaultDialTimeout)
+	}
+	return dialer
+}