Skip to content

ssh plugin: allow custom SSH port configuration #5047

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
volcano-sh-bot merged 1 commit into from
Feb 24, 2026
+85 −8
Merged

ssh plugin: allow custom SSH port configuration #5047

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion docs/user-guide/how_to_use_ssh_plugin.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,14 +21,16 @@ and subdomain.
|-----|----------------------|--------|---------------------|----------|-----------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| 1 | `ssh-key-file-path` | String | `/root/.ssh` | N | The path used to store ssh private and public keys. | ssh: ["--ssh-key-file-path=/home/user/.ssh"] |
| 2 | `ssh-private-key` | String | DEFAULT_PRIVATE_KEY | N | The input string of the private key. | ssh: ["--ssh-private-key=-----BEGIN RSA PRIVATE KEY-----\nMIIEpAIBAAKCAQEAyeyZjWDx5Na9bw1f61M4s+QlLT/kyrB37AR2j5Sb/A9hvJak\nLNQQpNC+KVfYNl4jePG+6lwHqye//pcC9+0SWsHWwgaahjMLnAthR2k8JAakNA9x\nV/wHz0YU99OKEetaOuxXpWZPXCHX0zuQO87YbdKzRbgxACirM3Phkwr7XLtQtWZk\nyXG34CQXZQWgBIS1Fl+PlGOpVpOPnWoZPMpbAK74i/Tz4sP8Zhqc6dya1hrbUwY3\nYfMZNYXpaAw7wWVjq8grfs0+Fl3SxHrzTXge2m+eZAZ6iPJ8cX4uYKxi0ZmxpM/a\ngI6Mmjq0MU75Vxpq22LaUvHIpOfX5UxhkrsxlwIDAQABAoIBAQDGOuIb6zpNn4rl\nBMpPqamW4LimjX08hrWUHGWQWyIu96LJk1GlOKMGSm8FA1odNZm5WApG5QYaPrG7\na+DcJ/7G3ljIrdbxPBd/n6RmiKcj7ukwuqBY8fFwyKo5CZEYOmagRfldRO1P02Gf\n22+jZ1MNrbWVElf4gfRgVLj0s+lEhFkzhi+QGMmMpjEJnnG98xxVGEvWMw1rnKJm\n3Gi771Gltbg3GuEPs3IeoBgba3EaHmSxJnBivAL4zsO8UUCAXB13cUiXx8qO7y1e\nCSWSenRmK2ugbL6v0co12O0n0pxF9xlJ6fALdRWzpJsFlN3ttkY9N5GrQc/pVjOa\nvqa172RRAoGBAOSAIMNLT6QjgYDk5Z7ZxjNnxH/lMso+cx6bxk9YMKRrw0fDQh8m\ncBAihXhuntCPDGhrzQ+Anqx4jJVDFqac0xBck90a8LmmzD0q72eDTCYPouDWe6DL\nJQAc/HDmIC13sADEXmGW3c0Qn4hjBnMd89ouYj7ZajU2sED2irPPc/HLAoGBAOI5\nruL4Q0FarGrP3a9z9EDrVJsK2OfSTaJ7rhZ+uvB838svbHU+4mEYPhx4PCwvrYyi\nFn4hyau003ZmLc1qTABjmwcO/PPiYyoRHJDUIIhiIyIL+id/G53uG2eTzqYtU6uS\nnAIB2rKwwhU8ek+zbJBLu5uxuxlf4mdZITdkwtXlAoGBALH3RQ02A9JgQQYFwP2G\nucLhx/6goX05RGoLg1na4w+8Sr0Cy+X9BvzaFkAlUBY5w700cOLpFyxXO48pUGP1\n8sFkiVmFGQZPbfUaEpn5ff6K4R3ijyk97xR2fvrjkR44gOEoECZL3XZQwx/zmFti\nccF1rNksdnb5oC8IliDTq4cfAoGANyy6asECJj5nLuXju5ccS3kZ+XZ70I6KQMbJ\nftMJ5P2P146JdU8RB31SKL9qbZxzR4mA0uKKvUYtDQN+yErUnoOsm9wb9Z+RcAEc\nZnZWOO02hGdHa7qkkbAxHuH91KnZbk8jnZm2LT7PFz7Y1fd80vSlnSOL7nRkU7B5\nWXlJy8ECgYA4g0wc0Jq8c1Q0FulMkOQqYRDXaDo34987L+mZ70i/RtdkKjK/IKJ9\n18UDCyEaDPD0BWBJGPejZkY8UD6FBG/5k7wNIbT7hHLRSRlw4iRmVX2hRVXrXzD8\nvc86Qyg2iG0JqkMAvRdH40amPKp5bW4VcfcvQo4TSsI972u12rgwtg==\n-----END RSA PRIVATE KEY-----\n"] |
| 3 | `ssh-public-key` | String | DEFAULT_PUBLIC_KEY | N | The input string of the public key. | ssh: ["--ssh-public-key=ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDJ7JmNYPHk1r1vDV/rUziz5CUtP+TKsHfsBHaPlJv8D2G8lqQs1BCk0L4pV9g2XiN48b7qXAerJ7/+lwL37RJawdbCBpqGMwucC2FHaTwkBqQ0D3FX/AfPRhT304oR61o67FelZk9cIdfTO5A7ztht0rNFuDEAKKszc+GTCvtcu1C1ZmTJcbfgJBdlBaAEhLUWX4+UY6lWk4+dahk8ylsArviL9PPiw/xmGpzp3JrWGttTBjdh8xk1heloDDvBZWOryCt+zT4WXdLEevNNeB7ab55kBnqI8nxxfi5grGLRmbGkz9qAjoyaOrQxTvlXGmrbYtpS8cik59flTGGSuzGX root@aiplatform"] |
| 3 | `ssh-public-key` | String | DEFAULT_PUBLIC_KEY | N | The input string of the public key. | ssh: ["--ssh-public-key=ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDJ7JmNYPHk1r1vDV/rUziz5CUtP+TKsHfsBHaPlJv8D2G8lqQs1BCk0L4pV9g2XiN48b7qXAerJ7/+lwL37RJawdbCBpqGMwucC2FHaTwkBqQ0D3FX/AfPRhT304oR61o67FelZk9cIdfTO5A7ztht0rNFuDEAKKszc+GTCvtcu1C1ZmTJcbfgJBdlBaAEhLUWX4+UY6lWk4+dahk8ylsArviL9PPiw/xmGpzp3JrWGttTBjdh8xk1heloDDvBZWOryCt+zT4WXdLEevNNeB7ab55kBnqI8nxxfi5grGLRmbGkz9qAjoyaOrQxTvlXGmrbYtpS8cik59flTGGSuzGX root@aiplatform"] |
| 4 | `ssh-port` | Int | 22 | N | SSH port used in generated ssh config. Useful when port 22 is already in use. | ssh: ["--ssh-port=2026"] |

Note:
* `DEFAULT_PRIVATE_KEY` and `DEFAULT_PUBLIC_KEY` are not fully listed for they are too long. Please refer to the examples
behind for cases.
* Volcano is not responsible for the validation of `ssh-key-file-path`. So please guarantee it correct yourself.
* Suggest keeping blank and making use of the default values in most scenarios just like the example behind. If that,
Volcano will help generate a pair of keys and finish all the configuration by default.
* A custom SSH port can be configured using `ssh-port`.

## Examples
```yaml
Expand Down Expand Up @@ -95,10 +97,13 @@ StrictHostKeyChecking no
UserKnownHostsFile /dev/null
Host mpi-job-mpimaster-0
HostName mpi-job-mpimaster-0.mpi-job
Port 22
Host mpi-job-mpiworker-0
HostName mpi-job-mpiworker-0.mpi-job
Port 22
Host mpi-job-mpiworker-1
HostName mpi-job-mpiworker-1.mpi-job
Port 22
```
* You can sign in other hosts in `master` pod as follows.
```
Expand Down
29 changes: 22 additions & 7 deletions pkg/controllers/job/plugins/ssh/ssh.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,16 @@ type sshPlugin struct {

// public key string
sshPublicKey string

// custom port
sshPort int
}

func (sp *sshPlugin) validateSSHPort() {
if sp.sshPort <= 0 || sp.sshPort > 65535 {
klog.Warningf("invalid ssh port %d, falling back to default port 22 ...", sp.sshPort)
sp.sshPort = 22
}
}

// New creates ssh plugin
Expand All @@ -56,6 +66,7 @@ func New(client pluginsinterface.PluginClientset, arguments []string) pluginsint
pluginArguments: arguments,
client: client,
sshKeyFilePath: SSHAbsolutePath,
sshPort: 22,
}

p.addFlags()
Expand All @@ -74,16 +85,18 @@ func (sp *sshPlugin) OnPodCreate(pod *v1.Pod, job *batch.Job) error {
}

func (sp *sshPlugin) OnJobAdd(job *batch.Job) error {
sp.validateSSHPort()

if job.Status.ControlledResources["plugin-"+sp.Name()] == sp.Name() {
return nil
}

var data map[string][]byte
var err error
if len(sp.sshPrivateKey) > 0 {
data, err = withUserProvidedRsaKey(job, sp.sshPrivateKey, sp.sshPublicKey)
data, err = withUserProvidedRsaKey(job, sp.sshPrivateKey, sp.sshPublicKey, sp.sshPort)
} else {
data, err = generateRsaKey(job)
data, err = generateRsaKey(job, sp.sshPort)
}
if err != nil {
return err
Expand Down Expand Up @@ -186,7 +199,7 @@ func (sp *sshPlugin) mountRsaKey(pod *v1.Pod, job *batch.Job) {
}
}

func generateRsaKey(job *batch.Job) (map[string][]byte, error) {
func generateRsaKey(job *batch.Job, port int) (map[string][]byte, error) {
bitSize := 2048

privateKey, err := rsa.GenerateKey(rand.Reader, bitSize)
Expand Down Expand Up @@ -214,17 +227,17 @@ func generateRsaKey(job *batch.Job) (map[string][]byte, error) {
data[SSHPrivateKey] = privateKeyBytes
data[SSHPublicKey] = publicKeyBytes
data[SSHAuthorizedKeys] = publicKeyBytes
data[SSHConfig] = []byte(generateSSHConfig(job))
data[SSHConfig] = []byte(generateSSHConfig(job, port))

return data, nil
}

func withUserProvidedRsaKey(job *batch.Job, sshPrivateKey string, sshPublicKey string) (map[string][]byte, error) {
func withUserProvidedRsaKey(job *batch.Job, sshPrivateKey string, sshPublicKey string, port int) (map[string][]byte, error) {
data := make(map[string][]byte)
data[SSHPrivateKey] = []byte(sshPrivateKey)
data[SSHPublicKey] = []byte(sshPublicKey)
data[SSHAuthorizedKeys] = []byte(sshPublicKey)
data[SSHConfig] = []byte(generateSSHConfig(job))
data[SSHConfig] = []byte(generateSSHConfig(job, port))

return data, nil
}
Expand All @@ -239,13 +252,14 @@ func (sp *sshPlugin) addFlags() {
"ssh private and public keys, it is `/root/.ssh` by default.")
flagSet.StringVar(&sp.sshPrivateKey, "ssh-private-key", sp.sshPrivateKey, "The input string of the private key")
flagSet.StringVar(&sp.sshPublicKey, "ssh-public-key", sp.sshPublicKey, "The input string of the public key")
flagSet.IntVar(&sp.sshPort, "ssh-port", sp.sshPort, "SSH port used in ssh config. Default Port: 22")

if err := flagSet.Parse(sp.pluginArguments); err != nil {
klog.Errorf("plugin %s flagset parse failed, err: %v", sp.Name(), err)
}
}

func generateSSHConfig(job *batch.Job) string {
func generateSSHConfig(job *batch.Job, port int) string {
config := "StrictHostKeyChecking no\nUserKnownHostsFile /dev/null\n"

for _, ts := range job.Spec.Tasks {
Expand All @@ -261,6 +275,7 @@ func generateSSHConfig(job *batch.Job) string {

config += "Host " + hostName + "\n"
config += " HostName " + hostName + "." + subdomain + "\n"
config += fmt.Sprintf(" Port %d\n", port)
if len(ts.Template.Spec.Hostname) != 0 {
break
}
Expand Down
57 changes: 57 additions & 0 deletions pkg/controllers/job/plugins/ssh/ssh_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,13 @@ limitations under the License.
package ssh

import (
"fmt"
"strings"
"testing"

v1 "k8s.io/api/core/v1"
Copy link
Contributor

@hajnalmt hajnalmt Feb 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will this pass gofmt? It looks wierd that it has an enter beneath.

Copy link
Contributor

@hajnalmt hajnalmt Feb 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Passed fmt, this is fine.


batch "volcano.sh/apis/pkg/apis/batch/v1alpha1"
pluginsinterface "volcano.sh/volcano/pkg/controllers/job/plugins/interface"
)

Expand All @@ -31,26 +36,38 @@ func TestSSHPlugin(t *testing.T) {
sshKeyFilePath string
sshPrivateKey string
sshPublicKey string
sshPort int
}{
{
name: "no params specified",
sshKeyFilePath: SSHAbsolutePath,
sshPrivateKey: "",
sshPublicKey: "",
sshPort: 22,
},
{
name: "--ssh-key-file-path=/a/b",
params: []string{"--ssh-key-file-path=/a/b"},
sshKeyFilePath: "/a/b",
sshPrivateKey: "",
sshPublicKey: "",
sshPort: 22,
},
{
name: "with user provided ssh keys",
params: []string{"--ssh-private-key=-----BEGIN RSA PRIVATE KEY-----\\nMIIEpAIBAAKCAQEAyeyZjWDx5Na9bw1f61M4s+QlLT/kyrB37AR2j5Sb/A9hvJak\\nLNQQpNC+KVfYNl4jePG+6lwHqye//pcC9+0SWsHWwgaahjMLnAthR2k8JAakNA9x\\nV/wHz0YU99OKEetaOuxXpWZPXCHX0zuQO87YbdKzRbgxACirM3Phkwr7XLtQtWZk\\nyXG34CQXZQWgBIS1Fl+PlGOpVpOPnWoZPMpbAK74i/Tz4sP8Zhqc6dya1hrbUwY3\\nYfMZNYXpaAw7wWVjq8grfs0+Fl3SxHrzTXge2m+eZAZ6iPJ8cX4uYKxi0ZmxpM/a\\ngI6Mmjq0MU75Vxpq22LaUvHIpOfX5UxhkrsxlwIDAQABAoIBAQDGOuIb6zpNn4rl\\nBMpPqamW4LimjX08hrWUHGWQWyIu96LJk1GlOKMGSm8FA1odNZm5WApG5QYaPrG7\\na+DcJ/7G3ljIrdbxPBd/n6RmiKcj7ukwuqBY8fFwyKo5CZEYOmagRfldRO1P02Gf\\n22+jZ1MNrbWVElf4gfRgVLj0s+lEhFkzhi+QGMmMpjEJnnG98xxVGEvWMw1rnKJm\\n3Gi771Gltbg3GuEPs3IeoBgba3EaHmSxJnBivAL4zsO8UUCAXB13cUiXx8qO7y1e\\nCSWSenRmK2ugbL6v0co12O0n0pxF9xlJ6fALdRWzpJsFlN3ttkY9N5GrQc/pVjOa\\nvqa172RRAoGBAOSAIMNLT6QjgYDk5Z7ZxjNnxH/lMso+cx6bxk9YMKRrw0fDQh8m\\ncBAihXhuntCPDGhrzQ+Anqx4jJVDFqac0xBck90a8LmmzD0q72eDTCYPouDWe6DL\\nJQAc/HDmIC13sADEXmGW3c0Qn4hjBnMd89ouYj7ZajU2sED2irPPc/HLAoGBAOI5\\nruL4Q0FarGrP3a9z9EDrVJsK2OfSTaJ7rhZ+uvB838svbHU+4mEYPhx4PCwvrYyi\\nFn4hyau003ZmLc1qTABjmwcO/PPiYyoRHJDUIIhiIyIL+id/G53uG2eTzqYtU6uS\\nnAIB2rKwwhU8ek+zbJBLu5uxuxlf4mdZITdkwtXlAoGBALH3RQ02A9JgQQYFwP2G\\nucLhx/6goX05RGoLg1na4w+8Sr0Cy+X9BvzaFkAlUBY5w700cOLpFyxXO48pUGP1\\n8sFkiVmFGQZPbfUaEpn5ff6K4R3ijyk97xR2fvrjkR44gOEoECZL3XZQwx/zmFti\\nccF1rNksdnb5oC8IliDTq4cfAoGANyy6asECJj5nLuXju5ccS3kZ+XZ70I6KQMbJ\\nftMJ5P2P146JdU8RB31SKL9qbZxzR4mA0uKKvUYtDQN+yErUnoOsm9wb9Z+RcAEc\\nZnZWOO02hGdHa7qkkbAxHuH91KnZbk8jnZm2LT7PFz7Y1fd80vSlnSOL7nRkU7B5\\nWXlJy8ECgYA4g0wc0Jq8c1Q0FulMkOQqYRDXaDo34987L+mZ70i/RtdkKjK/IKJ9\\n18UDCyEaDPD0BWBJGPejZkY8UD6FBG/5k7wNIbT7hHLRSRlw4iRmVX2hRVXrXzD8\\nvc86Qyg2iG0JqkMAvRdH40amPKp5bW4VcfcvQo4TSsI972u12rgwtg==\\n-----END RSA PRIVATE KEY-----\\n", "--ssh-public-key=ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDJ7JmNYPHk1r1vDV/rUziz5CUtP+TKsHfsBHaPlJv8D2G8lqQs1BCk0L4pV9g2XiN48b7qXAerJ7/+lwL37RJawdbCBpqGMwucC2FHaTwkBqQ0D3FX/AfPRhT304oR61o67FelZk9cIdfTO5A7ztht0rNFuDEAKKszc+GTCvtcu1C1ZmTJcbfgJBdlBaAEhLUWX4+UY6lWk4+dahk8ylsArviL9PPiw/xmGpzp3JrWGttTBjdh8xk1heloDDvBZWOryCt+zT4WXdLEevNNeB7ab55kBnqI8nxxfi5grGLRmbGkz9qAjoyaOrQxTvlXGmrbYtpS8cik59flTGGSuzGX root@aiplatform"},
sshKeyFilePath: SSHAbsolutePath,
sshPrivateKey: "-----BEGIN RSA PRIVATE KEY-----\\nMIIEpAIBAAKCAQEAyeyZjWDx5Na9bw1f61M4s+QlLT/kyrB37AR2j5Sb/A9hvJak\\nLNQQpNC+KVfYNl4jePG+6lwHqye//pcC9+0SWsHWwgaahjMLnAthR2k8JAakNA9x\\nV/wHz0YU99OKEetaOuxXpWZPXCHX0zuQO87YbdKzRbgxACirM3Phkwr7XLtQtWZk\\nyXG34CQXZQWgBIS1Fl+PlGOpVpOPnWoZPMpbAK74i/Tz4sP8Zhqc6dya1hrbUwY3\\nYfMZNYXpaAw7wWVjq8grfs0+Fl3SxHrzTXge2m+eZAZ6iPJ8cX4uYKxi0ZmxpM/a\\ngI6Mmjq0MU75Vxpq22LaUvHIpOfX5UxhkrsxlwIDAQABAoIBAQDGOuIb6zpNn4rl\\nBMpPqamW4LimjX08hrWUHGWQWyIu96LJk1GlOKMGSm8FA1odNZm5WApG5QYaPrG7\\na+DcJ/7G3ljIrdbxPBd/n6RmiKcj7ukwuqBY8fFwyKo5CZEYOmagRfldRO1P02Gf\\n22+jZ1MNrbWVElf4gfRgVLj0s+lEhFkzhi+QGMmMpjEJnnG98xxVGEvWMw1rnKJm\\n3Gi771Gltbg3GuEPs3IeoBgba3EaHmSxJnBivAL4zsO8UUCAXB13cUiXx8qO7y1e\\nCSWSenRmK2ugbL6v0co12O0n0pxF9xlJ6fALdRWzpJsFlN3ttkY9N5GrQc/pVjOa\\nvqa172RRAoGBAOSAIMNLT6QjgYDk5Z7ZxjNnxH/lMso+cx6bxk9YMKRrw0fDQh8m\\ncBAihXhuntCPDGhrzQ+Anqx4jJVDFqac0xBck90a8LmmzD0q72eDTCYPouDWe6DL\\nJQAc/HDmIC13sADEXmGW3c0Qn4hjBnMd89ouYj7ZajU2sED2irPPc/HLAoGBAOI5\\nruL4Q0FarGrP3a9z9EDrVJsK2OfSTaJ7rhZ+uvB838svbHU+4mEYPhx4PCwvrYyi\\nFn4hyau003ZmLc1qTABjmwcO/PPiYyoRHJDUIIhiIyIL+id/G53uG2eTzqYtU6uS\\nnAIB2rKwwhU8ek+zbJBLu5uxuxlf4mdZITdkwtXlAoGBALH3RQ02A9JgQQYFwP2G\\nucLhx/6goX05RGoLg1na4w+8Sr0Cy+X9BvzaFkAlUBY5w700cOLpFyxXO48pUGP1\\n8sFkiVmFGQZPbfUaEpn5ff6K4R3ijyk97xR2fvrjkR44gOEoECZL3XZQwx/zmFti\\nccF1rNksdnb5oC8IliDTq4cfAoGANyy6asECJj5nLuXju5ccS3kZ+XZ70I6KQMbJ\\nftMJ5P2P146JdU8RB31SKL9qbZxzR4mA0uKKvUYtDQN+yErUnoOsm9wb9Z+RcAEc\\nZnZWOO02hGdHa7qkkbAxHuH91KnZbk8jnZm2LT7PFz7Y1fd80vSlnSOL7nRkU7B5\\nWXlJy8ECgYA4g0wc0Jq8c1Q0FulMkOQqYRDXaDo34987L+mZ70i/RtdkKjK/IKJ9\\n18UDCyEaDPD0BWBJGPejZkY8UD6FBG/5k7wNIbT7hHLRSRlw4iRmVX2hRVXrXzD8\\nvc86Qyg2iG0JqkMAvRdH40amPKp5bW4VcfcvQo4TSsI972u12rgwtg==\\n-----END RSA PRIVATE KEY-----\\n",
sshPublicKey: "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDJ7JmNYPHk1r1vDV/rUziz5CUtP+TKsHfsBHaPlJv8D2G8lqQs1BCk0L4pV9g2XiN48b7qXAerJ7/+lwL37RJawdbCBpqGMwucC2FHaTwkBqQ0D3FX/AfPRhT304oR61o67FelZk9cIdfTO5A7ztht0rNFuDEAKKszc+GTCvtcu1C1ZmTJcbfgJBdlBaAEhLUWX4+UY6lWk4+dahk8ylsArviL9PPiw/xmGpzp3JrWGttTBjdh8xk1heloDDvBZWOryCt+zT4WXdLEevNNeB7ab55kBnqI8nxxfi5grGLRmbGkz9qAjoyaOrQxTvlXGmrbYtpS8cik59flTGGSuzGX root@aiplatform",
sshPort: 22,
},
{
name: "with custom ssh port",
params: []string{"--ssh-port=2026"},
sshKeyFilePath: SSHAbsolutePath,
sshPrivateKey: "",
sshPublicKey: "",
sshPort: 2026,
},
}

Expand All @@ -68,6 +85,46 @@ func TestSSHPlugin(t *testing.T) {
if plugin.sshPublicKey != test.sshPublicKey {
t.Errorf("Expected sshPublicKey=%s, got %s", test.sshPublicKey, plugin.sshPublicKey)
}
if plugin.sshPort != test.sshPort {
t.Errorf("Expected sshPort=%d, got %d", test.sshPort, plugin.sshPort)
}
})
}
}

func TestGenerateSSHConfigWithCustomPort(t *testing.T) {
job := &batch.Job{
Spec: batch.JobSpec{
Tasks: []batch.TaskSpec{
{
Name: "worker",
Replicas: 1,
Template: v1.PodTemplateSpec{
Spec: v1.PodSpec{},
},
},
},
},
}

port := 2026
config := generateSSHConfig(job, port)

expected := fmt.Sprintf("Port %d", port)
if !strings.Contains(config, expected) {
t.Errorf("expected SSH config to contain %q, got:\n%s", expected, config)
}
}

func TestInvalidSSHPortFallsBackToDefault(t *testing.T) {
params := []string{"--ssh-port=-10"}

pluginInterface := New(pluginsinterface.PluginClientset{}, params)
plugin := pluginInterface.(*sshPlugin)

plugin.validateSSHPort()

if plugin.sshPort != 22 {
t.Errorf("Expected sshPort to fall back to 22, got %d", plugin.sshPort)
}
}
Loading