From 476fd153a5669525847cc6438e8ee2ec954e28df Mon Sep 17 00:00:00 2001
From: Jeremi Piotrowski <jpiotrowski@microsoft.com>
Date: Mon, 5 Aug 2024 10:46:18 +0200
Subject: [PATCH] azure: Switch to managed boot diagnostics for console

This does not require that the user have RBAC permissions to a storage account
to fetch, because it uses SAS keys behind the scenes. The previous approach
used a kola created storage account has Shared Key Access disabled for security
reasons.

Signed-off-by: Jeremi Piotrowski <jpiotrowski@microsoft.com>
---
 platform/api/azure/instance.go | 43 ++++++++--------------------------
 1 file changed, 10 insertions(+), 33 deletions(-)

diff --git a/platform/api/azure/instance.go b/platform/api/azure/instance.go
index dc1f52b96..0be0f8055 100644
--- a/platform/api/azure/instance.go
+++ b/platform/api/azure/instance.go
@@ -19,7 +19,7 @@ import (
 	"encoding/base64"
 	"fmt"
 	"io"
-	"regexp"
+	"net/http"
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
@@ -148,8 +148,7 @@ func (a *API) getVMParameters(name, sshkey, storageAccountURI string, userdata *
 			},
 			DiagnosticsProfile: &armcompute.DiagnosticsProfile{
 				BootDiagnostics: &armcompute.BootDiagnostics{
-					Enabled:    to.Ptr(true),
-					StorageURI: &storageAccountURI,
+					Enabled: to.Ptr(true),
 				},
 			},
 		},
@@ -303,46 +302,24 @@ func (a *API) TerminateInstance(machine *Machine, resourceGroup string) error {
 
 func (a *API) GetConsoleOutput(name, resourceGroup, storageAccount string) ([]byte, error) {
 	vmResourceGroup := a.getVMRG(resourceGroup)
-	vm, err := a.compClient.Get(context.TODO(), vmResourceGroup, name, &armcompute.VirtualMachinesClientGetOptions{
-		Expand: to.Ptr(armcompute.InstanceViewTypesInstanceView),
-	})
+	param := &armcompute.VirtualMachinesClientRetrieveBootDiagnosticsDataOptions{
+		SasURIExpirationTimeInMinutes: to.Ptr[int32](5),
+	}
+	resp, err := a.compClient.RetrieveBootDiagnosticsData(context.TODO(), vmResourceGroup, name, param)
 	if err != nil {
 		return nil, fmt.Errorf("could not get VM: %v", err)
 	}
-
-	consoleURI := vm.Properties.InstanceView.BootDiagnostics.SerialConsoleLogBlobURI
-	if consoleURI == nil {
+	if resp.SerialConsoleLogBlobURI == nil {
 		return nil, fmt.Errorf("serial console URI is nil")
 	}
 
-	// Only the full URI to the logs are present in the virtual machine
-	// properties. Parse out the container & file name to use the GetBlob
-	// API call directly.
-	uri := []byte(*consoleURI)
-	containerPat := regexp.MustCompile(`bootdiagnostics-[a-z0-9\-]+`)
-	container := string(containerPat.Find(uri))
-	if container == "" {
-		return nil, fmt.Errorf("could not find container name in URI: %q", *consoleURI)
-	}
-	namePat := regexp.MustCompile(`[a-z0-9\-\.]+.serialconsole.log`)
-	blobname := string(namePat.Find(uri))
-	if blobname == "" {
-		return nil, fmt.Errorf("could not find blob name in URI: %q", *consoleURI)
-	}
-
-	client, err := a.GetBlobServiceClient(storageAccount)
-	if err != nil {
-		return nil, err
-	}
 	var data io.ReadCloser
 	err = util.Retry(6, 10*time.Second, func() error {
-		data, err = GetBlob(client, container, blobname)
+		reply, err := http.Get(*resp.SerialConsoleLogBlobURI)
 		if err != nil {
-			return fmt.Errorf("could not get blob for container %q, blobname %q: %v", container, blobname, err)
-		}
-		if data == nil {
-			return fmt.Errorf("empty data while getting blob for container %q, blobname %q", container, blobname)
+			return fmt.Errorf("could not GET console output: %v", err)
 		}
+		data = reply.Body
 		return nil
 	})
 	if err != nil {