From 5e72d368006247e664f5c9c2fe52d21aa91c2f88 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Sun, 7 Jul 2024 03:36:50 +0000
Subject: [PATCH] add math utils; integrate checked expression macros

Signed-off-by: Jason Volk <jason@zemos.net>
---
 Cargo.lock              | 10 ++++++++++
 Cargo.toml              |  3 +++
 src/core/Cargo.toml     |  1 +
 src/core/error.rs       |  8 ++++++++
 src/core/utils/math.rs  | 26 ++++++++++++++++++++++++++
 src/core/utils/mod.rs   |  1 +
 src/core/utils/tests.rs | 19 +++++++++++++++++++
 7 files changed, 68 insertions(+)
 create mode 100644 src/core/utils/math.rs

diff --git a/Cargo.lock b/Cargo.lock
index 3415cac0..7152c66d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -504,6 +504,15 @@ version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
 
+[[package]]
+name = "checked_ops"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b491d76efc1d99d74de3c8529bee64c62312c275c7eb124f9185291de45801d5"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "chrono"
 version = "0.4.38"
@@ -658,6 +667,7 @@ dependencies = [
  "argon2",
  "axum 0.7.5",
  "bytes",
+ "checked_ops",
  "chrono",
  "either",
  "figment",
diff --git a/Cargo.toml b/Cargo.toml
index f6b23102..c93c6736 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -411,6 +411,9 @@ rev = "de26100b0db03e419a3d8e1dd26895d170d1fe50"
 version = "0.29.4"
 default-features = false
 
+[workspace.dependencies.checked_ops]
+version = "0.1"
+
 
 #
 # Patches
diff --git a/src/core/Cargo.toml b/src/core/Cargo.toml
index 2c0fc47d..e47f673e 100644
--- a/src/core/Cargo.toml
+++ b/src/core/Cargo.toml
@@ -53,6 +53,7 @@ sha256_media = []
 argon2.workspace = true
 axum.workspace = true
 bytes.workspace = true
+checked_ops.workspace = true
 chrono.workspace = true
 either.workspace = true
 figment.workspace = true
diff --git a/src/core/error.rs b/src/core/error.rs
index ea7acda5..63729f31 100644
--- a/src/core/error.rs
+++ b/src/core/error.rs
@@ -35,6 +35,12 @@ pub enum Error {
 	FromUtf8Error(#[from] std::string::FromUtf8Error),
 	#[error("{0}")]
 	TryFromSliceError(#[from] std::array::TryFromSliceError),
+	#[error("{0}")]
+	TryFromIntError(#[from] std::num::TryFromIntError),
+	#[error("{0}")]
+	ParseIntError(#[from] std::num::ParseIntError),
+	#[error("{0}")]
+	ParseFloatError(#[from] std::num::ParseFloatError),
 
 	// third-party
 	#[error("Regex error: {0}")]
@@ -63,6 +69,8 @@ pub enum Error {
 	InconsistentRoomState(&'static str, ruma::OwnedRoomId),
 
 	// conduwuit
+	#[error("Arithmetic operation failed: {0}")]
+	Arithmetic(&'static str),
 	#[error("There was a problem with your configuration: {0}")]
 	BadConfig(String),
 	#[error("{0}")]
diff --git a/src/core/utils/math.rs b/src/core/utils/math.rs
new file mode 100644
index 00000000..d5dbf3a6
--- /dev/null
+++ b/src/core/utils/math.rs
@@ -0,0 +1,26 @@
+use std::{cmp, time::Duration};
+
+pub use checked_ops::checked_ops;
+
+/// Checked arithmetic expression. Returns a Result<R, Error::Arithmetic>
+#[macro_export]
+macro_rules! checked {
+	($($input:tt)*) => {
+		$crate::utils::math::checked_ops!($($input)*)
+			.ok_or_else(|| $crate::Error::Arithmetic("operation overflowed or result invalid"))
+	}
+}
+
+/// in release-mode. Use for performance when the expression is obviously safe.
+/// The check remains in debug-mode for regression analysis.
+#[cfg(not(debug_assertions))]
+#[macro_export]
+macro_rules! validated {
+	($($input:tt)*) => { Ok($($input)*) }
+}
+
+#[cfg(debug_assertions)]
+#[macro_export]
+macro_rules! validated {
+	($($input:tt)*) => { $crate::checked!($($input)*) }
+}
diff --git a/src/core/utils/mod.rs b/src/core/utils/mod.rs
index f9f1b87e..2b79c3c4 100644
--- a/src/core/utils/mod.rs
+++ b/src/core/utils/mod.rs
@@ -5,6 +5,7 @@ pub mod defer;
 pub mod hash;
 pub mod html;
 pub mod json;
+pub mod math;
 pub mod mutex_map;
 pub mod rand;
 pub mod string;
diff --git a/src/core/utils/tests.rs b/src/core/utils/tests.rs
index f5cd0a07..239e27e9 100644
--- a/src/core/utils/tests.rs
+++ b/src/core/utils/tests.rs
@@ -62,3 +62,22 @@ fn common_prefix_none() {
 	let output = string::common_prefix(&input);
 	assert_eq!(output, "");
 }
+
+#[test]
+fn checked_add() {
+	use utils::math::checked;
+
+	let a = 1234;
+	let res = checked!(a + 1).unwrap();
+	assert_eq!(res, 1235);
+}
+
+#[test]
+#[should_panic(expected = "overflow")]
+fn checked_add_overflow() {
+	use utils::math::checked;
+
+	let a: u64 = u64::MAX;
+	let res = checked!(a + 1).expect("overflow");
+	assert_eq!(res, 0);
+}