k3s/vendor/github.com/rancher/remotedialer/client.go
2019-07-03 13:11:54 -07:00

59 lines
1.4 KiB
Go

package remotedialer
import (
"context"
"net/http"
"time"
"github.com/gorilla/websocket"
"github.com/sirupsen/logrus"
)
type ConnectAuthorizer func(proto, address string) bool
func ClientConnect(ctx context.Context, wsURL string, headers http.Header, dialer *websocket.Dialer, auth ConnectAuthorizer, onConnect func(context.Context) error) {
if err := connectToProxy(ctx, wsURL, headers, auth, dialer, onConnect); err != nil {
logrus.WithError(err).Error("Remotedialer proxy error")
time.Sleep(time.Duration(5) * time.Second)
}
}
func connectToProxy(ctx context.Context, proxyURL string, headers http.Header, auth ConnectAuthorizer, dialer *websocket.Dialer, onConnect func(context.Context) error) error {
logrus.WithField("url", proxyURL).Info("Connecting to proxy")
if dialer == nil {
dialer = &websocket.Dialer{}
}
ws, _, err := dialer.Dial(proxyURL, headers)
if err != nil {
logrus.WithError(err).Error("Failed to connect to proxy")
return err
}
defer ws.Close()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
if onConnect != nil {
if err := onConnect(ctx); err != nil {
return err
}
}
session := NewClientSession(auth, ws)
defer session.Close()
result := make(chan error, 1)
go func() {
_, err = session.Serve()
result <- err
}()
select {
case <-ctx.Done():
return ctx.Err()
case err := <-result:
return err
}
}