k3s/vendor/github.com/rancher/remotedialer/client.go

60 lines
1.4 KiB
Go
Raw Normal View History

2019-01-12 04:58:27 +00:00
package remotedialer
import (
"context"
"net/http"
"time"
"github.com/gorilla/websocket"
"github.com/sirupsen/logrus"
)
type ConnectAuthorizer func(proto, address string) bool
2019-07-03 17:59:29 +00:00
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")
2019-01-12 04:58:27 +00:00
time.Sleep(time.Duration(5) * time.Second)
}
}
2019-07-18 11:58:06 +00:00
func connectToProxy(rootCtx context.Context, proxyURL string, headers http.Header, auth ConnectAuthorizer, dialer *websocket.Dialer, onConnect func(context.Context) error) error {
2019-01-12 04:58:27 +00:00
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()
2019-07-18 11:58:06 +00:00
ctx, cancel := context.WithCancel(rootCtx)
defer cancel()
2019-07-10 21:49:56 +00:00
2019-07-18 11:58:06 +00:00
if onConnect != nil {
if err := onConnect(ctx); err != nil {
2019-01-12 04:58:27 +00:00
return err
}
}
session := NewClientSession(auth, ws)
2019-07-03 17:59:29 +00:00
defer session.Close()
result := make(chan error, 1)
go func() {
2019-07-18 11:58:06 +00:00
_, err = session.Serve(ctx)
2019-07-03 17:59:29 +00:00
result <- err
}()
select {
case <-ctx.Done():
2019-07-10 21:49:56 +00:00
logrus.WithField("url", proxyURL).WithField("err", ctx.Err()).Info("Proxy done")
return nil
2019-07-03 17:59:29 +00:00
case err := <-result:
return err
}
2019-01-12 04:58:27 +00:00
}