feat: InsertAt patch location, secondary patch for service_account_mi…

…ssing

Cargo.toml

@@ -67,7 +67,7 @@ tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 url = "2.5.3"
 
 # test dependencies
-assertables = "8.18.0"
+assertables = "9.5.0"
 http = "1.1.0"
 httpmock = "0.6.8"
 hyper = "1.5.0"

sk-cli/src/validation/annotated_trace.rs

@@ -1,13 +1,16 @@
 use std::collections::BTreeMap; // BTreeMap sorts by key, HashMap doesn't
+use std::iter::once;
 use std::slice;
 
 use json_patch_ext::prelude::*;
+use serde_json::json;
 use sk_core::external_storage::{
     ObjectStoreWrapper,
     SkObjectStore,
 };
 use sk_core::prelude::*;
 use sk_store::{
+    TraceAction,
     TraceEvent,
     TraceStorable,
     TraceStore,
@@ -18,11 +21,12 @@ use super::validator::{
     ValidatorCode,
 };
 
+
 #[derive(Clone, Debug)]
 pub enum PatchLocations {
     Everywhere,
-    #[allow(dead_code)]
     ObjectReference(TypeMeta, String),
+    InsertAt(i64, TraceAction, TypeMeta, metav1::ObjectMeta),
 }
 
 #[derive(Clone, Debug)]
@@ -70,7 +74,7 @@ pub struct AnnotatedTrace {
     base: TraceStore,
     patches: Vec<AnnotatedTracePatch>,
 
-    events: Vec<AnnotatedTraceEvent>,
+    pub(super) events: Vec<AnnotatedTraceEvent>,
 }
 
 impl AnnotatedTrace {
@@ -89,21 +93,10 @@ impl AnnotatedTrace {
 
     pub fn apply_patch(&mut self, patch: AnnotatedTracePatch) -> anyhow::Result<usize> {
         let mut count = 0;
-        for event in self.events.iter_mut() {
-            for obj in event.data.applied_objs.iter_mut().chain(event.data.deleted_objs.iter_mut()) {
-                let should_apply_here = match patch.locations {
-                    PatchLocations::Everywhere => true,
-                    PatchLocations::ObjectReference(ref type_, ref ns_name) => {
-                        obj.types.as_ref().is_some_and(|t| t == type_) && &obj.namespaced_name() == ns_name
-                    },
-                };
-
-                if should_apply_here {
-                    count += 1;
-                    for op in &patch.ops {
-                        patch_ext(&mut obj.data, op.clone())?;
-                    }
-                }
+        for obj in self.matched_objects(&patch.locations) {
+            count += 1;
+            for op in &patch.ops {
+                patch_ext(&mut obj.data, op.clone())?;
             }
         }
         self.patches.push(patch);
@@ -186,6 +179,74 @@ impl AnnotatedTrace {
     pub fn start_ts(&self) -> Option<i64> {
         self.events.first().map(|evt| evt.data.ts)
     }
+
+    fn object_iter_mut(&mut self) -> impl Iterator<Item = &mut DynamicObject> {
+        self.events
+            .iter_mut()
+            .flat_map(|e| e.data.applied_objs.iter_mut().chain(e.data.deleted_objs.iter_mut()))
+    }
+}
+
+impl<'a> AnnotatedTrace {
+    fn matched_objects(
+        &'a mut self,
+        locations: &'a PatchLocations,
+    ) -> Box<dyn Iterator<Item = &mut DynamicObject> + 'a> {
+        match locations {
+            PatchLocations::Everywhere => Box::new(self.object_iter_mut()),
+            PatchLocations::ObjectReference(ref type_, ref ns_name) => {
+                Box::new(self.object_iter_mut().filter(move |obj| {
+                    obj.types.as_ref().is_some_and(|t| t == type_) && &obj.namespaced_name() == ns_name
+                }))
+            },
+            PatchLocations::InsertAt(relative_ts, action, type_meta, object_meta) => {
+                let insert_ts = self.start_ts().unwrap_or_default() + relative_ts;
+                let insert_idx = find_or_create_event_at_ts(&mut self.events, insert_ts);
+
+                let new_obj = DynamicObject {
+                    types: Some(type_meta.clone()),
+                    metadata: object_meta.clone(),
+                    data: json!({}),
+                };
+                let obj = match action {
+                    TraceAction::ObjectApplied => {
+                        self.events[insert_idx].data.applied_objs.push(new_obj);
+                        self.events[insert_idx].data.applied_objs.iter_mut().last().unwrap()
+                    },
+                    TraceAction::ObjectDeleted => {
+                        self.events[insert_idx].data.deleted_objs.push(new_obj);
+                        self.events[insert_idx].data.deleted_objs.iter_mut().last().unwrap()
+                    },
+                };
+                Box::new(once(obj))
+            },
+        }
+    }
+}
+
+pub(super) fn find_or_create_event_at_ts(events: &mut Vec<AnnotatedTraceEvent>, ts: i64) -> usize {
+    let new_event = AnnotatedTraceEvent {
+        data: TraceEvent { ts, ..Default::default() },
+        ..Default::default()
+    };
+    // Walk through the events list backwards until we find the first one less than the given ts
+    match events.iter().rposition(|e| e.data.ts <= ts) {
+        Some(i) => {
+            // If we found one, and the ts isn't equal, create an event with the specified
+            // timestamp; this goes at index i+1 since it needs to go after the lower (found) ts
+            if events[i].data.ts < ts {
+                events.insert(i + 1, new_event);
+                i + 1
+            } else {
+                i // otherwise the timestamp is equal so return this index
+            }
+        },
+        None => {
+            // In this case there are no events in the trace, so we add one at the beginning
+            events.push(new_event);
+            0
+        },
+    }
 }
 
 #[cfg(test)]

sk-cli/src/validation/rules/service_account_missing.rs

@@ -10,7 +10,10 @@ use std::sync::{
 use json_patch_ext::prelude::*;
 use sk_core::k8s::GVK;
 use sk_core::prelude::*;
-use sk_store::TracerConfig;
+use sk_store::{
+    TraceAction,
+    TracerConfig,
+};
 
 use crate::validation::validator::{
     CheckResult,
@@ -33,42 +36,63 @@ pub struct ServiceAccountMissing {
     pub(crate) seen_service_accounts: HashSet<String>,
 }
 
-impl Diagnostic for ServiceAccountMissing {
-    fn check_next_event(&mut self, event: &mut AnnotatedTraceEvent, config: &TracerConfig) -> CheckResult {
+impl ServiceAccountMissing {
+    fn record_service_accounts(&mut self, event: &mut AnnotatedTraceEvent) {
         for obj in &event.data.applied_objs {
             if let Some(ref type_meta) = obj.types {
-                if &type_meta.kind == "ServiceAccount" {
+                if type_meta.kind == SVC_ACCOUNT_KIND {
                     self.seen_service_accounts.insert(obj.namespaced_name());
                 }
             }
         }
         for obj in &event.data.deleted_objs {
             if let Some(ref type_meta) = obj.types {
-                if &type_meta.kind == "ServiceAccount" {
+                if type_meta.kind == SVC_ACCOUNT_KIND {
                     self.seen_service_accounts.remove(&obj.namespaced_name());
                 }
             }
         }
+    }
+}
+
+impl Diagnostic for ServiceAccountMissing {
+    fn check_next_event(&mut self, event: &mut AnnotatedTraceEvent, config: &TracerConfig) -> CheckResult {
+        // First we check all the objects in this event and record any service accounts we see
+        // (and remove any service accounts that got deleted); this way if the service account
+        // and the pod referencing it are created at the same time we don't fail (maybe we should,
+        // though?  not sure, anyways it's fine for now).
+        self.record_service_accounts(event);
 
         let mut patches = vec![];
         for (i, obj) in event.data.applied_objs.iter().enumerate() {
             let gvk = GVK::from_dynamic_obj(obj)?;
             if let Some(pod_spec_template_path) = config.pod_spec_template_path(&gvk) {
-                let sa_ptrs = [
+                let ptrs = [
                     // serviceAccount is deprecated but still supported (for now)
                     format_ptr!("{pod_spec_template_path}/spec/serviceAccount"),
                     format_ptr!("{pod_spec_template_path}/spec/serviceAccountName"),
                 ];
-                if let Some(sa) = sa_ptrs.iter().filter_map(|ptr| ptr.resolve(&obj.data).ok()).next() {
-                    if !self.seen_service_accounts.contains(sa.as_str().expect("expected string")) {
-                        let fix = AnnotatedTracePatch {
-                            locations: PatchLocations::ObjectReference(
-                                obj.types.clone().unwrap_or_default(),
-                                obj.namespaced_name(),
-                            ),
-                            ops: sa_ptrs.iter().map(|ptr| remove_operation(ptr.clone())).collect(),
-                        };
-                        patches.push((i, vec![fix]));
+
+                // If this obj references a service account that doesn't exist at this point in
+                // time, there are two possible fixes:
+                //
+                // 1) remove the reference to the service account from the pod template spec (recommended, because
+                //    the pod won't exist and can't actually _do_ anything anyways), or
+                // 2) add the service account object in at the beginning of the simulation
+                if let Some(sa) = ptrs.iter().filter_map(|ptr| ptr.resolve(&obj.data).ok()).next() {
+                    // if we're demanding a service account, we must have a namespace and a name,
+                    // these unwraps should be safe
+                    let svc_account = sa.as_str().unwrap();
+                    let svc_account_ns = &obj.namespace().unwrap();
+
+                    if !self.seen_service_accounts.contains(&format!("{svc_account_ns}/{svc_account}")) {
+                        patches.push((
+                            i,
+                            vec![
+                                construct_remove_svc_account_ref_patch(obj, &ptrs),
+                                construct_add_svc_account_patch(svc_account_ns, svc_account),
+                            ],
+                        ));
                     }
                 }
             }
@@ -82,6 +106,29 @@ impl Diagnostic for ServiceAccountMissing {
     }
 }
 
+fn construct_remove_svc_account_ref_patch(obj: &DynamicObject, ptrs: &[PointerBuf]) -> AnnotatedTracePatch {
+    AnnotatedTracePatch {
+        locations: PatchLocations::ObjectReference(obj.types.clone().unwrap_or_default(), obj.namespaced_name()),
+        ops: ptrs.iter().map(|ptr| remove_operation(ptr.clone())).collect(),
+    }
+}
+
+fn construct_add_svc_account_patch(svc_account_ns: &str, svc_account: &str) -> AnnotatedTracePatch {
+    AnnotatedTracePatch {
+        locations: PatchLocations::InsertAt(
+            0,
+            TraceAction::ObjectApplied,
+            SVC_ACCOUNT_GVK.into_type_meta(),
+            metav1::ObjectMeta {
+                name: Some(svc_account.into()),
+                namespace: Some(svc_account_ns.into()),
+                ..Default::default()
+            },
+        ),
+        ops: vec![],
+    }
+}
+
 pub fn validator() -> Validator {
     Validator {
         type_: ValidatorType::Error,

sk-cli/src/validation/rules/tests/mod.rs

@@ -25,7 +25,7 @@ fn test_trace_config() -> TracerConfig {
                     ..Default::default()
                 },
             ),
-            (SA_GVK.clone(), Default::default()),
+            (SVC_ACCOUNT_GVK.clone(), Default::default()),
         ]),
     }
 }

sk-cli/src/validation/rules/tests/service_account_missing_test.rs

@@ -12,8 +12,10 @@ use super::*;
 fn depl_event(test_deployment: DynamicObject, #[default("serviceAccount")] sa_key: &str) -> AnnotatedTraceEvent {
     AnnotatedTraceEvent {
         data: TraceEvent {
-            ts: 1,
-            applied_objs: vec![test_deployment.data(json!({"spec": {"template": {"spec": {sa_key: "foobar"}}}}))],
+            ts: 2,
+            applied_objs: vec![
+                test_deployment.data(json!({"spec": {"template": {"spec": {sa_key: TEST_SERVICE_ACCOUNT}}}}))
+            ],
             deleted_objs: vec![],
         },
         ..Default::default()
@@ -40,7 +42,7 @@ fn test_service_account_missing(test_deployment: DynamicObject, test_trace_confi
     let mut evt = depl_event(test_deployment, sa_key);
     let annotations = v.check_next_event(&mut evt, &test_trace_config).unwrap();
 
-    assert_eq!(annotations.keys().collect::<Vec<_>>(), vec![&0]);
+    assert_eq!(annotations.get(&0).unwrap().len(), 2);
 }
 
 #[rstest]
@@ -53,7 +55,7 @@ fn test_service_account_missing_deleted(
     let mut v = ServiceAccountMissing::default();
     let mut sa_event_del = AnnotatedTraceEvent {
         data: TraceEvent {
-            ts: 0,
+            ts: 1,
             applied_objs: vec![],
             deleted_objs: vec![test_service_account],
         },
@@ -63,7 +65,7 @@ fn test_service_account_missing_deleted(
     v.check_next_event(&mut sa_event_del, &test_trace_config).unwrap();
     let annotations = v.check_next_event(&mut depl_event, &test_trace_config).unwrap();
 
-    assert_eq!(annotations.keys().collect::<Vec<_>>(), vec![&0]);
+    assert_eq!(annotations.get(&0).unwrap().len(), 2);
 }
 
 #[rstest]
@@ -76,7 +78,7 @@ fn test_service_account_not_missing(
     v.check_next_event(&mut sa_event, &test_trace_config).unwrap();
     let annotations = v.check_next_event(&mut depl_event, &test_trace_config).unwrap();
 
-    assert_eq!(annotations.keys().collect::<Vec<_>>(), vec![&0]);
+    assert_none!(annotations.get(&0));
 }
 
 #[rstest]
@@ -100,7 +102,7 @@ fn test_service_account_not_missing_same_evt(
     };
     let annotations = v.check_next_event(&mut depl_evt, &test_trace_config).unwrap();
 
-    assert_eq!(annotations.keys().collect::<Vec<_>>(), vec![&0]);
+    assert_none!(annotations.get(&0));
 }
 
 #[rstest]