Merge pull request #374 from stakwork/ec2-limit

feat: ensure we have a limit to the number of ec2 to be created in a day

src/bin/super/state.rs

@@ -4,14 +4,15 @@ use serde::{Deserialize, Serialize};
 use sphinx_swarm::config::{Role, User};
 use sphinx_swarm::secrets;
 
-use crate::util::get_descriptive_instance_type;
+use crate::util::{get_descriptive_instance_type, get_today_dash_date};
 
 #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
 pub struct Super {
     pub stacks: Vec<RemoteStack>,
     pub users: Vec<User>,
     pub jwt_key: String,
     pub bots: Vec<BotCred>,
+    pub ec2_limit: Ec2Limit,
 }
 
 #[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Default)]
@@ -25,6 +26,12 @@ pub struct RemoteStack {
     pub ec2_instance_id: String,
 }
 
+#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Default)]
+pub struct Ec2Limit {
+    pub count: i32,
+    pub date: String,
+}
+
 #[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Default, Clone)]
 pub struct AwsInstanceType {
     pub name: String,
@@ -52,6 +59,7 @@ impl Default for Super {
             users: vec![default_superuser()],
             jwt_key: secrets::random_word(16),
             bots: Vec::new(),
+            ec2_limit: default_ec2_limit(),
         }
     }
 }
@@ -77,6 +85,14 @@ fn default_superuser() -> User {
     }
 }
 
+fn default_ec2_limit() -> Ec2Limit {
+    let today_dash_date = get_today_dash_date();
+    Ec2Limit {
+        count: 0,
+        date: today_dash_date,
+    }
+}
+
 impl Super {
     pub fn remove_tokens(&self) -> Super {
         let stacks = self
@@ -107,6 +123,10 @@ impl Super {
             users: vec![],
             jwt_key: "".to_string(),
             bots: bots,
+            ec2_limit: Ec2Limit {
+                count: 0,
+                date: "".to_string(),
+            },
         }
     }
 

src/bin/super/util.rs

@@ -11,6 +11,7 @@ use aws_sdk_ec2::types::{
 };
 use aws_sdk_ec2::Client;
 use aws_smithy_types::retry::RetryConfig;
+use chrono::Local;
 use futures_util::TryFutureExt;
 use reqwest::Response;
 use serde_json::Value;
@@ -605,6 +606,22 @@ pub async fn create_swarm_ec2(
     info: &CreateEc2InstanceInfo,
     state: &mut Super,
 ) -> Result<(), Error> {
+    let daily_limit = getenv("EC2_DAILY_LIMIT")
+        .unwrap_or("5".to_string())
+        .parse()
+        .unwrap_or(5);
+
+    let today_date = get_today_dash_date();
+    if today_date == state.ec2_limit.date {
+        if &state.ec2_limit.count < &daily_limit {
+            state.ec2_limit.count = state.ec2_limit.count + 1;
+        } else {
+            return Err(anyhow!("Daily limit for creating Ec2 Instance exceeded"));
+        }
+    } else {
+        state.ec2_limit.date = today_date;
+        state.ec2_limit.count = 1;
+    }
     let mut actual_vanity_address: Option<String> = None;
 
     let instance_type = get_instance(&info.instance_type);
@@ -902,3 +919,7 @@ pub async fn get_config(state: &mut Super) -> Result<Super, Error> {
     let res = state.remove_tokens();
     Ok(res)
 }
+
+pub fn get_today_dash_date() -> String {
+    Local::now().format("%d-%m-%Y").to_string()
+}

superadmin.yml

@@ -93,6 +93,7 @@ services:
       - SWARM_TAG_VALUE=$SWARM_TAG_VALUE
       - SWARM_TAG_KEY=$SWARM_TAG_KEY
       - SWARM_UPDATER_PASSWORD=$SWARM_UPDATER_PASSWORD
+      - EC2_DAILY_LIMIT=$EC2_DAILY_LIMIT
 
 networks:
   sphinx-swarm: