From ac3321c8570683a887f91c99206ddc2d3b5dde12 Mon Sep 17 00:00:00 2001
From: Ahson Khan <ahson_ahmedk@yahoo.com>
Date: Wed, 6 Nov 2024 18:32:09 -0800
Subject: [PATCH] Fix overflow issue in token cache. (#6190)

* Fix overflow issue in token cache.

* Add test
---
 sdk/identity/azure-identity/CHANGELOG.md        |  2 ++
 sdk/identity/azure-identity/src/token_cache.cpp |  4 +++-
 .../azure-identity/test/ut/token_cache_test.cpp | 17 +++++++++++++++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/sdk/identity/azure-identity/CHANGELOG.md b/sdk/identity/azure-identity/CHANGELOG.md
index 871f21e599..0a252a0274 100644
--- a/sdk/identity/azure-identity/CHANGELOG.md
+++ b/sdk/identity/azure-identity/CHANGELOG.md
@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Fix overflow issue in token cache.
+
 ### Other Changes
 
 - [[#6086]](https://github.com/Azure/azure-sdk-for-cpp/pull/6086) Correct minimum version specification for the Azure Core dependency. (A community contribution, courtesy of _[jdblischak](https://github.com/jdblischak)_)
diff --git a/sdk/identity/azure-identity/src/token_cache.cpp b/sdk/identity/azure-identity/src/token_cache.cpp
index 5c470d6584..20457f97e7 100644
--- a/sdk/identity/azure-identity/src/token_cache.cpp
+++ b/sdk/identity/azure-identity/src/token_cache.cpp
@@ -18,7 +18,9 @@ bool TokenCache::IsFresh(
     DateTime::duration minimumExpiration,
     std::chrono::system_clock::time_point now)
 {
-  return item->AccessToken.ExpiresOn > (DateTime(now) + minimumExpiration);
+  // Even if ExpiresOn is unset, the default and lowest value of DateTime is time_point of 0
+  // Therefore, there is no risk of underflow with subtracting duration::max().
+  return (item->AccessToken.ExpiresOn - minimumExpiration) > DateTime(now);
 }
 
 namespace {
diff --git a/sdk/identity/azure-identity/test/ut/token_cache_test.cpp b/sdk/identity/azure-identity/test/ut/token_cache_test.cpp
index 550dd7150f..bdf1ba058d 100644
--- a/sdk/identity/azure-identity/test/ut/token_cache_test.cpp
+++ b/sdk/identity/azure-identity/test/ut/token_cache_test.cpp
@@ -42,6 +42,23 @@ class TestableTokenCache final : public TokenCache {
 
 using namespace std::chrono_literals;
 
+TEST(TokenCache, IsFreshOverflow)
+{
+  TestableTokenCache tokenCache;
+
+  EXPECT_EQ(tokenCache.m_cache.size(), 0UL);
+
+  DateTime const Tomorrow = std::chrono::system_clock::now() + 24h;
+
+  auto const token1 = tokenCache.GetToken("A", {}, DateTime::duration::max(), [=]() {
+    AccessToken result;
+    result.Token = "T1";
+    result.ExpiresOn = Tomorrow;
+    return result;
+  });
+  EXPECT_EQ(token1.Token, "T1");
+}
+
 TEST(TokenCache, GetReuseRefresh)
 {
   TestableTokenCache tokenCache;