Add extensions to spec

crates/catalog/src/berachain/mod.rs

@@ -1,5 +1,8 @@
 use serde::{Deserialize, Serialize};
-use spec::{Arg, Artifacts, Babel, ComputeResource, Deployment, Manifest, Pod, Spec, Volume};
+use spec::{
+    Arg, Artifacts, Babel, ComputeResource, Deployment, DeploymentExtension, Manifest, Pod, Spec,
+    Volume,
+};
 use template::Template;
 use tokio::task;
 
@@ -42,28 +45,10 @@ impl Deployment for BerachainDeployment {
     fn manifest(&self, chain: Chains, input: BerachainDeploymentInput) -> eyre::Result<Manifest> {
         let mut manifest = Manifest::new("berachain".to_string());
 
-        let mut beaconkit_pod = input.beacon_kit.spec(chain.clone())?;
-        // Add Babel sidecar to BeaconKit pod
-        let babel_cosmos = Babel::new(
-            "cosmos",
-            Arg::Ref {
-                name: "beaconkit-node".to_string(),
-                port: "http".to_string(),
-            },
-        );
-        beaconkit_pod = beaconkit_pod.with_spec("babel", babel_cosmos.spec());
+        let beaconkit_pod = input.beacon_kit.spec(chain.clone())?;
         manifest.add_spec("beaconkit".to_string(), beaconkit_pod);
 
-        let mut berareth_pod = input.bera_reth.spec(chain)?;
-        // Add Babel sidecar to BeraReth pod
-        let babel_ethereum = Babel::new(
-            "ethereum",
-            Arg::Ref {
-                name: "berareth-reth".to_string(),
-                port: "http".to_string(),
-            },
-        );
-        berareth_pod = berareth_pod.with_spec("babel", babel_ethereum.spec());
+        let berareth_pod = input.bera_reth.spec(chain)?;
         manifest.add_spec("berareth".to_string(), berareth_pod);
 
         Ok(manifest)
@@ -124,6 +109,13 @@ impl ComputeResource for BeaconKit {
             )
             .env("EL_BOOTNODES", bootnodes.trim())
             .env("EL_PEERS", peers.trim())
+            .with_babel(Babel::new(
+                "cosmos",
+                Arg::Ref {
+                    name: "beaconkit-node".to_string(),
+                    port: "http".to_string(),
+                },
+            ))
             .artifact(Artifacts::File(spec::File {
                 name: "genesis".to_string(),
                 target_path: "/data/genesis.json".to_string(),
@@ -175,6 +167,13 @@ impl ComputeResource for BeraReth {
             )
             .arg2("--http.addr", "0.0.0.0")
             .arg("--http")
+            .with_babel(Babel::new(
+                "ethereum",
+                Arg::Ref {
+                    name: "berareth-reth".to_string(),
+                    port: "http".to_string(),
+                },
+            ))
             .artifact(Artifacts::File(spec::File {
                 name: "eth-genesis".to_string(),
                 target_path: "/data/eth-genesis.json".to_string(),

crates/catalog/src/ethereum/mod.rs

@@ -1,7 +1,7 @@
 use serde::Deserialize;
 use spec::{
-    Arg, Artifacts, Babel, ComputeResource, DEFAULT_JWT_TOKEN, Deployment, Manifest, Pod, Port,
-    Spec, Volume,
+    Arg, Artifacts, Babel, ComputeResource, DEFAULT_JWT_TOKEN, Deployment, DeploymentExtension,
+    Manifest, Pod, Port, Spec, Volume,
 };
 
 #[derive(Default, Clone)]
@@ -27,35 +27,17 @@ impl Deployment for EthereumDeployment {
     fn manifest(&self, chain: Chains, input: EthDeploymentInput) -> eyre::Result<Manifest> {
         let mut manifest = Manifest::new("eth".to_string());
 
-        let mut el_node = match input.el_node {
+        let el_node = match input.el_node {
             ELNode::Reth(reth) => reth.spec(chain.clone()),
         }?;
-
-        // Add Babel sidecar to EL pod
-        let babel_el = Babel::new(
-            "ethereum",
-            Arg::Ref {
-                name: "el-node".to_string(),
-                port: "http".to_string(),
-            },
-        );
-        el_node = el_node.with_spec("babel", babel_el.spec());
         manifest.add_spec("el".to_string(), el_node);
 
-        let mut cl_node = match input.cl_node {
+        let cl_node = match input.cl_node {
             CLNode::Lighthouse(lighthouse) => lighthouse.spec(chain.clone()),
             CLNode::Prysm(prysm) => prysm.spec(chain.clone()),
         }?;
 
         // Add Babel sidecar to CL pod
-        let babel_beacon = Babel::new(
-            "ethereum_beacon",
-            Arg::Ref {
-                name: "cl-node".to_string(),
-                port: "http".to_string(),
-            },
-        );
-        cl_node = cl_node.with_spec("babel", babel_beacon.spec());
         manifest.add_spec("cl".to_string(), cl_node);
 
         Ok(manifest)
@@ -115,6 +97,13 @@ impl ComputeResource for Reth {
                 port: 9090,
                 name: "metrics".to_string(),
             })
+            .with_babel(Babel::new(
+                "ethereum",
+                Arg::Ref {
+                    name: "el-node".to_string(),
+                    port: "http".to_string(),
+                },
+            ))
             .artifact(Artifacts::File(spec::File {
                 name: "jwt".to_string(),
                 target_path: "/data/jwt_secret".to_string(),
@@ -172,6 +161,13 @@ impl ComputeResource for Lighthouse {
             .arg2("--http-address", "0.0.0.0")
             .arg("--http")
             .arg2("--datadir", "/data")
+            .with_babel(Babel::new(
+                "ethereum_beacon",
+                Arg::Ref {
+                    name: "cl-node".to_string(),
+                    port: "http".to_string(),
+                },
+            ))
             .artifact(Artifacts::File(spec::File {
                 name: "jwt".to_string(),
                 target_path: "/data/jwt_secret".to_string(),
@@ -230,6 +226,13 @@ impl ComputeResource for Prysm {
                 },
             )
             .arg("--accept-terms-of-use")
+            .with_babel(Babel::new(
+                "ethereum_beacon",
+                Arg::Ref {
+                    name: "cl-node".to_string(),
+                    port: "http".to_string(),
+                },
+            ))
             .artifact(Artifacts::File(spec::File {
                 name: "jwt".to_string(),
                 target_path: "/data/jwt_secret".to_string(),

crates/catalog/src/polygon/mod.rs

@@ -1,6 +1,9 @@
 use cosmos_keys::{generate_cometbft_key, generate_tendermint_key};
 use serde::{Deserialize, Serialize};
-use spec::{Arg, Artifacts, Babel, ComputeResource, Deployment, Manifest, Pod, Spec, Volume};
+use spec::{
+    Arg, Artifacts, Babel, ComputeResource, Deployment, DeploymentExtension, Manifest, Pod, Spec,
+    Volume,
+};
 use template::Template;
 
 #[derive(Default, Clone)]
@@ -71,6 +74,13 @@ impl ComputeResource for Heimdall {
                     preferred: 1317,
                 },
             )
+            .with_babel(Babel::new(
+                "cosmos",
+                Arg::Ref {
+                    name: "heimdall-node".to_string(),
+                    port: "http".to_string(),
+                },
+            ))
             .artifact(Artifacts::File(spec::File{
                 name: "genesis".to_string(),
                 target_path: "/data/heimdall/config/genesis.json".to_string(),
@@ -182,18 +192,8 @@ impl Deployment for PolygonDeployment {
     fn manifest(&self, chain: Chains, input: PolygonDeploymentInput) -> eyre::Result<Manifest> {
         let mut manifest = Manifest::new("polygon".to_string());
 
-        let mut heimdall_pod = input.heimdall.spec(chain.clone())?;
-        // Add Babel sidecar to Heimdall pod
-        let babel_cosmos = Babel::new(
-            "cosmos",
-            Arg::Ref {
-                name: "heimdall-node".to_string(),
-                port: "http".to_string(),
-            },
-        );
-        heimdall_pod = heimdall_pod.with_spec("babel", babel_cosmos.spec());
+        let heimdall_pod = input.heimdall.spec(chain.clone())?;
         manifest.add_spec("heimdall".to_string(), heimdall_pod);
-
         manifest.add_spec("bor".to_string(), input.bor.spec(chain)?);
 
         Ok(manifest)

crates/spec/src/lib.rs

@@ -68,7 +68,18 @@ impl Manifest {
         }
     }
 
-    pub fn add_spec(&mut self, name: String, pod: Pod) {
+    pub fn add_spec(&mut self, name: String, mut pod: Pod) {
+        let mut new_specs = Vec::new();
+        for (spec_name, spec) in &pod.specs {
+            if let Some(babel) = spec.get_babel() {
+                let babel_spec = babel.spec();
+                let babel_name = format!("{}-babel", spec_name);
+                new_specs.push((babel_name, babel_spec));
+            }
+        }
+        for (babel_name, babel_spec) in new_specs {
+            pod.specs.insert(babel_name, babel_spec);
+        }
         self.pods.insert(name, pod);
     }
 
@@ -211,6 +222,7 @@ pub struct Spec {
     pub artifacts: Vec<Artifacts>,
     pub ports: Vec<Port>,
     pub volumes: Vec<Volume>,
+    pub extensions: HashMap<String, serde_json::Value>,
 }
 
 #[derive(Default)]
@@ -224,6 +236,7 @@ pub struct SpecBuilder {
     artifacts: Vec<Artifacts>,
     ports: Vec<Port>,
     volumes: Vec<Volume>,
+    extensions: HashMap<String, serde_json::Value>,
 }
 
 impl Spec {
@@ -305,6 +318,25 @@ impl SpecBuilder {
         self
     }
 
+    pub fn extension<K, V>(mut self, key: K, value: V) -> Self
+    where
+        K: Into<String>,
+        V: Into<serde_json::Value>,
+    {
+        self.extensions.insert(key.into(), value.into());
+        self
+    }
+
+    pub fn get_extension<R: serde::de::DeserializeOwned>(
+        &self,
+        name: String,
+    ) -> eyre::Result<Option<R>> {
+        match self.extensions.get(&name) {
+            Some(value) => Ok(Some(serde_json::from_value(value.clone())?)),
+            None => Ok(None),
+        }
+    }
+
     pub fn build(self) -> Spec {
         let mut ports = self.ports;
 
@@ -327,6 +359,7 @@ impl SpecBuilder {
             artifacts: self.artifacts,
             ports: ports,
             volumes: self.volumes,
+            extensions: self.extensions,
         }
     }
 }
@@ -337,7 +370,7 @@ impl Into<Spec> for SpecBuilder {
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Babel {
     pub node_type: String,
     pub rpc_url: Arg,
@@ -366,6 +399,52 @@ impl Babel {
     }
 }
 
+pub trait DeploymentExtension {
+    fn min_version(self, version: String) -> Self;
+    fn with_babel(self, babel: Babel) -> Self;
+    fn get_babel(&self) -> Option<Babel>;
+}
+
+impl DeploymentExtension for SpecBuilder {
+    fn min_version(self, version: String) -> Self {
+        self.extension("min_version", serde_json::Value::String(version))
+    }
+
+    fn with_babel(self, babel: Babel) -> Self {
+        self.extension("babel", serde_json::to_value(babel).unwrap())
+    }
+
+    fn get_babel(&self) -> Option<Babel> {
+        self.get_extension("babel".to_string()).ok().flatten()
+    }
+}
+
+impl DeploymentExtension for Spec {
+    fn min_version(self, _version: String) -> Self {
+        self
+    }
+
+    fn with_babel(self, _babel: Babel) -> Self {
+        self
+    }
+
+    fn get_babel(&self) -> Option<Babel> {
+        self.get_extension("babel".to_string()).ok().flatten()
+    }
+}
+
+impl Spec {
+    pub fn get_extension<R: serde::de::DeserializeOwned>(
+        &self,
+        name: String,
+    ) -> eyre::Result<Option<R>> {
+        match self.extensions.get(&name) {
+            Some(value) => Ok(Some(serde_json::from_value(value.clone())?)),
+            None => Ok(None),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;