From 2c94775ca8d872689e1d603bb76d7a1742efd23b Mon Sep 17 00:00:00 2001
From: Konstantin Knizhnik <knizhnik@garret.ru>
Date: Thu, 31 Mar 2022 12:53:41 +0300
Subject: [PATCH] Use big integer arithmetic to avoid multiplication overflow

---
 Cargo.lock                                     | 14 +++++++++++++-
 pageserver/Cargo.toml                          |  1 +
 pageserver/src/layered_repository/layer_map.rs | 12 +++++++++++-
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 5dbbd29ae7..41f14fe799 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1422,6 +1422,17 @@ dependencies = [
  "num-traits",
 ]
 
+[[package]]
+name = "num-bigint"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f93ab6289c7b344a8a9f60f88d80aa20032336fe78da341afc91c8a2341fc75f"
+dependencies = [
+ "autocfg",
+ "num-integer",
+ "num-traits",
+]
+
 [[package]]
 name = "num-integer"
 version = "0.1.44"
@@ -1522,6 +1533,7 @@ dependencies = [
  "lazy_static",
  "log",
  "nix",
+ "num-bigint 0.4.3",
  "num-traits",
  "once_cell",
  "postgres 0.19.1 (git+https://github.com/zenithdb/rust-postgres.git?rev=2949d98df52587d562986aad155dd4e889e408b7)",
@@ -2384,7 +2396,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "692ca13de57ce0613a363c8c2f1de925adebc81b04c923ac60c5488bb44abe4b"
 dependencies = [
  "chrono",
- "num-bigint",
+ "num-bigint 0.2.6",
  "num-traits",
 ]
 
diff --git a/pageserver/Cargo.toml b/pageserver/Cargo.toml
index 31ce4521cb..a555370e16 100644
--- a/pageserver/Cargo.toml
+++ b/pageserver/Cargo.toml
@@ -46,6 +46,7 @@ crossbeam-utils = "0.8.5"
 fail = "0.5.0"
 rstar = "0.9.2"
 num-traits = "0.2.14"
+num-bigint = "0.4.3"
 
 rust-s3 = { version = "0.28", default-features = false, features = ["no-verify-ssl", "tokio-rustls-tls"] }
 async-compression = {version = "0.3", features = ["zstd", "tokio"]}
diff --git a/pageserver/src/layered_repository/layer_map.rs b/pageserver/src/layered_repository/layer_map.rs
index 74e8d6c2df..6f9dedf855 100644
--- a/pageserver/src/layered_repository/layer_map.rs
+++ b/pageserver/src/layered_repository/layer_map.rs
@@ -16,6 +16,8 @@ use crate::layered_repository::InMemoryLayer;
 use crate::repository::Key;
 use anyhow::Result;
 use lazy_static::lazy_static;
+use num_bigint::BigInt;
+use num_traits::cast::ToPrimitive;
 use num_traits::identities::{One, Zero};
 use num_traits::{Bounded, Num, Signed};
 use rstar::{RTree, RTreeObject, AABB};
@@ -134,7 +136,15 @@ impl Sub for IntKey {
 impl Mul for IntKey {
     type Output = Self;
     fn mul(self, rhs: Self) -> Self::Output {
-        IntKey(self.0.wrapping_mul(rhs.0))
+        // Use big integer arithmetic to avoid overflow.
+        // We have to cacluate sqrt of the result to be able to store result of operation as i128.
+        // As far as multiplication is used by R-Tree to calculate area and distance, it should not be a problem.
+        IntKey(
+            (BigInt::from(self.0) * BigInt::from(rhs.0))
+                .sqrt()
+                .to_i128()
+                .unwrap(),
+        )
     }
 }