From 0e8a2e44a4c54f70d1de052d52bb1304148ddfa5 Mon Sep 17 00:00:00 2001
From: Darshan Chaudhary <deathbullet@gmail.com>
Date: Wed, 14 Feb 2024 22:29:01 +0530
Subject: [PATCH] Azure SD: check for nil before accessing and dereferencing 
 (#13578)

Signed-off-by: darshanime <deathbullet@gmail.com>
---
 discovery/azure/azure.go      |  16 +++++-
 discovery/azure/azure_test.go | 101 ++++++++++++++++++++++++++++++++++
 2 files changed, 115 insertions(+), 2 deletions(-)

diff --git a/discovery/azure/azure.go b/discovery/azure/azure.go
index e5e6109b2..a5d81f4ff 100644
--- a/discovery/azure/azure.go
+++ b/discovery/azure/azure.go
@@ -569,7 +569,7 @@ func (client *azureClient) getScaleSetVMs(ctx context.Context, scaleSet armcompu
 }
 
 func mapFromVM(vm armcompute.VirtualMachine) virtualMachine {
-	osType := string(*vm.Properties.StorageProfile.OSDisk.OSType)
+	var osType string
 	tags := map[string]*string{}
 	networkInterfaces := []string{}
 	var computerName string
@@ -580,6 +580,12 @@ func mapFromVM(vm armcompute.VirtualMachine) virtualMachine {
 	}
 
 	if vm.Properties != nil {
+		if vm.Properties.StorageProfile != nil &&
+			vm.Properties.StorageProfile.OSDisk != nil &&
+			vm.Properties.StorageProfile.OSDisk.OSType != nil {
+			osType = string(*vm.Properties.StorageProfile.OSDisk.OSType)
+		}
+
 		if vm.Properties.NetworkProfile != nil {
 			for _, vmNIC := range vm.Properties.NetworkProfile.NetworkInterfaces {
 				networkInterfaces = append(networkInterfaces, *vmNIC.ID)
@@ -608,7 +614,7 @@ func mapFromVM(vm armcompute.VirtualMachine) virtualMachine {
 }
 
 func mapFromVMScaleSetVM(vm armcompute.VirtualMachineScaleSetVM, scaleSetName string) virtualMachine {
-	osType := string(*vm.Properties.StorageProfile.OSDisk.OSType)
+	var osType string
 	tags := map[string]*string{}
 	networkInterfaces := []string{}
 	var computerName string
@@ -619,6 +625,12 @@ func mapFromVMScaleSetVM(vm armcompute.VirtualMachineScaleSetVM, scaleSetName st
 	}
 
 	if vm.Properties != nil {
+		if vm.Properties.StorageProfile != nil &&
+			vm.Properties.StorageProfile.OSDisk != nil &&
+			vm.Properties.StorageProfile.OSDisk.OSType != nil {
+			osType = string(*vm.Properties.StorageProfile.OSDisk.OSType)
+		}
+
 		if vm.Properties.NetworkProfile != nil {
 			for _, vmNIC := range vm.Properties.NetworkProfile.NetworkInterfaces {
 				networkInterfaces = append(networkInterfaces, *vmNIC.ID)
diff --git a/discovery/azure/azure_test.go b/discovery/azure/azure_test.go
index 4ff937e0b..1e437c75f 100644
--- a/discovery/azure/azure_test.go
+++ b/discovery/azure/azure_test.go
@@ -79,6 +79,55 @@ func TestMapFromVMWithEmptyTags(t *testing.T) {
 	require.Equal(t, expectedVM, actualVM)
 }
 
+func TestMapFromVMWithEmptyOSType(t *testing.T) {
+	id := "test"
+	name := "name"
+	size := "size"
+	vmSize := armcompute.VirtualMachineSizeTypes(size)
+	vmType := "type"
+	location := "westeurope"
+	computerName := "computer_name"
+	networkProfile := armcompute.NetworkProfile{
+		NetworkInterfaces: []*armcompute.NetworkInterfaceReference{},
+	}
+	properties := &armcompute.VirtualMachineProperties{
+		OSProfile: &armcompute.OSProfile{
+			ComputerName: &computerName,
+		},
+		StorageProfile: &armcompute.StorageProfile{
+			OSDisk: &armcompute.OSDisk{},
+		},
+		NetworkProfile: &networkProfile,
+		HardwareProfile: &armcompute.HardwareProfile{
+			VMSize: &vmSize,
+		},
+	}
+
+	testVM := armcompute.VirtualMachine{
+		ID:         &id,
+		Name:       &name,
+		Type:       &vmType,
+		Location:   &location,
+		Tags:       nil,
+		Properties: properties,
+	}
+
+	expectedVM := virtualMachine{
+		ID:                id,
+		Name:              name,
+		ComputerName:      computerName,
+		Type:              vmType,
+		Location:          location,
+		Tags:              map[string]*string{},
+		NetworkInterfaces: []string{},
+		Size:              size,
+	}
+
+	actualVM := mapFromVM(testVM)
+
+	require.Equal(t, expectedVM, actualVM)
+}
+
 func TestMapFromVMWithTags(t *testing.T) {
 	id := "test"
 	name := "name"
@@ -193,6 +242,58 @@ func TestMapFromVMScaleSetVMWithEmptyTags(t *testing.T) {
 	require.Equal(t, expectedVM, actualVM)
 }
 
+func TestMapFromVMScaleSetVMWithEmptyOSType(t *testing.T) {
+	id := "test"
+	name := "name"
+	size := "size"
+	vmSize := armcompute.VirtualMachineSizeTypes(size)
+	vmType := "type"
+	instanceID := "123"
+	location := "westeurope"
+	computerName := "computer_name"
+	networkProfile := armcompute.NetworkProfile{
+		NetworkInterfaces: []*armcompute.NetworkInterfaceReference{},
+	}
+	properties := &armcompute.VirtualMachineScaleSetVMProperties{
+		OSProfile: &armcompute.OSProfile{
+			ComputerName: &computerName,
+		},
+		StorageProfile: &armcompute.StorageProfile{},
+		NetworkProfile: &networkProfile,
+		HardwareProfile: &armcompute.HardwareProfile{
+			VMSize: &vmSize,
+		},
+	}
+
+	testVM := armcompute.VirtualMachineScaleSetVM{
+		ID:         &id,
+		Name:       &name,
+		Type:       &vmType,
+		InstanceID: &instanceID,
+		Location:   &location,
+		Tags:       nil,
+		Properties: properties,
+	}
+
+	scaleSet := "testSet"
+	expectedVM := virtualMachine{
+		ID:                id,
+		Name:              name,
+		ComputerName:      computerName,
+		Type:              vmType,
+		Location:          location,
+		Tags:              map[string]*string{},
+		NetworkInterfaces: []string{},
+		ScaleSet:          scaleSet,
+		InstanceID:        instanceID,
+		Size:              size,
+	}
+
+	actualVM := mapFromVMScaleSetVM(testVM, scaleSet)
+
+	require.Equal(t, expectedVM, actualVM)
+}
+
 func TestMapFromVMScaleSetVMWithTags(t *testing.T) {
 	id := "test"
 	name := "name"