Made error message on MSI download errors more user friendly, made retries more selective

pkg/download/blobwithmsitoken.go

@@ -57,20 +57,38 @@ func NewBlobWithMsiDownload(url string, msiProvider MsiProvider) Downloader {
 
 func GetMsiProviderForStorageAccountsImplicitly(blobUri string) MsiProvider {
 	msiProvider := msi.NewMsiProvider(httputil.NewSecureHttpClient(httputil.DefaultRetryBehavior))
-	return func() (msi.Msi, error) { return msiProvider.GetMsiForResource(GetResourceNameFromBlobUri(blobUri)) }
+	return func() (msi.Msi, error) {
+		msi, err := msiProvider.GetMsiForResource(GetResourceNameFromBlobUri(blobUri))
+		if err != nil {
+			return msi, errors.Wrapf(err, "Unable to get managed identity. "+
+				"Please make sure that system assigned managed identity is enabled on the VM"+
+				"or user assigned identity is added to the system.")
+		}
+		return msi, nil
+	}
 }
 
 func GetMsiProviderForStorageAccountsWithClientId(blobUri, clientId string) MsiProvider {
 	msiProvider := msi.NewMsiProvider(httputil.NewSecureHttpClient(httputil.DefaultRetryBehavior))
 	return func() (msi.Msi, error) {
-		return msiProvider.GetMsiUsingClientId(clientId, GetResourceNameFromBlobUri(blobUri))
+		msi, err := msiProvider.GetMsiUsingClientId(clientId, GetResourceNameFromBlobUri(blobUri))
+		if err != nil {
+			return msi, errors.Wrapf(err, "Unable to get managed identity with client id %s. "+
+				"Please make sure that the user assigned managed identity is added to the VM ", clientId)
+		}
+		return msi, nil
 	}
 }
 
 func GetMsiProviderForStorageAccountsWithObjectId(blobUri, objectId string) MsiProvider {
 	msiProvider := msi.NewMsiProvider(httputil.NewSecureHttpClient(httputil.DefaultRetryBehavior))
 	return func() (msi.Msi, error) {
-		return msiProvider.GetMsiUsingObjectId(objectId, GetResourceNameFromBlobUri(blobUri))
+		msi, err := msiProvider.GetMsiUsingObjectId(objectId, GetResourceNameFromBlobUri(blobUri))
+		if err != nil {
+			return msi, errors.Wrapf(err, "Unable to get managed identity with object id %s. "+
+				"Please make sure that the user assigned managed identity is added to the VM ", objectId)
+		}
+		return msi, nil
 	}
 }
 

pkg/download/downloader.go

@@ -52,5 +52,13 @@ func Download(d Downloader) (int, io.ReadCloser, error) {
 	if resp.StatusCode == http.StatusOK {
 		return resp.StatusCode, resp.Body, nil
 	}
-	return resp.StatusCode, nil, fmt.Errorf("unexpected status code: got=%d expected=%d", resp.StatusCode, http.StatusOK)
+
+	err = fmt.Errorf("unexpected status code: got=%d expected=%d", resp.StatusCode, http.StatusOK)
+	switch d.(type) {
+	case *blobWithMsiToken:
+		if resp.StatusCode == http.StatusForbidden {
+			return resp.StatusCode, nil, errors.Wrapf(err, "please ensure that the specified Managed Identity has read permissions to the storage blob")
+		}
+	}
+	return resp.StatusCode, nil, err
 }

pkg/download/retry.go

@@ -48,8 +48,8 @@ func WithRetries(ctx *log.Context, downloaders []Downloader, sf SleepFunc) (io.R
 				out.Close()
 			}
 
-			if status == http.StatusForbidden {
-				ctx.Log("info", fmt.Sprintf("downloader %T returned 403, skipping retries", d))
+			if !isTransientHttpStatusCode(status) {
+				ctx.Log("info", fmt.Sprintf("downloader %T returned %v, skipping retries", d, status))
 				break
 			}
 
@@ -63,3 +63,18 @@ func WithRetries(ctx *log.Context, downloaders []Downloader, sf SleepFunc) (io.R
 	}
 	return nil, lastErr
 }
+
+func isTransientHttpStatusCode(statusCode int) bool {
+	switch statusCode {
+	case
+		http.StatusRequestTimeout,      // 408
+		http.StatusTooManyRequests,     // 429
+		http.StatusInternalServerError, // 500
+		http.StatusBadGateway,          // 502
+		http.StatusServiceUnavailable,  // 503
+		http.StatusGatewayTimeout:      // 504
+		return true // timeout and too many requests
+	default:
+		return false
+	}
+}