fixed passing by value vs ref bug and added UTs

Author: alsanmsft

internal/cmds/cmds.go

@@ -221,7 +221,7 @@ func enable(ctx *log.Context, h types.HandlerEnvironment, report *types.RunComma
 	var rceps *extensionpolicysettingsrc.RCv2ExtensionPolicySettings
 
 	if _, err := os.Stat(policyPath); err == nil {
-		err = extensionpolicysettingsrc.InitializeExtensionPolicySettings(ExtensionPolicyManagerPtr, policyPath, rceps)
+		ExtensionPolicyManagerPtr, rceps, err = extensionpolicysettingsrc.InitializeExtensionPolicySettings(policyPath)
 		if err != nil {
 			return "", "", errors.Wrap(err, "failed to initialize extension policy settings"), constants.ExitCode_LoadExtensionPolicySettingsFailed
 		}

internal/cmds/cmds_test.go

@@ -1,6 +1,8 @@
 package commands
 
 import (
+	"crypto/sha256"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"io/ioutil"
@@ -17,6 +19,7 @@ import (
 	"github.com/Azure/azure-extension-platform/pkg/handlerenv"
 	"github.com/Azure/azure-extension-platform/pkg/logging"
 	"github.com/Azure/run-command-handler-linux/internal/constants"
+	"github.com/Azure/run-command-handler-linux/internal/extensionpolicysettingsrc"
 	"github.com/Azure/run-command-handler-linux/internal/files"
 	"github.com/Azure/run-command-handler-linux/internal/handlersettings"
 	"github.com/Azure/run-command-handler-linux/internal/settings"
@@ -1445,3 +1448,77 @@ func mustReadFile(t *testing.T, p string) string {
 	}
 	return string(b)
 }
+
+// Test_downloadScript_BlockedByAllowlist verifies that downloadScript returns an error
+// when the policy allows downloaded scripts (alloweddownloaded) but the script's
+// SHA256 hash is not in the DownloadedScriptsAllowlist.
+func Test_downloadScript_BlockedByAllowlist(t *testing.T) {
+	dir, err := ioutil.TempDir("", "")
+	require.Nil(t, err)
+	defer os.RemoveAll(dir)
+
+	scriptContent := []byte("#!/bin/bash\necho hello\n")
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write(scriptContent)
+	}))
+	defer srv.Close()
+
+	policy := &extensionpolicysettingsrc.RCv2ExtensionPolicySettings{
+		LimitScripts: "alloweddownloaded",
+		// A wrong hash — the script's actual hash is not this.
+		DownloadedScriptsAllowlist: []string{"0000000000000000000000000000000000000000000000000000000000000000"},
+	}
+
+	_, err = downloadScript(log.NewContext(log.NewNopLogger()),
+		dir,
+		&handlersettings.HandlerSettings{
+			PublicSettings: handlersettings.PublicSettings{
+				Source: &handlersettings.ScriptSource{ScriptURI: srv.URL + "/script.sh"},
+			},
+		},
+		policy,
+	)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "blocked by policy")
+	require.Contains(t, err.Error(), "item is not in the allowlist")
+}
+
+// Test_downloadScript_AllowedByAllowlist verifies that downloadScript succeeds
+// when the policy allows downloaded scripts and the script's SHA256 hash IS
+// present in the DownloadedScriptsAllowlist.
+func Test_downloadScript_AllowedByAllowlist(t *testing.T) {
+	dir, err := os.MkdirTemp("", "")
+	require.Nil(t, err)
+	defer os.RemoveAll(dir)
+
+	// Content uses Unix LF only and has no BOM, so PostProcessFile leaves bytes
+	// unchanged, making the pre-computed hash match the on-disk file hash.
+	scriptContent := []byte("#!/bin/bash\necho hello\n")
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write(scriptContent)
+	}))
+	defer srv.Close()
+
+	// Compute the SHA256 hash that ValidateFileHashInAllowlist will compare against.
+	h := sha256.New()
+	h.Write(scriptContent)
+	correctHash := hex.EncodeToString(h.Sum(nil))
+
+	policy := &extensionpolicysettingsrc.RCv2ExtensionPolicySettings{
+		LimitScripts:               "alloweddownloaded",
+		DownloadedScriptsAllowlist: []string{correctHash},
+	}
+
+	_, err = downloadScript(log.NewContext(log.NewNopLogger()),
+		dir,
+		&handlersettings.HandlerSettings{
+			PublicSettings: handlersettings.PublicSettings{
+				Source: &handlersettings.ScriptSource{ScriptURI: srv.URL + "/script.sh"},
+			},
+		},
+		policy,
+	)
+	require.NoError(t, err)
+}

internal/extensionpolicysettingsrc/extensionpolicysettingsrc.go

@@ -9,25 +9,26 @@ import (
 	"github.com/pkg/errors"
 )
 
-func InitializeExtensionPolicySettings(ExtensionPolicyManagerPtr *extensionpolicysettings.ExtensionPolicySettingsManager[RCv2ExtensionPolicySettings],
-	policyPath string,
-	rceps *RCv2ExtensionPolicySettings) error {
+func InitializeExtensionPolicySettings(policyPath string) (*extensionpolicysettings.ExtensionPolicySettingsManager[RCv2ExtensionPolicySettings], *RCv2ExtensionPolicySettings, error) {
+	var ExtensionPolicyManagerPtr *extensionpolicysettings.ExtensionPolicySettingsManager[RCv2ExtensionPolicySettings]
+	var rceps *RCv2ExtensionPolicySettings
+
 	ExtensionPolicyManagerPtr, err := extensionpolicysettings.NewExtensionPolicySettingsManager[RCv2ExtensionPolicySettings](policyPath)
 	if err != nil {
-		return errors.Wrap(err, "failed to create extension policy settings manager")
+		return nil, nil, errors.Wrap(err, "failed to create extension policy settings manager")
 	}
 
 	err = ExtensionPolicyManagerPtr.LoadExtensionPolicySettings()
 	if err != nil {
-		return errors.Wrap(err, "failed to load extension policy settings")
+		return nil, nil, errors.Wrap(err, "failed to load extension policy settings")
 	} else {
 		rceps, err = ExtensionPolicyManagerPtr.GetSettings()
 
 		if err != nil {
-			return errors.Wrap(err, "failed to get extension policy settings")
+			return nil, nil, errors.Wrap(err, "failed to get extension policy settings")
 		}
 	}
-	return nil
+	return ExtensionPolicyManagerPtr, rceps, nil
 }
 
 func InitialValidateHandlerSettingsAgainstPolicy(settings *handlersettings.HandlerSettings, policy *RCv2ExtensionPolicySettings) error {
@@ -47,9 +48,9 @@ func InitialValidateHandlerSettingsAgainstPolicy(settings *handlersettings.Handl
 			return err
 		}
 	}
-	if policy.DisableOutputBlobs {
-		ValidateOutputBlob(settings, policy)
-	}
+
+	// TO-DO: Validate Disable Outputblob and RequireSigning once those features are implemented for RCv2.
+
 	return nil
 }
 
@@ -71,7 +72,11 @@ func ValidateCommandId(settings *handlersettings.HandlerSettings, policy *RCv2Ex
 		// if list is empty, all commandIds are allowed
 		return nil
 	}
-	return extensionpolicysettings.ValidateValueInAllowlist(settingsCommandId, allowedCommandIds)
+	err := extensionpolicysettings.ValidateValueInAllowlist(settingsCommandId, allowedCommandIds)
+	if err != nil {
+		return errors.Wrapf(err, "command ID %s is not allowed by policy", settingsCommandId)
+	}
+	return nil
 }
 
 func ValidateRunAsUser(settings *handlersettings.HandlerSettings, policy *RCv2ExtensionPolicySettings) error {
@@ -83,14 +88,3 @@ func ValidateRunAsUser(settings *handlersettings.HandlerSettings, policy *RCv2Ex
 	}
 	return nil
 }
-
-func ValidateOutputBlob(settings *handlersettings.HandlerSettings, policy *RCv2ExtensionPolicySettings) {
-	if policy.DisableOutputBlobs {
-		// Log a warning that output blobs are disabled by policy. The command will still execute, but no output blobs will be created.
-		if settings.OutputBlobURI != "" {
-			fmt.Println("Warning: Output blobs are disabled by policy. The provided output blob URI will be ignored and no output blobs will be created for this command.")
-		} else {
-			fmt.Println("Warning: Output blobs are disabled by policy. No output blobs will be created for this command.")
-		}
-	}
-}

internal/extensionpolicysettingsrc/extensionpolicysettingsrc_test.go

@@ -6,7 +6,6 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/Azure/azure-extension-platform/pkg/extensionpolicysettings"
 	"github.com/Azure/run-command-handler-linux/internal/handlersettings"
 	"github.com/stretchr/testify/require"
 )
@@ -44,10 +43,7 @@ func captureStdout(t *testing.T, fn func()) string {
 }
 
 func TestInitializeExtensionPolicySettings_InvalidPath_ReturnsError(t *testing.T) {
-	var mgr *extensionpolicysettings.ExtensionPolicySettingsManager[RCv2ExtensionPolicySettings]
-	out := &RCv2ExtensionPolicySettings{}
-
-	err := InitializeExtensionPolicySettings(mgr, "/definitely/not/found/policy.json", out)
+	_, _, err := InitializeExtensionPolicySettings("/definitely/not/found/policy.json")
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "failed to")
 }
@@ -60,10 +56,7 @@ func TestInitializeExtensionPolicySettings_ValidFile_ReturnsNil(t *testing.T) {
 	err := os.WriteFile(policyPath, []byte("{}"), 0600)
 	require.NoError(t, err)
 
-	var mgr *extensionpolicysettings.ExtensionPolicySettingsManager[RCv2ExtensionPolicySettings]
-	out := &RCv2ExtensionPolicySettings{}
-
-	err = InitializeExtensionPolicySettings(mgr, policyPath, out)
+	_, _, err = InitializeExtensionPolicySettings(policyPath)
 	require.NoError(t, err)
 }
 
@@ -75,17 +68,16 @@ func TestInitializeExtensionPolicySettings_CurrentBehavior_DoesNotPopulateOutput
 	err := os.WriteFile(policyPath, []byte(payload), 0600)
 	require.NoError(t, err)
 
-	var mgr *extensionpolicysettings.ExtensionPolicySettingsManager[RCv2ExtensionPolicySettings]
 	out := &RCv2ExtensionPolicySettings{}
 
-	err = InitializeExtensionPolicySettings(mgr, policyPath, out)
+	_, out, err = InitializeExtensionPolicySettings(policyPath)
 	require.NoError(t, err)
 
-	// Documents current implementation behavior (pointer reassignment inside function).
-	require.Equal(t, "", out.LimitScripts)
-	require.Equal(t, "", out.RunAsUser)
+	require.Equal(t, "inline", out.LimitScripts)
+	require.Equal(t, "alice", out.RunAsUser)
 }
 
+// Test that validation passes and fails as expected.
 func TestInitialValidateHandlerSettingsAgainstPolicy(t *testing.T) {
 	t.Run("nil policy", func(t *testing.T) {
 		settings := makeSettings(handlersettings.InlineScript, "", "", "")
@@ -94,6 +86,8 @@ func TestInitialValidateHandlerSettingsAgainstPolicy(t *testing.T) {
 		require.Contains(t, err.Error(), "no policy provided")
 	})
 
+	// This test mimicks running an inline script, but policy only allows gallery scripts.
+	// Validation fails.
 	t.Run("script type blocked by policy", func(t *testing.T) {
 		settings := makeSettings(handlersettings.InlineScript, "", "", "")
 		policy := &RCv2ExtensionPolicySettings{
@@ -105,7 +99,9 @@ func TestInitialValidateHandlerSettingsAgainstPolicy(t *testing.T) {
 		require.Contains(t, err.Error(), "script type inline is not allowed by policy")
 	})
 
-	t.Run("command id not in allowlist", func(t *testing.T) {
+	// This test mimicks running a commandId that is not in the allowlist.
+	// Additionally, only commandId types are allowed.
+	t.Run("command ID not in allowlist", func(t *testing.T) {
 		settings := makeSettings(handlersettings.CommandIdScript, "restartVM", "", "")
 		policy := &RCv2ExtensionPolicySettings{
 			LimitScripts:       "allowedcommandid",
@@ -128,7 +124,33 @@ func TestInitialValidateHandlerSettingsAgainstPolicy(t *testing.T) {
 		require.Contains(t, err.Error(), "does not match")
 	})
 
-	t.Run("all checks pass", func(t *testing.T) {
+	t.Run("enforce limitScripts must be set. If not set, all commands fail", func(t *testing.T) {
+		settings := makeSettings(handlersettings.CommandIdScript, "safeCommand", " Alice ", "https://example/blob")
+		policy := &RCv2ExtensionPolicySettings{
+			LimitScripts:       "",
+			CommandIdAllowlist: []string{"safeCommand"},
+			RunAsUser:          "Alice",
+			DisableOutputBlobs: true,
+		}
+
+		err := InitialValidateHandlerSettingsAgainstPolicy(settings, policy)
+		require.Contains(t, err.Error(), "script type commandId is not allowed by policy")
+	})
+
+	t.Run("all checks pass commandId", func(t *testing.T) {
+		settings := makeSettings(handlersettings.CommandIdScript, "safeCommand", " Alice ", "https://example/blob")
+		policy := &RCv2ExtensionPolicySettings{
+			LimitScripts:       "allowall",
+			CommandIdAllowlist: []string{"safeCommand"},
+			RunAsUser:          "alice",
+			DisableOutputBlobs: true,
+		}
+
+		err := InitialValidateHandlerSettingsAgainstPolicy(settings, policy)
+		require.NoError(t, err)
+	})
+
+	t.Run("all checks pass commandId", func(t *testing.T) {
 		settings := makeSettings(handlersettings.CommandIdScript, "safeCommand", " Alice ", "https://example/blob")
 		policy := &RCv2ExtensionPolicySettings{
 			LimitScripts:       "allowall",
@@ -140,6 +162,19 @@ func TestInitialValidateHandlerSettingsAgainstPolicy(t *testing.T) {
 		err := InitialValidateHandlerSettingsAgainstPolicy(settings, policy)
 		require.NoError(t, err)
 	})
+
+	t.Run("all checks pass downloadedScript", func(t *testing.T) {
+		settings := makeSettings(handlersettings.DownloadedScript, "safeCommand", " Alice ", "https://example/blob")
+		policy := &RCv2ExtensionPolicySettings{
+			LimitScripts:       "alloweddownloaded",
+			CommandIdAllowlist: []string{"safeCommand"},
+			RunAsUser:          "alice",
+			DisableOutputBlobs: true,
+		}
+
+		err := InitialValidateHandlerSettingsAgainstPolicy(settings, policy)
+		require.NoError(t, err)
+	})
 }
 
 func TestValidateScriptTypeAgainstPolicy(t *testing.T) {
@@ -154,7 +189,8 @@ func TestValidateScriptTypeAgainstPolicy(t *testing.T) {
 		require.Contains(t, err.Error(), "script type gallery is not allowed by policy")
 	})
 
-	t.Run("invalid policy token currently treated as blocked", func(t *testing.T) {
+	// This tests edge case where policy has an invalid script type token.
+	t.Run("invalid policy token is treated as blocked", func(t *testing.T) {
 		err := ValidateScriptTypeAgainstPolicy(handlersettings.InlineScript, "notARealScriptType")
 		require.Error(t, err)
 		require.Contains(t, err.Error(), "script type inline is not allowed by policy")
@@ -186,7 +222,8 @@ func TestValidateCommandId(t *testing.T) {
 			CommandIdAllowlist: []string{"safeCommand", "other"},
 		}
 		err := ValidateCommandId(settings, policy)
-		require.Error(t, err)
+		require.Contains(t, err.Error(), "command ID restartVM is not allowed by policy")
+		require.Contains(t, err.Error(), "item is not in the allowlist")
 	})
 }
 
@@ -207,46 +244,6 @@ func TestValidateRunAsUser(t *testing.T) {
 		}
 		err := ValidateRunAsUser(settings, policy)
 		require.Error(t, err)
-		require.Contains(t, err.Error(), "does not match")
-	})
-}
-
-func TestValidateOutputBlob(t *testing.T) {
-	t.Run("policy does not disable output blobs prints nothing", func(t *testing.T) {
-		settings := makeSettings(handlersettings.InlineScript, "", "", "https://example/blob")
-		policy := &RCv2ExtensionPolicySettings{
-			DisableOutputBlobs: false,
-		}
-
-		out := captureStdout(t, func() {
-			ValidateOutputBlob(settings, policy)
-		})
-		require.Equal(t, "", out)
-	})
-
-	t.Run("disabled with output blob uri prints ignore warning", func(t *testing.T) {
-		settings := makeSettings(handlersettings.InlineScript, "", "", "https://example/blob")
-		policy := &RCv2ExtensionPolicySettings{
-			DisableOutputBlobs: true,
-		}
-
-		out := captureStdout(t, func() {
-			ValidateOutputBlob(settings, policy)
-		})
-		require.Contains(t, out, "Output blobs are disabled by policy")
-		require.Contains(t, out, "provided output blob URI will be ignored")
-	})
-
-	t.Run("disabled without output blob uri prints no blob warning", func(t *testing.T) {
-		settings := makeSettings(handlersettings.InlineScript, "", "", "")
-		policy := &RCv2ExtensionPolicySettings{
-			DisableOutputBlobs: true,
-		}
-
-		out := captureStdout(t, func() {
-			ValidateOutputBlob(settings, policy)
-		})
-		require.Contains(t, out, "Output blobs are disabled by policy")
-		require.Contains(t, out, "No output blobs will be created")
+		require.Contains(t, err.Error(), "RunAsUser 'bob' in settings does not match RunAsUser 'alice' in policy")
 	})
 }

internal/extensionpolicysettingsrc/types.go

@@ -86,6 +86,7 @@ func (rceps RCv2ExtensionPolicySettings) ValidateFormat() error {
 	flag, err := StringToAllowedScriptTypeFlag(string(rceps.LimitScripts))
 	// Requirements:
 	// 1. If RequireSigning is not "none", FileRootCert must be present and non-empty.
+	// 	TO-DO: implement RequireSigning and FileRootCert validation once signature verification is implemented for RCv2.
 	// 2. LimitScripts must be a valid AllowedScriptType value. so map/check the value to the AllowedScriptTypeFlag bitmask.
 	if rceps.LimitScripts != "" {
 		if err != nil {
@@ -107,9 +108,10 @@ func (rceps RCv2ExtensionPolicySettings) ValidateFormat() error {
 	return nil
 }
 
-// This function compares a script type (of type ScriptType, defined in this file) to the allowed script types
-// (of type AllowedScriptTypeFlag, also defined in this file) listed in the policy. These values and mappings
-// are specific to Run Command, hence why they are defined here and not in the shared library.
+// This function compares a script type (of string type ScriptType, defined in this file) to the allowed script types
+// (of type AllowedScriptTypeFlag, also defined in this file) listed in the policy.
+// Depending on the string case (the value of scriptType), it checks if the corresponding bit is enabled in the allowed script types bitmask.
+// These values and mappings are specific to Run Command, hence why they are defined here and not in the shared library.
 func CompareScriptTypeToAllowedScriptType(scriptType handlersettings.ScriptType, allowedScriptTypes AllowedScriptTypeFlag) error {
 	switch scriptType {
 	case handlersettings.InlineScript: