From 98b2f2ba20523fc82287b3d91e179a6a2eb0ce16 Mon Sep 17 00:00:00 2001
From: Oluwatobi Bamidele <teebams49@gmail.com>
Date: Wed, 30 Oct 2024 22:42:36 +0100
Subject: [PATCH 1/2] feat: ensure we have a limit to the number of ec2 to be
 created in a day

---
 src/bin/super/state.rs | 22 +++++++++++++++++++++-
 src/bin/super/util.rs  | 21 +++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/src/bin/super/state.rs b/src/bin/super/state.rs
index cc8f98c7..b351955b 100644
--- a/src/bin/super/state.rs
+++ b/src/bin/super/state.rs
@@ -4,7 +4,7 @@ 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 {
@@ -12,6 +12,7 @@ pub struct Super {
     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(),
+            },
         }
     }
 
diff --git a/src/bin/super/util.rs b/src/bin/super/util.rs
index b1494f3f..2c7f9e7c 100644
--- a/src/bin/super/util.rs
+++ b/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()
+}

From ce1d2c17aac5b1f09505ae3692c51651d82314ae Mon Sep 17 00:00:00 2001
From: Oluwatobi Bamidele <teebams49@gmail.com>
Date: Wed, 30 Oct 2024 22:46:14 +0100
Subject: [PATCH 2/2] update: added EC2_DAILY_LIMIT to docker compose

---
 superadmin.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/superadmin.yml b/superadmin.yml
index 80a14154..e4ac700b 100644
--- a/superadmin.yml
+++ b/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: