generated from brevdev/seed
-
Notifications
You must be signed in to change notification settings - Fork 9
/
huproxyclient.go
120 lines (101 loc) · 2.83 KB
/
huproxyclient.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
package huproxyclient
// https://github.com/google/huproxy/blob/master/huproxyclient/client.go
import (
"context"
"crypto/tls"
"fmt"
"io"
"io/ioutil"
"net/http"
"os"
"time"
"github.com/brevdev/brev-cli/pkg/entity"
"github.com/brevdev/brev-cli/pkg/errors"
"github.com/gorilla/websocket"
log "github.com/sirupsen/logrus"
huproxy "github.com/google/huproxy/lib"
)
var writeTimeout = 10 * time.Second
type HubProxyStore interface {
GetAuthTokens() (*entity.AuthTokens, error)
GetCurrentWorkspaceGroupID() (string, error)
}
func dialError(url string, resp *http.Response, err error) {
if resp != nil {
extra := ""
b, err1 := ioutil.ReadAll(resp.Body)
if err1 != nil {
log.Warningf("Failed to read HTTP body: %v", err1)
}
extra = "Body:\n" + string(b)
log.Fatalf("%s: HTTP error: %d %s\n%s", err, resp.StatusCode, resp.Status, extra)
}
log.Fatalf("Dial to %q fail: %v", url, err)
}
func Run(url string, store HubProxyStore) error {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
dialer := websocket.Dialer{}
dialer.TLSClientConfig = new(tls.Config)
head := map[string][]string{}
token, err := store.GetAuthTokens()
if err != nil {
return errors.WrapAndTrace(err)
}
workspaceGroupID, err := store.GetCurrentWorkspaceGroupID()
if err != nil {
fmt.Printf("%v\n", err)
}
if workspaceGroupID != "" {
head["X-Workspace-Group-ID"] = []string{workspaceGroupID}
}
head["Authorization"] = []string{
"Bearer " + token.AccessToken,
}
conn, resp, err := dialer.Dial(url, head)
if err != nil {
dialError(url, resp, err)
}
defer conn.Close() //nolint:errcheck // lazy to refactor
RunProxy(ctx, conn, cancel)
if ctx.Err() != nil {
return errors.WrapAndTrace(ctx.Err())
}
return nil
}
func RunProxy(ctx context.Context, conn *websocket.Conn, cancel context.CancelFunc) {
// websocket -> stdout
go func() {
for {
mt, r, err := conn.NextReader()
if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
return
}
if err != nil {
log.Warn("Workspace disconnect: may be from network failure or workspace was stopped/deleted")
log.Fatal(err)
}
if mt != websocket.BinaryMessage {
log.Fatal("non-binary websocket message received")
}
if _, err := io.Copy(os.Stdout, r); err != nil {
log.Errorf("Reading from websocket: %v", err)
cancel()
}
}
}()
// stdin -> websocket
// TODO: NextWriter() seems to be broken.
if err := huproxy.File2WS(ctx, cancel, os.Stdin, conn); err == io.EOF {
if err1 := conn.WriteControl(websocket.CloseMessage,
websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
time.Now().Add(writeTimeout)); err1 == websocket.ErrCloseSent {
_ = ""
} else if err1 != nil {
log.Errorf("Error sending 'close' message: %v", err1)
}
} else if err != nil {
log.Errorf("reading from stdin: %v", err)
cancel()
}
}