From 7ca2e2436fa697c568f5996e74f4e34473875378 Mon Sep 17 00:00:00 2001
From: Cecile Robert-Michon <cerobert@microsoft.com>
Date: Fri, 29 Oct 2021 11:37:13 -0700
Subject: [PATCH] Re-add defaulting for AzureMachineTemplate ssh key

---
 api/v1beta1/azuremachinetemplate_webhook.go   | 10 ++--
 .../azuremachinetemplate_webhook_test.go      | 53 +++++++++++++++++--
 2 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/api/v1beta1/azuremachinetemplate_webhook.go b/api/v1beta1/azuremachinetemplate_webhook.go
index a19819b1649..4ca25d9deac 100644
--- a/api/v1beta1/azuremachinetemplate_webhook.go
+++ b/api/v1beta1/azuremachinetemplate_webhook.go
@@ -69,6 +69,12 @@ func (r *AzureMachineTemplate) ValidateUpdate(oldRaw runtime.Object) error {
 		// This means if the old object was in v1alpha3, it would not get the new defaults set in v1beta1 resulting
 		// in object inequality. To workaround this, we set the v1beta1 defaults here so that the old object also gets
 		// the new defaults.
+
+		// We need to set ssh key explicitly, otherwise Default() will create a new one.
+		if old.Spec.Template.Spec.SSHPublicKey == "" {
+			old.Spec.Template.Spec.SSHPublicKey = r.Spec.Template.Spec.SSHPublicKey
+		}
+
 		old.Default()
 
 		// if it's still not equal, return error.
@@ -93,7 +99,5 @@ func (r *AzureMachineTemplate) ValidateDelete() error {
 // Default implements webhookutil.defaulter so a webhook will be registered for the type.
 func (r *AzureMachineTemplate) Default() {
 	machinetemplatelog.Info("default", "name", r.Name)
-	r.Spec.Template.Spec.SetDefaultCachingType()
-	r.Spec.Template.Spec.SetDataDisksDefaults()
-	r.Spec.Template.Spec.SetIdentityDefaults()
+	r.Spec.Template.Spec.SetDefaults(machinetemplatelog)
 }
diff --git a/api/v1beta1/azuremachinetemplate_webhook_test.go b/api/v1beta1/azuremachinetemplate_webhook_test.go
index 596929ed498..bce60acc324 100644
--- a/api/v1beta1/azuremachinetemplate_webhook_test.go
+++ b/api/v1beta1/azuremachinetemplate_webhook_test.go
@@ -177,7 +177,7 @@ func TestAzureMachineTemplate_ValidateUpdate(t *testing.T) {
 								DiskSizeGB: to.Int32Ptr(11),
 							},
 							DataDisks:    []DataDisk{},
-							SSHPublicKey: "",
+							SSHPublicKey: "fake ssh key",
 						},
 					},
 				},
@@ -197,7 +197,7 @@ func TestAzureMachineTemplate_ValidateUpdate(t *testing.T) {
 								DiskSizeGB: to.Int32Ptr(11),
 							},
 							DataDisks:    []DataDisk{},
-							SSHPublicKey: "",
+							SSHPublicKey: "fake ssh key",
 						},
 					},
 				},
@@ -216,7 +216,7 @@ func TestAzureMachineTemplate_ValidateUpdate(t *testing.T) {
 								DiskSizeGB: to.Int32Ptr(11),
 							},
 							DataDisks:    []DataDisk{},
-							SSHPublicKey: "",
+							SSHPublicKey: "fake ssh key",
 						},
 					},
 				},
@@ -259,7 +259,8 @@ func TestAzureMachineTemplate_ValidateUpdate(t *testing.T) {
 								DiskSizeGB:  to.Int32Ptr(11),
 								CachingType: "None",
 							},
-							DataDisks: []DataDisk{},
+							DataDisks:    []DataDisk{},
+							SSHPublicKey: "fake ssh key",
 						},
 					},
 				},
@@ -269,6 +270,50 @@ func TestAzureMachineTemplate_ValidateUpdate(t *testing.T) {
 			},
 			wantErr: false,
 		},
+		{
+			name: "AzureMachineTemplate ssh key removed",
+			oldTemplate: &AzureMachineTemplate{
+				Spec: AzureMachineTemplateSpec{
+					Template: AzureMachineTemplateResource{
+						Spec: AzureMachineSpec{
+							VMSize:        "size",
+							FailureDomain: &failureDomain,
+							OSDisk: OSDisk{
+								OSType:      "type",
+								DiskSizeGB:  to.Int32Ptr(11),
+								CachingType: "None",
+							},
+							DataDisks:    []DataDisk{},
+							SSHPublicKey: "some key",
+						},
+					},
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "OldTemplate",
+				},
+			},
+			template: &AzureMachineTemplate{
+				Spec: AzureMachineTemplateSpec{
+					Template: AzureMachineTemplateResource{
+						Spec: AzureMachineSpec{
+							VMSize:        "size",
+							FailureDomain: &failureDomain,
+							OSDisk: OSDisk{
+								OSType:      "type",
+								DiskSizeGB:  to.Int32Ptr(11),
+								CachingType: "None",
+							},
+							DataDisks:    []DataDisk{},
+							SSHPublicKey: "",
+						},
+					},
+				},
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "NewTemplate",
+				},
+			},
+			wantErr: true,
+		},
 	}
 
 	for _, amt := range tests {