feat: enhance plugin networking capabilities

- Enable SSH server to listen on all interfaces by uncommenting relevant lines in `sshd_config` - Add new `protocol` and `proxy.protocol` flags to `main.go` with usage information and default values - Change the `Port` field type from `string` to `int` in `plugin.go` and `plugin_test.go` - Refactor variable name from `host` to `h` and add `port` variable in `plugin.go` loop - Remove commented-out code and refactor `hostPort` function in `plugin.go` - Add import for `io` package in `plugin_test.go` - Add new test function `TestPlugin_hostPort` with multiple test cases in `plugin_test.go` Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>
appleboy · Jun 1, 2024 · 05b4e1c · 05b4e1c
1 parent 45ef128
commit 05b4e1c
Show file tree

Hide file tree

Showing 4 changed files with 113 additions and 27 deletions.
diff --git a/Makefile b/Makefile
@@ -108,6 +108,8 @@ ssh-server:
 	rm -rf /etc/ssh/ssh_host_rsa_key /etc/ssh/ssh_host_dsa_key
 	sed -i 's/^#PubkeyAuthentication yes/PubkeyAuthentication yes/g' /etc/ssh/sshd_config
 	sed -i 's/AllowTcpForwarding no/AllowTcpForwarding yes/g' /etc/ssh/sshd_config
+	sed -i 's/^#ListenAddress 0.0.0.0/ListenAddress 0.0.0.0/g' /etc/ssh/sshd_config
+	sed -i 's/^#ListenAddress ::/ListenAddress ::/g' /etc/ssh/sshd_config
 	./tests/entrypoint.sh /usr/sbin/sshd -D &
 
 coverage:

diff --git a/main.go b/main.go
@@ -50,6 +50,12 @@ func main() {
 			EnvVars: []string{"PLUGIN_PORT", "SSH_PORT", "INPUT_PORT"},
 			Value:   22,
 		},
+		&cli.StringFlag{
+			Name:    "protocol",
+			Usage:   "The IP protocol to use. Valid values are \"tcp\". \"tcp4\" or \"tcp6\". Default to tcp.",
+			EnvVars: []string{"PLUGIN_PROTOCOL", "SSH_PROTOCOL", "INPUT_PROTOCOL"},
+			Value:   "tcp",
+		},
 		&cli.StringFlag{
 			Name:    "username",
 			Aliases: []string{"user", "u"},
@@ -135,6 +141,12 @@ func main() {
 			EnvVars: []string{"PLUGIN_PROXY_PORT", "PROXY_SSH_PORT", "INPUT_PROXY_PORT"},
 			Value:   "22",
 		},
+		&cli.StringFlag{
+			Name:    "proxy.protocol",
+			Usage:   "The IP protocol to use for the proxy. Valid values are \"tcp\". \"tcp4\" or \"tcp6\". Default to tcp.",
+			EnvVars: []string{"PLUGIN_PROXY_PROTOCOL", "SSH_PROXY_PROTOCOL", "INPUT_PROXY_PROTOCOL"},
+			Value:   "tcp",
+		},
 		&cli.StringFlag{
 			Name:    "proxy.username",
 			Usage:   "connect as user of proxy",
@@ -261,7 +273,7 @@ func run(c *cli.Context) error {
 	plugin := Plugin{
 		Config: Config{
 			Host:              c.StringSlice("host"),
-			Port:              c.String("port"),
+			Port:              c.Int("port"),
 			Username:          c.String("username"),
 			Password:          c.String("password"),
 			Passphrase:        c.String("ssh-passphrase"),

diff --git a/plugin.go b/plugin.go
@@ -26,7 +26,8 @@ type (
 	// Config for the plugin.
 	Config struct {
 		Host              []string
-		Port              string
+		Port              int
+		Protocol          easyssh.Protocol
 		Username          string
 		Password          string
 		Key               string
@@ -115,12 +116,13 @@ func (p *Plugin) removeDestFile(os string, ssh *easyssh.MakeConfig) error {
 }
 
 func (p *Plugin) removeAllDestFile() error {
-	for _, host := range trimValues(p.Config.Host) {
+	for _, h := range trimValues(p.Config.Host) {
+		host, port := p.hostPort(h)
 		ssh := &easyssh.MakeConfig{
 			Server:            host,
 			User:              p.Config.Username,
 			Password:          p.Config.Password,
-			Port:              p.Config.Port,
+			Port:              port,
 			Key:               p.Config.Key,
 			KeyPath:           p.Config.KeyPath,
 			Passphrase:        p.Config.Passphrase,
@@ -389,19 +391,16 @@ func (p *Plugin) Exec() error {
 	return nil
 }
 
-// This function takes a Plugin struct and a host string and returns the host and port as separate strings.
 func (p Plugin) hostPort(host string) (string, string) {
-	// Split the host string by colon (":") to get the host and port
 	hosts := strings.Split(host, ":")
-	// Get the default port from the Plugin's Config field
-	port := p.Config.Port
-	// If the host string contains a port (i.e. it has more than one element after splitting), set the port to that value
-	if len(hosts) > 1 {
+	port := strconv.Itoa(p.Config.Port)
+	if len(hosts) > 1 &&
+		(p.Config.Protocol == easyssh.PROTOCOL_TCP ||
+			p.Config.Protocol == easyssh.PROTOCOL_TCP4) {
 		host = hosts[0]
 		port = hosts[1]
 	}
 
-	// Return the host and port as separate strings
 	return host, port
 }
 

diff --git a/plugin_test.go b/plugin_test.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"io"
 	"log"
 	"os"
 	"os/exec"
@@ -41,7 +42,7 @@ func TestMissingSourceConfig(t *testing.T) {
 		Config: Config{
 			Host:     []string{"example.com"},
 			Username: "ubuntu",
-			Port:     "443",
+			Port:     443,
 			Password: "1234",
 		},
 	}
@@ -81,7 +82,7 @@ func TestSCPFileFromPublicKey(t *testing.T) {
 		Config: Config{
 			Host:           []string{"localhost"},
 			Username:       "drone-scp",
-			Port:           "22",
+			Port:           22,
 			KeyPath:        "tests/.ssh/id_rsa",
 			Source:         []string{"tests/a.txt", "tests/b.txt"},
 			Target:         []string{filepath.Join(u.HomeDir, "/test")},
@@ -131,7 +132,7 @@ func TestSCPFileFromPublicKeyWithPassphrase(t *testing.T) {
 		Config: Config{
 			Host:           []string{"localhost"},
 			Username:       "drone-scp",
-			Port:           "22",
+			Port:           22,
 			KeyPath:        "tests/.ssh/test",
 			Passphrase:     "1234",
 			Source:         []string{"tests/a.txt", "tests/b.txt"},
@@ -164,7 +165,7 @@ func TestWrongFingerprint(t *testing.T) {
 		Config: Config{
 			Host:           []string{"localhost"},
 			Username:       "drone-scp",
-			Port:           "22",
+			Port:           22,
 			KeyPath:        "./tests/.ssh/id_rsa",
 			Source:         []string{"tests/a.txt", "tests/b.txt"},
 			Target:         []string{filepath.Join(u.HomeDir, "/test2")},
@@ -188,7 +189,6 @@ func getHostPublicKeyFile(keypath string) (ssh.PublicKey, error) {
 	}
 
 	pubkey, _, _, _, err = ssh.ParseAuthorizedKey(buf)
-
 	if err != nil {
 		return nil, err
 	}
@@ -215,7 +215,7 @@ func TestSCPFileFromPublicKeyWithFingerprint(t *testing.T) {
 		Config: Config{
 			Host:           []string{"localhost"},
 			Username:       "drone-scp",
-			Port:           "22",
+			Port:           22,
 			KeyPath:        "./tests/.ssh/id_rsa",
 			Fingerprint:    ssh.FingerprintSHA256(hostKey),
 			Source:         []string{"tests/a.txt", "tests/b.txt"},
@@ -254,7 +254,7 @@ func TestSCPWildcardFileList(t *testing.T) {
 		Config: Config{
 			Host:           []string{"localhost"},
 			Username:       "drone-scp",
-			Port:           "22",
+			Port:           22,
 			KeyPath:        "tests/.ssh/id_rsa",
 			Source:         []string{"tests/global/*"},
 			Target:         []string{filepath.Join(u.HomeDir, "abc")},
@@ -286,7 +286,7 @@ func TestSCPFromProxySetting(t *testing.T) {
 		Config: Config{
 			Host:           []string{"localhost"},
 			Username:       "drone-scp",
-			Port:           "22",
+			Port:           22,
 			KeyPath:        "tests/.ssh/id_rsa",
 			Source:         []string{"tests/global/*"},
 			Target:         []string{filepath.Join(u.HomeDir, "def")},
@@ -330,7 +330,7 @@ func TestStripComponentsFlag(t *testing.T) {
 		Config: Config{
 			Host:            []string{"localhost"},
 			Username:        "drone-scp",
-			Port:            "22",
+			Port:            22,
 			KeyPath:         "tests/.ssh/id_rsa",
 			Source:          []string{"tests/global/*"},
 			StripComponents: 2,
@@ -363,7 +363,7 @@ func TestUseInsecureCipherFlag(t *testing.T) {
 		Config: Config{
 			Host:              []string{"localhost"},
 			Username:          "drone-scp",
-			Port:              "22",
+			Port:              22,
 			KeyPath:           "tests/.ssh/id_rsa",
 			Source:            []string{"tests/global/*"},
 			StripComponents:   2,
@@ -403,7 +403,7 @@ func TestIgnoreList(t *testing.T) {
 		Config: Config{
 			Host:            []string{"localhost"},
 			Username:        "drone-scp",
-			Port:            "22",
+			Port:            22,
 			KeyPath:         "tests/.ssh/id_rsa",
 			Source:          []string{"tests/global/*", "!tests/global/c.txt", "!tests/global/e.txt"},
 			StripComponents: 2,
@@ -483,7 +483,7 @@ func TestIncorrectPassword(t *testing.T) {
 		Config: Config{
 			Host:           []string{"localhost"},
 			Username:       "drone-scp",
-			Port:           "22",
+			Port:           22,
 			Password:       "123456",
 			Source:         []string{"tests/a.txt", "tests/b.txt"},
 			Target:         []string{"/home"},
@@ -506,7 +506,7 @@ func TestNoPermissionCreateFolder(t *testing.T) {
 		Config: Config{
 			Host:           []string{"localhost"},
 			Username:       "drone-scp",
-			Port:           "22",
+			Port:           22,
 			KeyPath:        "tests/.ssh/id_rsa",
 			Source:         []string{"tests/a.txt", "tests/b.txt"},
 			Target:         []string{"/etc/test"},
@@ -782,7 +782,7 @@ func TestTargetFolderWithSpaces(t *testing.T) {
 		Config: Config{
 			Host:            []string{"localhost"},
 			Username:        "drone-scp",
-			Port:            "22",
+			Port:            22,
 			KeyPath:         "tests/.ssh/id_rsa",
 			Source:          []string{"tests/global/*"},
 			StripComponents: 2,
@@ -821,7 +821,7 @@ func TestHostPortString(t *testing.T) {
 		Config: Config{
 			Host:            []string{"localhost:22", "localhost:22"},
 			Username:        "drone-scp",
-			Port:            "8080",
+			Port:            8080,
 			KeyPath:         "tests/.ssh/id_rsa",
 			Source:          []string{"tests/global/*"},
 			StripComponents: 2,
@@ -848,7 +848,7 @@ func TestHostPortString(t *testing.T) {
 func TestHostPort(t *testing.T) {
 	p := Plugin{
 		Config: Config{
-			Port: "8080",
+			Port: 8080,
 		},
 	}
 
@@ -870,3 +870,76 @@ func TestHostPort(t *testing.T) {
 		t.Errorf("hostPort(%s) = (%s, %s); expected (%s, %s)", host2, actualHost2, actualPort2, expectedHost2, expectedPort2)
 	}
 }
+
+func TestPlugin_hostPort(t *testing.T) {
+	type fields struct {
+		Config Config
+		Writer io.Writer
+	}
+	type args struct {
+		h string
+	}
+	tests := []struct {
+		name     string
+		fields   fields
+		args     args
+		wantHost string
+		wantPort string
+	}{
+		{
+			name: "default host and port",
+			fields: fields{
+				Config: Config{
+					Port: 22,
+				},
+			},
+			args: args{
+				h: "localhost",
+			},
+			wantHost: "localhost",
+			wantPort: "22",
+		},
+		{
+			name: "different port",
+			fields: fields{
+				Config: Config{
+					Port:     22,
+					Protocol: easyssh.PROTOCOL_TCP4,
+				},
+			},
+			args: args{
+				h: "localhost:443",
+			},
+			wantHost: "localhost",
+			wantPort: "443",
+		},
+		{
+			name: "ipv6",
+			fields: fields{
+				Config: Config{
+					Port:     22,
+					Protocol: easyssh.PROTOCOL_TCP6,
+				},
+			},
+			args: args{
+				h: "::1",
+			},
+			wantHost: "::1",
+			wantPort: "22",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			p := Plugin{
+				Config: tt.fields.Config,
+			}
+			gotHost, gotPort := p.hostPort(tt.args.h)
+			if gotHost != tt.wantHost {
+				t.Errorf("Plugin.hostPort() gotHost = %v, want %v", gotHost, tt.wantHost)
+			}
+			if gotPort != tt.wantPort {
+				t.Errorf("Plugin.hostPort() gotPort = %v, want %v", gotPort, tt.wantPort)
+			}
+		})
+	}
+}