Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit eeb2f7b

Browse files
kylecarbskylecarbs
committed
Improve testing of config-ssh
1 parent 008ba69 commit eeb2f7b

File tree

5 files changed

+197
-94
lines changed

5 files changed

+197
-94
lines changed

‎agent/agent_test.go

Lines changed: 77 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,13 @@ package agent_test
22

33
import (
44
"context"
5+
"fmt"
6+
"io"
7+
"net"
8+
"os/exec"
9+
"path/filepath"
510
"runtime"
11+
"strconv"
612
"strings"
713
"testing"
814

@@ -29,7 +35,8 @@ func TestAgent(t *testing.T) {
2935
t.Parallel()
3036
t.Run("SessionExec", func(t *testing.T) {
3137
t.Parallel()
32-
session := setupSSH(t)
38+
session := setupSSHSession(t)
39+
3340
command := "echo test"
3441
if runtime.GOOS == "windows" {
3542
command = "cmd.exe /c echo test"
@@ -41,7 +48,7 @@ func TestAgent(t *testing.T) {
4148

4249
t.Run("GitSSH", func(t *testing.T) {
4350
t.Parallel()
44-
session := setupSSH(t)
51+
session := setupSSHSession(t)
4552
command := "sh -c 'echo $GIT_SSH_COMMAND'"
4653
if runtime.GOOS == "windows" {
4754
command = "cmd.exe /c echo %GIT_SSH_COMMAND%"
@@ -53,7 +60,7 @@ func TestAgent(t *testing.T) {
5360

5461
t.Run("SessionTTY", func(t *testing.T) {
5562
t.Parallel()
56-
session := setupSSH(t)
63+
session := setupSSHSession(t)
5764
prompt := "$"
5865
command := "bash"
5966
if runtime.GOOS == "windows" {
@@ -76,9 +83,73 @@ func TestAgent(t *testing.T) {
7683
err = session.Wait()
7784
require.NoError(t, err)
7885
})
86+
87+
t.Run("LocalForwarding", func(t *testing.T) {
88+
t.Parallel()
89+
random, err := net.Listen("tcp", "127.0.0.1:0")
90+
require.NoError(t, err)
91+
_ = random.Close()
92+
tcpAddr, valid := random.Addr().(*net.TCPAddr)
93+
require.True(t, valid)
94+
randomPort := tcpAddr.Port
95+
96+
local, err := net.Listen("tcp", "127.0.0.1:0")
97+
require.NoError(t, err)
98+
tcpAddr, valid = local.Addr().(*net.TCPAddr)
99+
require.True(t, valid)
100+
localPort := tcpAddr.Port
101+
done := make(chan struct{})
102+
go func() {
103+
conn, err := local.Accept()
104+
require.NoError(t, err)
105+
_ = conn.Close()
106+
close(done)
107+
}()
108+
109+
err = setupSSHCommand(t, []string{"-L", fmt.Sprintf("%d:127.0.0.1:%d", randomPort, localPort)}, []string{"echo", "test"}).Start()
110+
require.NoError(t, err)
111+
112+
conn, err := net.Dial("tcp", "127.0.0.1:"+strconv.Itoa(localPort))
113+
require.NoError(t, err)
114+
conn.Close()
115+
<-done
116+
})
79117
}
80118

81-
func setupSSH(t *testing.T) *ssh.Session {
119+
func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
120+
agentConn := setupAgent(t)
121+
socket := filepath.Join(t.TempDir(), "ssh")
122+
listener, err := net.Listen("unix", socket)
123+
require.NoError(t, err)
124+
go func() {
125+
for {
126+
conn, err := listener.Accept()
127+
if err != nil {
128+
return
129+
}
130+
ssh, err := agentConn.SSH()
131+
require.NoError(t, err)
132+
go io.Copy(conn, ssh)
133+
go io.Copy(ssh, conn)
134+
}
135+
}()
136+
t.Cleanup(func() {
137+
_ = listener.Close()
138+
})
139+
args := append(beforeArgs, "-o", "ProxyCommand socat - UNIX-CLIENT:"+socket, "-o", "StrictHostKeyChecking=no", "host")
140+
args = append(args, afterArgs...)
141+
return exec.Command("ssh", args...)
142+
}
143+
144+
func setupSSHSession(t *testing.T) *ssh.Session {
145+
sshClient, err := setupAgent(t).SSHClient()
146+
require.NoError(t, err)
147+
session, err := sshClient.NewSession()
148+
require.NoError(t, err)
149+
return session
150+
}
151+
152+
func setupAgent(t *testing.T) *agent.Conn {
82153
client, server := provisionersdk.TransportPipe()
83154
closer := agent.New(func(ctx context.Context, opts *peer.ConnOptions) (*peerbroker.Listener, error) {
84155
return peerbroker.Listen(server, nil, opts)
@@ -100,14 +171,9 @@ func setupSSH(t *testing.T) *ssh.Session {
100171
t.Cleanup(func() {
101172
_ = conn.Close()
102173
})
103-
agentClient := &agent.Conn{
174+
175+
return &agent.Conn{
104176
Negotiator: api,
105177
Conn: conn,
106178
}
107-
sshClient, err := agentClient.SSHClient()
108-
require.NoError(t, err)
109-
session, err := sshClient.NewSession()
110-
require.NoError(t, err)
111-
112-
return session
113179
}

‎cli/configssh.go

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@ const sshEndToken = "# ------------END-CODER------------"
3131

3232
func configSSH() *cobra.Command {
3333
var (
34-
binaryFile string
3534
sshConfigFile string
35+
sshOptions []string
3636
)
3737
cmd := &cobra.Command{
3838
Use: "config-ssh",
@@ -62,11 +62,9 @@ func configSSH() *cobra.Command {
6262
return xerrors.New("You don't have any workspaces!")
6363
}
6464

65-
if binaryFile == "" {
66-
binaryFile, err = currentBinPath(cmd)
67-
if err != nil {
68-
return err
69-
}
65+
binaryFile, err := currentBinPath(cmd)
66+
if err != nil {
67+
return err
7068
}
7169

7270
root := createConfig(cmd)
@@ -90,13 +88,19 @@ func configSSH() *cobra.Command {
9088
if len(resource.Agents) > 1 {
9189
hostname += "." + agent.Name
9290
}
93-
sshConfigContent += strings.Join([]string{
91+
configOptions := []string{
9492
"Host coder." + hostname,
9593
"\tHostName coder." + hostname,
94+
}
95+
for _, option := range sshOptions {
96+
configOptions = append(configOptions, "\t"+option)
97+
}
98+
configOptions = append(configOptions,
9699
fmt.Sprintf("\tProxyCommand %q --global-config %q ssh --stdio %s", binaryFile, root, hostname),
97100
"\tConnectTimeout=0",
98101
"\tStrictHostKeyChecking=no",
99-
}, "\n") + "\n"
102+
)
103+
sshConfigContent += strings.Join(configOptions, "\n") + "\n"
100104
sshConfigContentMutex.Unlock()
101105
}
102106
}
@@ -122,9 +126,8 @@ func configSSH() *cobra.Command {
122126
return nil
123127
},
124128
}
125-
cliflag.StringVarP(cmd.Flags(), &binaryFile, "binary-file", "", "CODER_BINARY_PATH", "", "Specifies the path to a Coder binary.")
126-
_ = cmd.Flags().MarkHidden("binary-file")
127129
cliflag.StringVarP(cmd.Flags(), &sshConfigFile, "ssh-config-file", "", "CODER_SSH_CONFIG_FILE", "~/.ssh/config", "Specifies the path to an SSH config.")
130+
cmd.Flags().StringArrayVarP(&sshOptions, "ssh-option", "o", []string{}, "Specifies additional options embedded in each host.")
128131

129132
return cmd
130133
}

‎cli/configssh_test.go

Lines changed: 91 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,20 @@
11
package cli_test
22

33
import (
4+
"context"
5+
"io"
6+
"net"
47
"os"
58
"os/exec"
69
"path/filepath"
710
"strings"
811
"testing"
9-
"time"
1012

11-
"cdr.dev/slog/sloggers/slogtest"
1213
"github.com/google/uuid"
1314
"github.com/stretchr/testify/require"
1415

16+
"cdr.dev/slog/sloggers/slogtest"
17+
1518
"github.com/coder/coder/agent"
1619
"github.com/coder/coder/cli/clitest"
1720
"github.com/coder/coder/coderd/coderdtest"
@@ -24,82 +27,98 @@ import (
2427

2528
func TestConfigSSH(t *testing.T) {
2629
t.Parallel()
27-
binPath := filepath.Join(t.TempDir(), "coder")
28-
_, err := exec.Command("go", "build", "-o", binPath, "github.com/coder/coder/cmd/coder").CombinedOutput()
29-
require.NoError(t, err)
30-
31-
t.Run("Dial", func(t *testing.T) {
32-
t.Parallel()
33-
client := coderdtest.New(t, nil)
34-
user := coderdtest.CreateFirstUser(t, client)
35-
coderdtest.NewProvisionerDaemon(t, client)
36-
authToken := uuid.NewString()
37-
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
38-
Parse: echo.ParseComplete,
39-
ProvisionDryRun: []*proto.Provision_Response{{
40-
Type: &proto.Provision_Response_Complete{
41-
Complete: &proto.Provision_Complete{
42-
Resources: []*proto.Resource{{
30+
client := coderdtest.New(t, nil)
31+
user := coderdtest.CreateFirstUser(t, client)
32+
coderdtest.NewProvisionerDaemon(t, client)
33+
authToken := uuid.NewString()
34+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
35+
Parse: echo.ParseComplete,
36+
ProvisionDryRun: []*proto.Provision_Response{{
37+
Type: &proto.Provision_Response_Complete{
38+
Complete: &proto.Provision_Complete{
39+
Resources: []*proto.Resource{{
40+
Name: "example",
41+
Type: "aws_instance",
42+
Agents: []*proto.Agent{{
43+
Id: uuid.NewString(),
4344
Name: "example",
44-
Type: "aws_instance",
45-
Agents: []*proto.Agent{{
46-
Id: uuid.NewString(),
47-
Name: "example",
48-
}},
4945
}},
50-
},
46+
}},
5147
},
52-
}},
53-
Provision: []*proto.Provision_Response{{
54-
Type: &proto.Provision_Response_Complete{
55-
Complete: &proto.Provision_Complete{
56-
Resources: []*proto.Resource{{
48+
},
49+
}},
50+
Provision: []*proto.Provision_Response{{
51+
Type: &proto.Provision_Response_Complete{
52+
Complete: &proto.Provision_Complete{
53+
Resources: []*proto.Resource{{
54+
Name: "example",
55+
Type: "aws_instance",
56+
Agents: []*proto.Agent{{
57+
Id: uuid.NewString(),
5758
Name: "example",
58-
Type: "aws_instance",
59-
Agents: []*proto.Agent{{
60-
Id: uuid.NewString(),
61-
Name: "example",
62-
Auth: &proto.Agent_Token{
63-
Token: authToken,
64-
},
65-
}},
59+
Auth: &proto.Agent_Token{
60+
Token: authToken,
61+
},
6662
}},
67-
},
63+
}},
6864
},
69-
}},
70-
})
71-
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
72-
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
73-
workspace := coderdtest.CreateWorkspace(t, client, codersdk.Me, template.ID)
74-
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
75-
agentClient := codersdk.New(client.URL)
76-
agentClient.SessionToken = authToken
77-
agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &peer.ConnOptions{
78-
Logger: slogtest.Make(t, nil),
79-
})
80-
t.Cleanup(func() {
81-
_ = agentCloser.Close()
82-
})
83-
tempFile, err := os.CreateTemp(t.TempDir(), "")
84-
require.NoError(t, err)
85-
_ = tempFile.Close()
86-
cmd, root := clitest.New(t, "config-ssh", "--binary-file", binPath, "--ssh-config-file", tempFile.Name())
87-
clitest.SetupConfig(t, client, root)
88-
doneChan := make(chan struct{})
89-
pty := ptytest.New(t)
90-
cmd.SetIn(pty.Input())
91-
cmd.SetOut(pty.Output())
92-
go func() {
93-
defer close(doneChan)
94-
err := cmd.Execute()
65+
},
66+
}},
67+
})
68+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
69+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
70+
workspace := coderdtest.CreateWorkspace(t, client, codersdk.Me, template.ID)
71+
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
72+
agentClient := codersdk.New(client.URL)
73+
agentClient.SessionToken = authToken
74+
agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &peer.ConnOptions{
75+
Logger: slogtest.Make(t, nil),
76+
})
77+
t.Cleanup(func() {
78+
_ = agentCloser.Close()
79+
})
80+
tempFile, err := os.CreateTemp(t.TempDir(), "")
81+
require.NoError(t, err)
82+
_ = tempFile.Close()
83+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
84+
agentConn, err := client.DialWorkspaceAgent(context.Background(), resources[0].Agents[0].ID, nil, nil)
85+
require.NoError(t, err)
86+
defer agentConn.Close()
87+
socket := filepath.Join(t.TempDir(), "ssh")
88+
listener, err := net.Listen("unix", socket)
89+
require.NoError(t, err)
90+
go func() {
91+
for {
92+
conn, err := listener.Accept()
93+
if err != nil {
94+
return
95+
}
96+
ssh, err := agentConn.SSH()
9597
require.NoError(t, err)
96-
}()
97-
<-doneChan
98-
t.Log(tempFile.Name())
99-
time.Sleep(time.Hour)
100-
output, err := exec.Command("ssh", "-F", tempFile.Name(), "coder."+workspace.Name, "echo", "test").Output()
101-
t.Log(string(output))
102-
require.NoError(t, err)
103-
require.Equal(t, "test", strings.TrimSpace(string(output)))
98+
go io.Copy(conn, ssh)
99+
go io.Copy(ssh, conn)
100+
}
101+
}()
102+
t.Cleanup(func() {
103+
_ = listener.Close()
104104
})
105+
106+
cmd, root := clitest.New(t, "config-ssh", "--ssh-option", "ProxyCommand socat - UNIX-CLIENT:"+socket, "--ssh-config-file", tempFile.Name())
107+
clitest.SetupConfig(t, client, root)
108+
doneChan := make(chan struct{})
109+
pty := ptytest.New(t)
110+
cmd.SetIn(pty.Input())
111+
cmd.SetOut(pty.Output())
112+
go func() {
113+
defer close(doneChan)
114+
err := cmd.Execute()
115+
require.NoError(t, err)
116+
}()
117+
<-doneChan
118+
119+
t.Log(tempFile.Name())
120+
// #nosec
121+
data, err := exec.Command("ssh", "-F", tempFile.Name(), "coder."+workspace.Name, "echo", "test").Output()
122+
require.NoError(t, err)
123+
require.Equal(t, "test", strings.TrimSpace(string(data)))
105124
}

‎cli/templatelist.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,11 @@ package cli
33
import (
44
"fmt"
55

6-
"github.com/coder/coder/cli/cliui"
76
"github.com/fatih/color"
87
"github.com/jedib0t/go-pretty/v6/table"
98
"github.com/spf13/cobra"
9+
10+
"github.com/coder/coder/cli/cliui"
1011
)
1112

1213
func templateList() *cobra.Command {

0 commit comments

Comments
 (0)