From fe5cf8a530cb3888451612465afa9ba0f44dc3a7 Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Thu, 9 Jun 2022 15:56:25 +0200
Subject: [PATCH] CLI operation bucket set-quotas

---
 src/garage/admin.rs       | 55 +++++++++++++++++++++++++++++++++++++++
 src/garage/cli/structs.rs | 19 ++++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/src/garage/admin.rs b/src/garage/admin.rs
index 53c8be8f..6630ae16 100644
--- a/src/garage/admin.rs
+++ b/src/garage/admin.rs
@@ -80,6 +80,7 @@ impl AdminRpcHandler {
 			BucketOperation::Allow(query) => self.handle_bucket_allow(query).await,
 			BucketOperation::Deny(query) => self.handle_bucket_deny(query).await,
 			BucketOperation::Website(query) => self.handle_bucket_website(query).await,
+			BucketOperation::SetQuotas(query) => self.handle_bucket_set_quotas(query).await,
 		}
 	}
 
@@ -470,6 +471,60 @@ impl AdminRpcHandler {
 		Ok(AdminRpc::Ok(msg))
 	}
 
+	async fn handle_bucket_set_quotas(&self, query: &SetQuotasOpt) -> Result<AdminRpc, Error> {
+		let bucket_id = self
+			.garage
+			.bucket_helper()
+			.resolve_global_bucket_name(&query.bucket)
+			.await?
+			.ok_or_bad_request("Bucket not found")?;
+
+		let mut bucket = self
+			.garage
+			.bucket_helper()
+			.get_existing_bucket(bucket_id)
+			.await?;
+		let bucket_state = bucket.state.as_option_mut().unwrap();
+
+		if query.max_size.is_none() && query.max_objects.is_none() {
+			return Err(Error::BadRequest(
+				"You must specify either --max-size or --max-objects (or both) for this command to do something.".to_string(),
+			));
+		}
+
+		let mut quotas = bucket_state.quotas.get().clone();
+
+		match query.max_size.as_ref().map(String::as_ref) {
+			Some("none") => quotas.max_size = None,
+			Some(v) => {
+				let bs = v
+					.parse::<bytesize::ByteSize>()
+					.ok_or_bad_request(format!("Invalid size specified: {}", v))?;
+				quotas.max_size = Some(bs.as_u64());
+			}
+			_ => (),
+		}
+
+		match query.max_objects.as_ref().map(String::as_ref) {
+			Some("none") => quotas.max_objects = None,
+			Some(v) => {
+				let mo = v
+					.parse::<u64>()
+					.ok_or_bad_request(format!("Invalid number specified: {}", v))?;
+				quotas.max_objects = Some(mo);
+			}
+			_ => (),
+		}
+
+		bucket_state.quotas.update(quotas);
+		self.garage.bucket_table.insert(&bucket).await?;
+
+		Ok(AdminRpc::Ok(format!(
+			"Quotas updated for {}",
+			&query.bucket
+		)))
+	}
+
 	async fn handle_key_cmd(&self, cmd: &KeyOperation) -> Result<AdminRpc, Error> {
 		match cmd {
 			KeyOperation::List => self.handle_list_keys().await,
diff --git a/src/garage/cli/structs.rs b/src/garage/cli/structs.rs
index cdaa09be..28228b07 100644
--- a/src/garage/cli/structs.rs
+++ b/src/garage/cli/structs.rs
@@ -180,6 +180,10 @@ pub enum BucketOperation {
 	/// Expose as website or not
 	#[structopt(name = "website")]
 	Website(WebsiteOpt),
+
+	/// Set the quotas for this bucket
+	#[structopt(name = "set-quotas")]
+	SetQuotas(SetQuotasOpt),
 }
 
 #[derive(Serialize, Deserialize, StructOpt, Debug)]
@@ -266,6 +270,21 @@ pub struct PermBucketOpt {
 	pub bucket: String,
 }
 
+#[derive(Serialize, Deserialize, StructOpt, Debug)]
+pub struct SetQuotasOpt {
+	/// Bucket name
+	pub bucket: String,
+
+	/// Set a maximum size for the bucket (specify a size e.g. in MiB or GiB,
+	/// or `none` for no size restriction)
+	#[structopt(long = "max-size")]
+	pub max_size: Option<String>,
+
+	/// Set a maximum number of objects for the bucket (or `none` for no restriction)
+	#[structopt(long = "max-objects")]
+	pub max_objects: Option<String>,
+}
+
 #[derive(Serialize, Deserialize, StructOpt, Debug)]
 pub enum KeyOperation {
 	/// List keys