-
Notifications
You must be signed in to change notification settings - Fork 235
/
connect.go
150 lines (124 loc) · 3.6 KB
/
connect.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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
package proxy
import (
"context"
"fmt"
"net"
"strconv"
"github.com/AlecAivazis/survey/v2"
"github.com/superfly/flyctl/agent"
"github.com/superfly/flyctl/client"
"github.com/superfly/flyctl/iostreams"
"github.com/superfly/flyctl/ip"
)
type ConnectParams struct {
AppName string
OrganizationSlug string
Dialer agent.Dialer
BindAddr string
Ports []string
RemoteHost string
PromptInstance bool
DisableSpinner bool
}
// Binds to a local port and runs a proxy to a remote address over Wireguard.
// Blocks until context is cancelled.
func Connect(ctx context.Context, p *ConnectParams) (err error) {
server, err := NewServer(ctx, p)
if err != nil {
return err
}
return server.ProxyServer(ctx)
}
// Binds to a local port and then starts a goroutine to run a proxy to a remote
// address over Wireguard. Proxy runs until context is cancelled.
// Blocks only until local listener is bound and ready to accept connections.
func Start(ctx context.Context, p *ConnectParams) error {
server, err := NewServer(ctx, p)
if err != nil {
return err
}
// currently ignores any error returned by ProxyServer
// TODO return a channel to caller for async error notification
go server.ProxyServer(ctx)
return nil
}
func NewServer(ctx context.Context, p *ConnectParams) (*Server, error) {
var (
io = iostreams.FromContext(ctx)
client = client.FromContext(ctx).API()
orgSlug = p.OrganizationSlug
localBindAddr = p.BindAddr
localPort = p.Ports[0]
remotePort = localPort
remoteAddr string
)
if len(p.Ports) > 1 {
remotePort = p.Ports[1]
}
agentclient, err := agent.Establish(ctx, client)
if err != nil {
return nil, err
}
// Prompt for a specific instance and set it as the remote target
if p.PromptInstance {
instance, err := selectInstance(ctx, p.OrganizationSlug, p.AppName, agentclient)
if err != nil {
return nil, err
}
remoteAddr = fmt.Sprintf("[%s]:%s", instance, remotePort)
}
if remoteAddr == "" && p.RemoteHost != "" {
// If a host is specified that isn't an IpV6 address, assume it's a DNS entry and wait for that
// entry to resolve
if !ip.IsV6(p.RemoteHost) {
if err := agentclient.WaitForDNS(ctx, p.Dialer, orgSlug, p.RemoteHost); err != nil {
return nil, fmt.Errorf("%s: %w", p.RemoteHost, err)
}
}
remoteAddr = fmt.Sprintf("[%s]:%s", p.RemoteHost, remotePort)
}
var listener net.Listener
if _, err := strconv.Atoi(localPort); err == nil {
// just numbers
addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%s", localBindAddr, localPort))
if err != nil {
return nil, err
}
listener, err = net.ListenTCP("tcp", addr)
if err != nil {
return nil, err
}
} else {
// probably a unix path
addr, err := net.ResolveUnixAddr("unix", localPort)
if err != nil {
return nil, err
}
listener, err = net.ListenUnix("unix", addr)
if err != nil {
return nil, err
}
}
fmt.Fprintf(io.Out, "Proxying local port %s to remote %s\n", localPort, remoteAddr)
return &Server{
Addr: remoteAddr,
Listener: listener,
Dial: p.Dialer.DialContext,
}, nil
}
func selectInstance(ctx context.Context, org, app string, c *agent.Client) (instance string, err error) {
instances, err := c.Instances(ctx, org, app)
if err != nil {
return "", fmt.Errorf("look up %s: %w", app, err)
}
selected := 0
prompt := &survey.Select{
Message: "Select instance:",
Options: instances.Labels,
PageSize: 15,
}
if err := survey.AskOne(prompt, &selected); err != nil {
return "", fmt.Errorf("selecting instance: %w", err)
}
return instances.Addresses[selected], nil
}