feat: InsertAt patch location + secondary patch for service_account_m…

…issing

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)]
@@ -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,63 @@ 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 = match self.events.iter().rposition(|e| e.data.ts <= insert_ts) {
+                    Some(i) => {
+                        if self.events[i].data.ts < insert_ts {
+                            self.events.insert(i + 1, AnnotatedTraceEvent::default());
+                            i + 1
+                        } else {
+                            i
+                        }
+                    },
+                    None => {
+                        self.events.push(AnnotatedTraceEvent::default());
+                        0
+                    },
+                };
+
+
+                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))
+            },
+        }
+    }
 }
 
 #[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,59 @@ 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() {
+                    let svc_account = sa.as_str().expect("expected string for service account");
+                    if !self.seen_service_accounts.contains(svc_account) {
+                        patches.push((
+                            i,
+                            vec![
+                                construct_remove_svc_account_ref_patch(obj, &ptrs),
+                                construct_add_svc_account_patch(obj, svc_account),
+                            ],
+                        ));
                     }
                 }
             }
@@ -82,6 +102,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(obj: &DynamicObject, 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: obj.namespace(),
+                ..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-core/Cargo.toml

@@ -9,7 +9,7 @@ readme.workspace = true
 edition.workspace = true
 
 [features]
-testutils = ["dep:http", "dep:httpmock", "dep:lazy_static", "dep:mockall", "dep:rstest"]
+testutils = ["dep:http", "dep:httpmock", "dep:mockall", "dep:rstest"]
 
 [dependencies]
 anyhow = { workspace = true }
@@ -19,6 +19,7 @@ bytes = { workspace = true }
 clockabilly = { workspace = true }
 kube = { workspace = true }
 k8s-openapi = { workspace = true }
+lazy_static = { workspace = true }
 object_store = { workspace = true }
 parse_datetime_fork = { workspace = true }
 paste = { workspace = true }
@@ -37,7 +38,6 @@ url = { workspace = true }
 # testutils dependencies
 http = { workspace = true, optional = true }
 httpmock = { workspace = true,  optional = true }
-lazy_static = { workspace = true, optional = true }
 mockall = { workspace = true, optional = true }
 rstest = { workspace = true, optional = true }
 

sk-core/src/constants.rs

@@ -1,3 +1,7 @@
+use lazy_static::lazy_static;
+
+use crate::k8s::GVK;
+
 // Well-known labels, annotations, and taints
 pub const KUBERNETES_IO_METADATA_NAME_KEY: &str = "kubernetes.io/metadata.name";
 pub const APP_KUBERNETES_IO_NAME_KEY: &str = "app.kubernetes.io/name";
@@ -28,11 +32,17 @@ pub const SK_LEASE_NAME: &str = "sk-lease";
 pub const RETRY_DELAY_SECONDS: u64 = 5;
 pub const ERROR_RETRY_DELAY_SECONDS: u64 = 30;
 
+// Kinds
+pub const SVC_ACCOUNT_KIND: &str = "ServiceAccount";
+
+// Built-in GVKs
+lazy_static! {
+    pub static ref SVC_ACCOUNT_GVK: GVK = GVK::new("", "v1", SVC_ACCOUNT_KIND);
+}
+
 #[cfg(feature = "testutils")]
 mod test_constants {
-    use lazy_static::lazy_static;
-
-    use crate::k8s::GVK;
+    use super::*;
 
     pub const EMPTY_POD_SPEC_HASH: u64 = 17506812802394981455;
     pub const TEST_DEPLOYMENT: &str = "the-deployment";
@@ -49,7 +59,6 @@ mod test_constants {
     lazy_static! {
         pub static ref DEPL_GVK: GVK = GVK::new("apps", "v1", "Deployment");
         pub static ref DS_GVK: GVK = GVK::new("apps", "v1", "DaemonSet");
-        pub static ref SA_GVK: GVK = GVK::new("", "v1", "ServiceAccount");
     }
 }
 

sk-core/src/k8s/testutils/objs.rs

@@ -24,5 +24,5 @@ pub fn test_daemonset(#[default(TEST_DAEMONSET)] name: &str) -> DynamicObject {
 
 #[fixture]
 pub fn test_service_account(#[default(TEST_SERVICE_ACCOUNT)] name: &str) -> DynamicObject {
-    DynamicObject::new(&name, &ApiResource::from_gvk(&SA_GVK)).within(TEST_NAMESPACE)
+    DynamicObject::new(&name, &ApiResource::from_gvk(&SVC_ACCOUNT_GVK)).within(TEST_NAMESPACE)
 }

sk-store/src/lib.rs

@@ -29,8 +29,8 @@ pub use crate::index::TraceIndex;
 use crate::pod_owners_map::PodLifecyclesMap;
 pub use crate::store::TraceStore;
 
-#[derive(Debug)]
-enum TraceAction {
+#[derive(Clone, Debug)]
+pub enum TraceAction {
     ObjectApplied,
     ObjectDeleted,
 }