From 7cb84d45c126f5009db5556c82987e9c2ad83451 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Tue, 28 Nov 2023 06:03:00 -0800
Subject: [PATCH] Prevent integer overflow in history buffer for large models.

PiperOrigin-RevId: 585949038
Change-Id: Iabec5feedd1632b635583a608a5488b39e6bb9a2
---
 simulate/simulate.cc | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/simulate/simulate.cc b/simulate/simulate.cc
index f1d27cf3..0b2a150a 100644
--- a/simulate/simulate.cc
+++ b/simulate/simulate.cc
@@ -17,6 +17,7 @@
 #include <algorithm>
 #include <atomic>
 #include <chrono>
+#include <climits>
 #include <cstdio>
 #include <cstring>
 #include <memory>
@@ -2228,13 +2229,13 @@ void Simulate::LoadOnRenderThread() {
 
   // allocate history buffer: smaller of {2000 states, 100 MB}
   if (!this->is_passive_) {
-    constexpr int kHistoryLength = 2000;
     constexpr int kMaxHistoryBytes = 1e8;
 
     // get state size, size of history buffer
     state_size_ = mj_stateSize(this->m_, mjSTATE_INTEGRATION);
     int state_bytes = state_size_ * sizeof(mjtNum);
-    int history_bytes = mjMIN(state_bytes * kHistoryLength, kMaxHistoryBytes);
+    int history_length = mjMIN(INT_MAX / state_bytes, 2000);
+    int history_bytes = mjMIN(state_bytes * history_length, kMaxHistoryBytes);
     nhistory_ = history_bytes / state_bytes;
 
     // allocate history buffer, reset cursor and UI slider