wip added new arg validation by substitution

microsoft · Feb 5, 2025 · a2f0a20 · a2f0a20
1 parent eaf7179
commit a2f0a20
Show file tree

Hide file tree

Showing 4 changed files with 163 additions and 49 deletions.
diff --git a/src/agent/src/policy.rs b/src/agent/src/policy.rs
@@ -68,17 +68,6 @@ mod tests {
         // Call the function
         let result = policy.allow_request(ep, &request).await;
 
-        let x = "--nodeid=aks-nodepool1-40948945-vmss000000";
-
-        let split = shlex::split(x).unwrap();
-        println!("shlex split {:?}", split);
-
-        let split = shell_words::split(x).unwrap();
-        println!("shell_words split {:?}", split);
-
-        let split = shellwords::split(x).unwrap();
-        println!("shellwords split {:?}", split);
-
         // Assert the expected result
         match result {
             Ok((allowed, _)) => assert!(allowed, "Expected the request to be allowed"),

diff --git a/src/tools/genpolicy/exec2.rego b/src/tools/genpolicy/exec2.rego
@@ -61,19 +61,82 @@ ALLOWED_SUBDOMAIN_NAMES := ["$(host-name)", "$(node-name)", "$(pod-uid)"]
 ALWAYS_ALLOWED = ["$(resource-field)", "$(todo-annotation)"]
 
 PolicyCreateContainerRequest:= resp {
+  # resp = {
+  #   "ops": [],
+  #   "allowed": true,
+  # }
   resp = CreateContainerRequestCommon(input.base)
   i_env_map := input.env_map
-  i_tokenized_args := input.tokenized_args
+  # i_tokenized_args := input.tokenized_args
   some p_container in policy_data.containers
   p_env_map := p_container.env_map  
   allow_env_map(p_env_map, i_env_map)
 
-  p_tokenized_args := p_container.tokenized_args
-  allow_tokenized_args(p_tokenized_args, i_tokenized_args)
+  # p_tokenized_args := p_container.tokenized_args
+  # allow_tokenized_args(p_tokenized_args, i_tokenized_args)
+
+  i_process := input.base.OCI.Process
+  p_process := p_container.OCI.Process
+
+  allow_args_by_equality(i_process, p_process, i_env_map)
 
   print("PolicyCreateContainerRequest: true")
 }
 
+allow_args_by_equality(i_process, p_process, i_env_map) {
+  not i_process.Args
+  not p_process.Args
+  print("allow_args_by_equality: no args")
+}
+
+allow_args_by_equality(i_process, p_process, i_env_map) {
+  i_args := i_process.Args
+  p_args := p_process.Args
+  print("allow_args_by_equality: i_args =", i_args, "p_args =", p_args)
+  count(i_args) == count(p_args)
+  every i, i_arg in i_args {
+    print("allow_args_by_equality: i_arg =", i_arg)
+    p_arg_replaced := replace_env_variables(p_args[i], i_env_map)
+    p_arg_replaced2 := replace(p_arg_replaced, "$$", "$")
+    print("allow_args_by_equality: p_arg_replaced2 =", p_arg_replaced2)
+    i_arg == p_arg_replaced2
+    print("allow_args_by_equality: true")
+  }
+}
+
+# this function replaces all the environment variables in a string, given a map of environment keys to environment values
+# eg str = "echo $(CLUSTER_ID); echo ${NODE_NAME};"
+# env_map = {"CLUSTER_ID": "abc", "NODE_NAME" : "xyz"}
+# result = "echo abc; echo xyz;"
+replace_env_variables(str, env_map) = result {
+  keys := [x | some x in object.keys(env_map)]
+  result := replace_str_rec(str, env_map, keys, count(keys) - 1)
+}
+
+# base case
+replace_str_rec(str, env_map, arr_keys, i) = result {
+  i < 0
+  result = str
+}
+# recursive step
+replace_str_rec(str,env_map, arr_keys, i) = result {
+  i >= 0
+  key := arr_keys[i]
+
+  # eg $(CLUSTER_ID)
+  env_key1 := concat("", ["$(", key, ")"])
+  new_str1 := replace(str, env_key1, env_map[key])
+
+  # eg ${CLUSTER_ID}
+  env_key2 := concat("", ["${", key, "}"])
+  new_str2 := replace(new_str1, env_key2, env_map[key])
+
+  # eg $CLUSTER_ID
+  env_key3 := concat("", ["$", key])
+  new_str3 := replace(new_str2, env_key3, env_map[key])
+  result = replace_str_rec(new_str3, env_map, arr_keys, i - 1)
+}
+
 allow_tokenized_args(p_tokenized_args, i_tokenized_args) {
     every i, i_tokenized_arg in i_tokenized_args {
       allow_tokenized_arg(p_tokenized_args[i], i_tokenized_arg)
@@ -724,7 +787,7 @@ allow_process_common(p_process, i_process, s_name, s_namespace) {
 allow_process(p_process, i_process, s_name, s_namespace) {
     print("allow_process: start")
 
-    allow_args(p_process, i_process, s_name)
+    # allow_args(p_process, i_process, s_name)
     allow_process_common(p_process, i_process, s_name, s_namespace)
     allow_caps(p_process.Capabilities, i_process.Capabilities)
     p_process.Terminal == i_process.Terminal
@@ -1776,7 +1839,7 @@ policy_data := {
           "Args": [
             "/bin/sh",
             "-c",
-            "while true; do echo Kubernetes; echo $(node-name); sleep 10; done"
+            "while true; do echo $(ISTIO_META_CLUSTER_ID); echo $(ISTIO_META_NODE_NAME); sleep 10; done"
           ],
           "Env": [
             "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",

diff --git a/src/tools/genpolicy/rules.rego b/src/tools/genpolicy/rules.rego
@@ -61,20 +61,82 @@ ALLOWED_SUBDOMAIN_NAMES := ["$(host-name)", "$(node-name)", "$(pod-uid)"]
 ALWAYS_ALLOWED = ["$(resource-field)", "$(todo-annotation)"]
 
 PolicyCreateContainerRequest:= resp {
+  # resp = {
+  #   "ops": [],
+  #   "allowed": true,
+  # }
   resp = CreateContainerRequestCommon(input.base)
-#   resp :=  {"ops": [], "allowed": true}
   i_env_map := input.env_map
-  i_tokenized_args := input.tokenized_args
+  # i_tokenized_args := input.tokenized_args
   some p_container in policy_data.containers
   p_env_map := p_container.env_map  
   allow_env_map(p_env_map, i_env_map)
 
-  p_tokenized_args := p_container.tokenized_args
-  allow_tokenized_args(p_tokenized_args, i_tokenized_args)
+  # p_tokenized_args := p_container.tokenized_args
+  # allow_tokenized_args(p_tokenized_args, i_tokenized_args)
+
+  i_process := input.base.OCI.Process
+  p_process := p_container.OCI.Process
+
+  allow_args_by_equality(i_process, p_process, i_env_map)
 
   print("PolicyCreateContainerRequest: true")
 }
 
+allow_args_by_equality(i_process, p_process, i_env_map) {
+  not i_process.Args
+  not p_process.Args
+  print("allow_args_by_equality: no args")
+}
+
+allow_args_by_equality(i_process, p_process, i_env_map) {
+  i_args := i_process.Args
+  p_args := p_process.Args
+  print("allow_args_by_equality: i_args =", i_args, "p_args =", p_args)
+  count(i_args) == count(p_args)
+  every i, i_arg in i_args {
+    print("allow_args_by_equality: i_arg =", i_arg)
+    p_arg_replaced := replace_env_variables(p_args[i], i_env_map)
+    p_arg_replaced2 := replace(p_arg_replaced, "$$", "$")
+    print("allow_args_by_equality: p_arg_replaced2 =", p_arg_replaced2)
+    i_arg == p_arg_replaced2
+    print("allow_args_by_equality: true")
+  }
+}
+
+# this function replaces all the environment variables in a string, given a map of environment keys to environment values
+# eg str = "echo $(CLUSTER_ID); echo ${NODE_NAME};"
+# env_map = {"CLUSTER_ID": "abc", "NODE_NAME" : "xyz"}
+# result = "echo abc; echo xyz;"
+replace_env_variables(str, env_map) = result {
+  keys := [x | some x in object.keys(env_map)]
+  result := replace_str_rec(str, env_map, keys, count(keys) - 1)
+}
+
+# base case
+replace_str_rec(str, env_map, arr_keys, i) = result {
+  i < 0
+  result = str
+}
+# recursive step
+replace_str_rec(str,env_map, arr_keys, i) = result {
+  i >= 0
+  key := arr_keys[i]
+
+  # eg $(CLUSTER_ID)
+  env_key1 := concat("", ["$(", key, ")"])
+  new_str1 := replace(str, env_key1, env_map[key])
+
+  # eg ${CLUSTER_ID}
+  env_key2 := concat("", ["${", key, "}"])
+  new_str2 := replace(new_str1, env_key2, env_map[key])
+
+  # eg $CLUSTER_ID
+  env_key3 := concat("", ["$", key])
+  new_str3 := replace(new_str2, env_key3, env_map[key])
+  result = replace_str_rec(new_str3, env_map, arr_keys, i - 1)
+}
+
 allow_tokenized_args(p_tokenized_args, i_tokenized_args) {
     every i, i_tokenized_arg in i_tokenized_args {
       allow_tokenized_arg(p_tokenized_args[i], i_tokenized_arg)
@@ -743,7 +805,7 @@ allow_process_common(p_process, i_process, s_name, s_namespace) {
 allow_process(p_process, i_process, s_name, s_namespace) {
     print("allow_process: start")
 
-    allow_args(p_process, i_process, s_name)
+    # allow_args(p_process, i_process, s_name)
     allow_process_common(p_process, i_process, s_name, s_namespace)
     allow_caps(p_process.Capabilities, i_process.Capabilities)
     p_process.Terminal == i_process.Terminal

diff --git a/src/tools/genpolicy/src/policy.rs b/src/tools/genpolicy/src/policy.rs
@@ -666,7 +666,7 @@ impl AgentPolicy {
         );
 
         substitute_env_variables(&mut process.Env);
-        substitute_args_env_variables(&mut process.Args, &process.Env);
+        // substitute_args_env_variables(&mut process.Args, &process.Env);
 
         c_settings.get_process_fields(&mut process);
         resource.get_process_fields(&mut process);
@@ -883,33 +883,33 @@ fn substitute_variable(
     None
 }
 
-fn substitute_args_env_variables(args: &mut Vec<String>, env: &Vec<String>) {
-    for arg in args {
-        substitute_arg_env_variables(arg, env);
-    }
-}
-
-fn substitute_arg_env_variables(arg: &mut String, env: &Vec<String>) {
-    loop {
-        let mut substituted = false;
-
-        if let Some((start, end)) = find_subst_target(arg) {
-            if let Some(new_value) = substitute_variable(arg, start, end, env) {
-                debug!(
-                    "substitute_arg_env_variables: replacing {} with {}",
-                    &arg[start..end],
-                    &new_value
-                );
-                *arg = new_value;
-                substituted = true;
-            }
-        }
-
-        if !substituted {
-            break;
-        }
-    }
-}
+// fn substitute_args_env_variables(args: &mut Vec<String>, env: &Vec<String>) {
+//     for arg in args {
+//         substitute_arg_env_variables(arg, env);
+//     }
+// }
+
+// fn substitute_arg_env_variables(arg: &mut String, env: &Vec<String>) {
+//     loop {
+//         let mut substituted = false;
+
+//         if let Some((start, end)) = find_subst_target(arg) {
+//             if let Some(new_value) = substitute_variable(arg, start, end, env) {
+//                 debug!(
+//                     "substitute_arg_env_variables: replacing {} with {}",
+//                     &arg[start..end],
+//                     &new_value
+//                 );
+//                 *arg = new_value;
+//                 substituted = true;
+//             }
+//         }
+
+//         if !substituted {
+//             break;
+//         }
+//     }
+// }
 
 fn get_container_annotations(
     resource: &dyn yaml::K8sResource,