Skip to content

Commit ff448db

Browse files
enocomenocom
committed
feat: add port flag and query param (#1139)
1 parent feee0d6 commit ff448db

File tree

4 files changed

+133
-20
lines changed

4 files changed

+133
-20
lines changed

‎cmd/root.go

Lines changed: 28 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"net/url"
2323
"os"
2424
"os/signal"
25+
"strconv"
2526
"strings"
2627
"syscall"
2728

@@ -76,6 +77,8 @@ any client SSL certificates.`,
7677

7778
cmd.PersistentFlags().StringVarP(&c.conf.Addr, "address", "a", "127.0.0.1",
7879
"Address on which to bind Cloud SQL instance listeners.")
80+
cmd.PersistentFlags().IntVarP(&c.conf.Port, "port", "p", 5000,
81+
"Initial port to use for listeners. Subsequent listeners increment from this value.")
7982

8083
c.Command = cmd
8184
return c
@@ -104,16 +107,33 @@ func parseConfig(conf *proxy.Config, args []string) error {
104107
if err != nil {
105108
return newBadCommandError(fmt.Sprintf("could not parse query: %q", res[1]))
106109
}
107-
if len(q["address"]) != 1 {
108-
return newBadCommandError(fmt.Sprintf("address query param should be only one value: %q", q["address"]))
110+
111+
if a, ok := q["address"]; ok {
112+
if len(a) != 1 {
113+
return newBadCommandError(fmt.Sprintf("address query param should be only one value: %q", a))
114+
}
115+
if ip := net.ParseIP(a[0]); ip == nil {
116+
return newBadCommandError(
117+
fmt.Sprintf("address query param is not a valid IP address: %q",
118+
a[0],
119+
))
120+
}
121+
ic.Addr = a[0]
109122
}
110-
if ip := net.ParseIP(q["address"][0]); ip == nil {
111-
return newBadCommandError(
112-
fmt.Sprintf("address query param is not a valid IP address: %q",
113-
q["address"][0],
114-
))
123+
124+
if p, ok := q["port"]; ok {
125+
if len(p) != 1 {
126+
return newBadCommandError(fmt.Sprintf("port query param should be only one value: %q", a))
127+
}
128+
pp, err := strconv.Atoi(p[0])
129+
if err != nil {
130+
return newBadCommandError(
131+
fmt.Sprintf("port query param is not a valid integer: %q",
132+
p[0],
133+
))
134+
}
135+
ic.Port = pp
115136
}
116-
ic.Addr = q["address"][0]
117137
}
118138
ics = append(ics, ic)
119139
}

‎cmd/root_test.go

Lines changed: 57 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,21 @@ import (
2323
)
2424

2525
func TestNewCommandArguments(t *testing.T) {
26+
withDefaults := func(c *proxy.Config) *proxy.Config {
27+
if c.Addr == "" {
28+
c.Addr = "127.0.0.1"
29+
}
30+
if c.Port == 0 {
31+
c.Port = 5000
32+
}
33+
if c.Instances == nil {
34+
c.Instances = []proxy.InstanceConnConfig{{}}
35+
}
36+
if i := &c.Instances[0]; i.Name == "" {
37+
i.Name = "proj:region:inst"
38+
}
39+
return c
40+
}
2641
tcs := []struct {
2742
desc string
2843
args []string
@@ -31,43 +46,69 @@ func TestNewCommandArguments(t *testing.T) {
3146
{
3247
desc: "basic invocation with defaults",
3348
args: []string{"proj:region:inst"},
34-
want: &proxy.Config{
49+
want: withDefaults(&proxy.Config{
3550
Addr: "127.0.0.1",
3651
Instances: []proxy.InstanceConnConfig{{Name: "proj:region:inst"}},
37-
},
52+
}),
3853
},
3954
{
4055
desc: "using the address flag",
4156
args: []string{"--address", "0.0.0.0", "proj:region:inst"},
42-
want: &proxy.Config{
57+
want: withDefaults(&proxy.Config{
4358
Addr: "0.0.0.0",
4459
Instances: []proxy.InstanceConnConfig{{Name: "proj:region:inst"}},
45-
},
60+
}),
4661
},
4762
{
4863
desc: "using the address (short) flag",
4964
args: []string{"-a", "0.0.0.0", "proj:region:inst"},
50-
want: &proxy.Config{
65+
want: withDefaults(&proxy.Config{
5166
Addr: "0.0.0.0",
5267
Instances: []proxy.InstanceConnConfig{{Name: "proj:region:inst"}},
53-
},
68+
}),
5469
},
5570
{
5671
desc: "using the address query param",
5772
args: []string{"proj:region:inst?address=0.0.0.0"},
58-
want: &proxy.Config{
73+
want: withDefaults(&proxy.Config{
5974
Addr: "127.0.0.1",
6075
Instances: []proxy.InstanceConnConfig{{
6176
Addr: "0.0.0.0",
6277
Name: "proj:region:inst",
6378
}},
64-
},
79+
}),
80+
},
81+
{
82+
desc: "using the port flag",
83+
args: []string{"--port", "6000", "proj:region:inst"},
84+
want: withDefaults(&proxy.Config{
85+
Port: 6000,
86+
}),
87+
},
88+
{
89+
desc: "using the port (short) flag",
90+
args: []string{"-p", "6000", "proj:region:inst"},
91+
want: withDefaults(&proxy.Config{
92+
Port: 6000,
93+
}),
94+
},
95+
{
96+
desc: "using the port query param",
97+
args: []string{"proj:region:inst?port=6000"},
98+
want: withDefaults(&proxy.Config{
99+
Instances: []proxy.InstanceConnConfig{{
100+
Port: 6000,
101+
}},
102+
}),
65103
},
66104
}
67105

68106
for _, tc := range tcs {
69107
t.Run(tc.desc, func(t *testing.T) {
70108
c := NewCommand()
109+
// Keep the test output quiet
110+
c.SilenceUsage = true
111+
c.SilenceErrors = true
71112
// Disable execute behavior
72113
c.RunE = func(*cobra.Command, []string) error {
73114
return nil
@@ -119,6 +160,14 @@ func TestNewCommandWithErrors(t *testing.T) {
119160
desc: "when the query string is invalid",
120161
args: []string{"proj:region:inst?address=1.1.1.1?foo=2.2.2.2"},
121162
},
163+
{
164+
desc: "when the port query param contains multiple values",
165+
args: []string{"proj:region:inst?port=1&port=2"},
166+
},
167+
{
168+
desc: "when the port query param is not a number",
169+
args: []string{"proj:region:inst?port=hi"},
170+
},
122171
}
123172

124173
for _, tc := range tcs {

‎internal/proxy/proxy.go

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,13 +33,19 @@ type InstanceConnConfig struct {
3333
Name string
3434
// Addr is the address on which to bind a listener for the instance.
3535
Addr string
36+
// Port is the port on which to bind a listener for the instance.
37+
Port int
3638
}
3739

3840
// Config contains all the configuration provided by the caller.
3941
type Config struct {
4042
// Addr is the address on which to bind all instances.
4143
Addr string
4244

45+
// Port is the first port to bind to. Subsequent ports will increment from
46+
// this value.
47+
Port int
48+
4349
// Instances are configuration for individual instances. Instance
4450
// configuration takes precedence over global configuration.
4551
Instances []InstanceConnConfig
@@ -62,17 +68,24 @@ func NewClient(ctx context.Context, cmd *cobra.Command, conf *Config) (*Client,
6268
}
6369
c := &Client{cmd: cmd, dialer: d}
6470

65-
port := 5000 // TODO: figure out better port allocation strategy
66-
for i, inst := range conf.Instances {
71+
port := conf.Port
72+
var inc int
73+
for _, inst := range conf.Instances {
6774
m := &socketMount{inst: inst.Name}
6875
a := conf.Addr
6976
if inst.Addr != "" {
7077
a = inst.Addr
7178
}
72-
addr, err := m.listen(ctx, "tcp", net.JoinHostPort(a, fmt.Sprint(port+i)))
79+
nextPort := port + inc
80+
if inst.Port != 0 {
81+
nextPort = inst.Port
82+
} else {
83+
inc++
84+
}
85+
addr, err := m.listen(ctx, "tcp", net.JoinHostPort(a, fmt.Sprint(nextPort)))
7386
if err != nil {
7487
c.Close()
75-
return nil, fmt.Errorf("[%s] Unable to mount socket: %v", inst, err)
88+
return nil, fmt.Errorf("[%v] Unable to mount socket: %v", inst.Name, err)
7689
}
7790
c.cmd.Printf("[%s] Listening on %s\n", inst.Name, addr.String())
7891
c.mnts = append(c.mnts, m)

‎internal/proxy/proxy_test.go

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ func TestClientInitialization(t *testing.T) {
3333
desc: "multiple instances",
3434
in: &proxy.Config{
3535
Addr: "127.0.0.1",
36+
Port: 5000,
3637
Instances: []proxy.InstanceConnConfig{
3738
{Name: "proj:region:inst1"},
3839
{Name: "proj:region:inst2"},
@@ -44,6 +45,7 @@ func TestClientInitialization(t *testing.T) {
4445
desc: "with instance address",
4546
in: &proxy.Config{
4647
Addr: "1.1.1.1", // bad address, binding shouldn't happen here.
48+
Port: 5000,
4749
Instances: []proxy.InstanceConnConfig{
4850
{Addr: "0.0.0.0", Name: "proj:region:inst1"},
4951
},
@@ -54,12 +56,41 @@ func TestClientInitialization(t *testing.T) {
5456
desc: "IPv6 support",
5557
in: &proxy.Config{
5658
Addr: "::1",
59+
Port: 5000,
5760
Instances: []proxy.InstanceConnConfig{
5861
{Name: "proj:region:inst1"},
5962
},
6063
},
6164
wantAddrs: []string{"[::1]:5000"},
6265
},
66+
{
67+
desc: "with instance port",
68+
in: &proxy.Config{
69+
Addr: "127.0.0.1",
70+
Port: 5000,
71+
Instances: []proxy.InstanceConnConfig{
72+
{Name: "proj:region:inst1", Port: 6000},
73+
},
74+
},
75+
wantAddrs: []string{"127.0.0.1:6000"},
76+
},
77+
{
78+
desc: "with global port and instance port",
79+
in: &proxy.Config{
80+
Addr: "127.0.0.1",
81+
Port: 5000,
82+
Instances: []proxy.InstanceConnConfig{
83+
{Name: "proj:region:inst1"},
84+
{Name: "proj:region:inst2", Port: 6000},
85+
{Name: "proj:region:inst3"},
86+
},
87+
},
88+
wantAddrs: []string{
89+
"127.0.0.1:5000",
90+
"127.0.0.1:6000",
91+
"127.0.0.1:5001",
92+
},
93+
},
6394
}
6495
for _, tc := range tcs {
6596
t.Run(tc.desc, func(t *testing.T) {

0 commit comments

Comments
 (0)