From 4c9124c4b3d81e5adbf88ec0ae9e523b94b5c14e Mon Sep 17 00:00:00 2001 From: Carlos Eduardo Arango Gutierrez Date: Mon, 26 May 2025 12:58:11 +0200 Subject: [PATCH] Add a safely exit to securly close the ssh connection Signed-off-by: Carlos Eduardo Arango Gutierrez --- .../templates/container-toolkit.go | 3 +++ .../templates/container-toolkit_test.go | 5 +++++ pkg/provisioner/templates/containerd.go | 3 +++ pkg/provisioner/templates/containerd_test.go | 6 +++-- pkg/provisioner/templates/crio.go | 3 +++ pkg/provisioner/templates/crio_test.go | 5 +++++ pkg/provisioner/templates/docker.go | 3 +++ pkg/provisioner/templates/docker_test.go | 5 +++++ pkg/provisioner/templates/kernel.go | 3 +++ pkg/provisioner/templates/kernel_test.go | 4 ++++ pkg/provisioner/templates/kubernetes.go | 9 ++++++++ pkg/provisioner/templates/kubernetes_test.go | 22 ++++++++++++------- pkg/provisioner/templates/nv-driver.go | 3 +++ pkg/provisioner/templates/nv-driver_test.go | 16 +++++++++++++- 14 files changed, 79 insertions(+), 11 deletions(-) diff --git a/pkg/provisioner/templates/container-toolkit.go b/pkg/provisioner/templates/container-toolkit.go index ddf377381..e5145accd 100644 --- a/pkg/provisioner/templates/container-toolkit.go +++ b/pkg/provisioner/templates/container-toolkit.go @@ -39,6 +39,9 @@ install_packages_with_retry nvidia-container-toolkit # Configure container runtime sudo nvidia-ctk runtime configure --runtime={{.ContainerRuntime}} --set-as-default --enable-cdi={{.EnableCDI}} sudo systemctl restart {{.ContainerRuntime}} + +# safely close the ssh connection +exit 0 ` type ContainerToolkit struct { diff --git a/pkg/provisioner/templates/container-toolkit_test.go b/pkg/provisioner/templates/container-toolkit_test.go index 38c7e797b..17396b24a 100644 --- a/pkg/provisioner/templates/container-toolkit_test.go +++ b/pkg/provisioner/templates/container-toolkit_test.go @@ -60,4 +60,9 @@ func TestContainerToolkit_Execute(t *testing.T) { if !strings.Contains(out, "nvidia-ctk runtime configure --runtime=containerd --set-as-default --enable-cdi=true") { t.Errorf("template output missing expected runtime config: %s", out) } + + // Test safe exit + if !strings.Contains(out, "exit 0") { + t.Errorf("template output missing safe exit: %s", out) + } } diff --git a/pkg/provisioner/templates/containerd.go b/pkg/provisioner/templates/containerd.go index b6960312c..17751b8ad 100644 --- a/pkg/provisioner/templates/containerd.go +++ b/pkg/provisioner/templates/containerd.go @@ -315,6 +315,9 @@ sudo ctr images pull docker.io/library/hello-world:latest sudo ctr run --rm docker.io/library/hello-world:latest test echo "Containerd installation completed successfully!" + +# safely close the ssh connection +exit 0 ` type Containerd struct { diff --git a/pkg/provisioner/templates/containerd_test.go b/pkg/provisioner/templates/containerd_test.go index fd8d32dbc..9a80013cd 100644 --- a/pkg/provisioner/templates/containerd_test.go +++ b/pkg/provisioner/templates/containerd_test.go @@ -198,8 +198,10 @@ func TestContainerd_Execute_SystemChecks(t *testing.T) { if !strings.Contains(out, "TMP_DIR=$(mktemp -d)") { t.Error("template output missing temporary directory creation") } - if !strings.Contains(out, "rm -rf $TMP_DIR") { - t.Error("template output missing temporary directory cleanup") + + // Test safe exit + if !strings.Contains(out, "exit 0") { + t.Error("template output missing safe exit") } // Test error handling diff --git a/pkg/provisioner/templates/crio.go b/pkg/provisioner/templates/crio.go index baf57f23b..2858bc097 100644 --- a/pkg/provisioner/templates/crio.go +++ b/pkg/provisioner/templates/crio.go @@ -38,6 +38,9 @@ apt install -y cri-o # Start and enable Service systemctl daemon-reload systemctl start crio.service + +# safely close the ssh connection +exit 0 ` type CriO struct { diff --git a/pkg/provisioner/templates/crio_test.go b/pkg/provisioner/templates/crio_test.go index 16a5c1dad..55c8bb893 100644 --- a/pkg/provisioner/templates/crio_test.go +++ b/pkg/provisioner/templates/crio_test.go @@ -43,4 +43,9 @@ func TestCriO_Execute(t *testing.T) { if !strings.Contains(out, "systemctl start crio.service") { t.Errorf("template output missing crio start: %s", out) } + + // Test safe exit + if !strings.Contains(out, "exit 0") { + t.Errorf("template output missing safe exit: %s", out) + } } diff --git a/pkg/provisioner/templates/docker.go b/pkg/provisioner/templates/docker.go index 0cc033266..8b1163871 100644 --- a/pkg/provisioner/templates/docker.go +++ b/pkg/provisioner/templates/docker.go @@ -73,6 +73,9 @@ sudo systemctl restart docker # Post-installation steps for Linux sudo usermod -aG docker $USER newgrp docker + +# safely close the ssh connection +exit 0 ` type Docker struct { diff --git a/pkg/provisioner/templates/docker_test.go b/pkg/provisioner/templates/docker_test.go index 6988dbbd8..6de4d4ef7 100644 --- a/pkg/provisioner/templates/docker_test.go +++ b/pkg/provisioner/templates/docker_test.go @@ -54,4 +54,9 @@ func TestDocker_Execute(t *testing.T) { if !strings.Contains(out, "systemctl enable docker") { t.Errorf("template output missing enable docker: %s", out) } + + // Test safe exit + if !strings.Contains(out, "exit 0") { + t.Errorf("template output missing safe exit: %s", out) + } } diff --git a/pkg/provisioner/templates/kernel.go b/pkg/provisioner/templates/kernel.go index 293853660..d76e95e99 100644 --- a/pkg/provisioner/templates/kernel.go +++ b/pkg/provisioner/templates/kernel.go @@ -60,6 +60,9 @@ if [ "${CURRENT_KERNEL}" != "${KERNEL_VERSION}" ]; then # Run the reboot command with nohup to avoid abrupt SSH closure issues nohup sudo reboot & fi + +# safely close the ssh connection +exit 0 {{- end }} ` diff --git a/pkg/provisioner/templates/kernel_test.go b/pkg/provisioner/templates/kernel_test.go index 8da7231c3..30e1cc510 100644 --- a/pkg/provisioner/templates/kernel_test.go +++ b/pkg/provisioner/templates/kernel_test.go @@ -155,6 +155,10 @@ func TestKernelTemplateContent(t *testing.T) { name: "reboot command", contains: "nohup sudo reboot", }, + { + name: "safe exit", + contains: "# safely close the ssh connection\nexit 0", + }, } for _, tt := range tests { diff --git a/pkg/provisioner/templates/kubernetes.go b/pkg/provisioner/templates/kubernetes.go index 7fee9bac3..9398e23d4 100644 --- a/pkg/provisioner/templates/kubernetes.go +++ b/pkg/provisioner/templates/kubernetes.go @@ -110,6 +110,9 @@ with_retry 5 10s kubectl --kubeconfig $KUBECONFIG create -f https://raw.githubus kubectl taint nodes --all node-role.kubernetes.io/control-plane:NoSchedule- kubectl label node --all node-role.kubernetes.io/worker= kubectl label node --all nvidia.com/holodeck.managed=true + +# safely close the ssh connection +exit 0 ` const KindTemplate = ` @@ -152,6 +155,9 @@ with_retry 3 10s kind create cluster --name holodeck $KIND_CONFIG --kubeconfig=" echo "KIND installed successfully" echo "you can now access the cluster with:" echo "ssh -i ubuntu@${INSTANCE_ENDPOINT_HOST}" + +# safely close the ssh connection +exit 0 ` const microk8sTemplate = ` @@ -173,6 +179,9 @@ sudo snap alias microk8s.kubectl kubectl echo "Microk8s {{.Version}} installed successfully" echo "you can now access the cluster with:" echo "ssh -i ubuntu@${INSTANCE_ENDPOINT_HOST}" + +# safely close the ssh connection +exit 0 ` const kubeadmTemplate = `apiVersion: kubeadm.k8s.io/v1beta4 diff --git a/pkg/provisioner/templates/kubernetes_test.go b/pkg/provisioner/templates/kubernetes_test.go index 40c871309..41bf7dd82 100644 --- a/pkg/provisioner/templates/kubernetes_test.go +++ b/pkg/provisioner/templates/kubernetes_test.go @@ -178,6 +178,7 @@ func TestKubernetes_Execute(t *testing.T) { wantErr bool checkTemplate bool expectedString string + checkSafeExit bool }{ { name: "kubeadm installer", @@ -192,7 +193,8 @@ func TestKubernetes_Execute(t *testing.T) { }, }, }, - wantErr: false, + wantErr: false, + checkSafeExit: true, }, { name: "legacy kubeadm installer", @@ -211,6 +213,7 @@ func TestKubernetes_Execute(t *testing.T) { wantErr: false, checkTemplate: true, expectedString: "kubeadm init \\\n --kubernetes-version=${K8S_VERSION} \\\n --pod-network-cidr=192.168.0.0/16 \\\n --control-plane-endpoint=test-host:6443 \\\n --ignore-preflight-errors=all", + checkSafeExit: true, }, { name: "kind installer", @@ -225,7 +228,8 @@ func TestKubernetes_Execute(t *testing.T) { }, }, }, - wantErr: false, + wantErr: false, + checkSafeExit: true, }, { name: "microk8s installer", @@ -240,7 +244,8 @@ func TestKubernetes_Execute(t *testing.T) { }, }, }, - wantErr: false, + wantErr: false, + checkSafeExit: true, }, { name: "invalid installer", @@ -276,13 +281,14 @@ func TestKubernetes_Execute(t *testing.T) { return } assert.NoError(t, err) - assert.NotEmpty(t, buf.String()) + out := buf.String() if tt.checkTemplate { - // Check if the template contains the expected kubeadm init command - assert.Contains(t, buf.String(), tt.expectedString) - // Verify that it doesn't use the config file - assert.NotContains(t, buf.String(), "kubeadm init --config /etc/kubernetes/kubeadm-config.yaml") + assert.Contains(t, out, tt.expectedString) + } + + if tt.checkSafeExit { + assert.Contains(t, out, "exit 0", "template output missing safe exit") } }) } diff --git a/pkg/provisioner/templates/nv-driver.go b/pkg/provisioner/templates/nv-driver.go index 2436a7672..fce3d3417 100644 --- a/pkg/provisioner/templates/nv-driver.go +++ b/pkg/provisioner/templates/nv-driver.go @@ -46,6 +46,9 @@ sudo nvidia-persistenced --persistence-mode # Quick check to see if the driver is installed nvidia-smi + +# safely close the ssh connection +exit 0 ` type NvDriver v1alpha1.NVIDIADriver diff --git a/pkg/provisioner/templates/nv-driver_test.go b/pkg/provisioner/templates/nv-driver_test.go index 40ce51ef8..502bfed9a 100644 --- a/pkg/provisioner/templates/nv-driver_test.go +++ b/pkg/provisioner/templates/nv-driver_test.go @@ -19,6 +19,7 @@ package templates import ( "bytes" + "strings" "testing" "github.com/stretchr/testify/require" @@ -59,6 +60,9 @@ sudo nvidia-persistenced --persistence-mode # Quick check to see if the driver is installed nvidia-smi + +# safely close the ssh connection +exit 0 `, }, { @@ -87,6 +91,9 @@ sudo nvidia-persistenced --persistence-mode # Quick check to see if the driver is installed nvidia-smi + +# safely close the ssh connection +exit 0 `, }, { @@ -116,6 +123,9 @@ sudo nvidia-persistenced --persistence-mode # Quick check to see if the driver is installed nvidia-smi + +# safely close the ssh connection +exit 0 `, }, } @@ -128,7 +138,11 @@ nvidia-smi err := tc.driver.Execute(&output, v1alpha1.Environment{}) require.EqualValues(t, tc.expecteError, err) - require.EqualValues(t, tc.expectedOutput, output.String()) + // Compare trimmed strings to avoid whitespace issues + require.EqualValues(t, strings.TrimSpace(tc.expectedOutput), strings.TrimSpace(output.String())) + + // Test safe exit + require.Contains(t, output.String(), "exit 0", "template output missing safe exit") }) }