From 129bbd3ca3f56f128d1c657159263fc326b459a3 Mon Sep 17 00:00:00 2001
From: Yuki Nishiwaki <uckey.1067@gmail.com>
Date: Wed, 26 Jun 2019 00:55:50 +0900
Subject: [PATCH] Configure HandshakeTimeout in websocket.Dialer

To prevent agent from being hung in the middle of Websocket Session
Handshake, We'd better to set HandshakeTimeout

Related to https://github.com/rancher/rancher/issues/21117
---
 clientbase/common.go       | 2 +-
 pkg/remotedialer/client.go | 2 +-
 pkg/remotedialer/peer.go   | 1 +
 3 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/clientbase/common.go b/clientbase/common.go
index e1ab8f93..8107952f 100644
--- a/clientbase/common.go
+++ b/clientbase/common.go
@@ -277,7 +277,7 @@ func NewAPIClient(opts *ClientOpts) (APIBaseClient, error) {
 	result.Ops = &APIOperations{
 		Opts:   opts,
 		Client: client,
-		Dialer: &websocket.Dialer{},
+		Dialer: &websocket.Dialer{HandshakeTimeout: 10 * time.Second},
 		Types:  result.Types,
 	}
 
diff --git a/pkg/remotedialer/client.go b/pkg/remotedialer/client.go
index 86dacd11..b53d0aed 100644
--- a/pkg/remotedialer/client.go
+++ b/pkg/remotedialer/client.go
@@ -22,7 +22,7 @@ func connectToProxy(proxyURL string, headers http.Header, auth ConnectAuthorizer
 	logrus.WithField("url", proxyURL).Info("Connecting to proxy")
 
 	if dialer == nil {
-		dialer = &websocket.Dialer{}
+		dialer = &websocket.Dialer{HandshakeTimeout: 10 * time.Second}
 	}
 	ws, _, err := dialer.Dial(proxyURL, headers)
 	if err != nil {
diff --git a/pkg/remotedialer/peer.go b/pkg/remotedialer/peer.go
index 17de7d7f..452ea57d 100644
--- a/pkg/remotedialer/peer.go
+++ b/pkg/remotedialer/peer.go
@@ -79,6 +79,7 @@ func (p *peer) start(ctx context.Context, s *Server) {
 		TLSClientConfig: &tls.Config{
 			InsecureSkipVerify: true,
 		},
+		HandshakeTimeout: 10 * time.Second,
 	}
 
 outer: