From 3bbaf55e3a07f656beb47fc2404ea0dee7442fea Mon Sep 17 00:00:00 2001
From: Android HW Trust Team <hwtrust@google.com>
Date: Wed, 29 Jul 2026 17:14:34 -0700
Subject: [PATCH] Add DataItem.asLong() extension function and refactor asInt
 function.

PiperOrigin-RevId: 956169666
---
 src/main/kotlin/Extension.kt     | 31 ++++++++++++++++++++-----
 src/test/kotlin/ExtensionTest.kt | 40 ++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/src/main/kotlin/Extension.kt b/src/main/kotlin/Extension.kt
index bb1e4a5..a0036dd 100644
--- a/src/main/kotlin/Extension.kt
+++ b/src/main/kotlin/Extension.kt
@@ -38,6 +38,7 @@ import java.math.BigInteger
 import java.nio.ByteBuffer
 import java.nio.charset.CodingErrorAction
 import java.security.cert.X509Certificate
+import java.time.Instant
 import java.time.YearMonth
 import java.time.format.DateTimeFormatter
 import java.time.format.DateTimeParseException
@@ -913,16 +914,34 @@ fun cborEncode(dataItem: DataItem): ByteArray {
   return baos.toByteArray()
 }
 
-fun DataItem.asInteger(): Int {
-  if (this.majorType == MajorType.UNSIGNED_INTEGER) {
-    return (this as UnsignedInteger).value.toInt()
+private fun BigInteger.toIntExact(): Int {
+  if (bitLength() < Integer.SIZE) {
+    return toInt()
   }
-  if (this.majorType == MajorType.NEGATIVE_INTEGER) {
-    return (this as NegativeInteger).value.toInt()
+  throw CborException("BigInteger out of int range")
+}
+
+private fun BigInteger.toLongExact(): Long {
+  if (bitLength() < java.lang.Long.SIZE) {
+    return toLong()
   }
-  throw CborException("Expected a number, got ${this.majorType}")
+  throw CborException("BigInteger out of long range")
 }
 
+fun DataItem.asInteger(): Int =
+  when (this) {
+    is UnsignedInteger -> value.toIntExact()
+    is NegativeInteger -> value.toIntExact()
+    else -> throw CborException("Expected a number, got $majorType")
+  }
+
+fun DataItem.asLong(): Long =
+  when (this) {
+    is UnsignedInteger -> value.toLongExact()
+    is NegativeInteger -> value.toLongExact()
+    else -> throw CborException("Expected a number, got $majorType")
+  }
+
 fun Int.asDataItem() =
   when {
     this >= 0 -> UnsignedInteger(this.toLong())
diff --git a/src/test/kotlin/ExtensionTest.kt b/src/test/kotlin/ExtensionTest.kt
index 082d948..59648af 100644
--- a/src/test/kotlin/ExtensionTest.kt
+++ b/src/test/kotlin/ExtensionTest.kt
@@ -15,6 +15,9 @@
  */
 package com.android.keyattestation.verifier
 
+import co.nstant.`in`.cbor.CborException
+import co.nstant.`in`.cbor.model.NegativeInteger
+import co.nstant.`in`.cbor.model.UnsignedInteger
 import com.android.keyattestation.verifier.testing.Chains
 import com.android.keyattestation.verifier.testing.FakeLogHook
 import com.android.keyattestation.verifier.testing.TestUtils.TESTDATA_PATH
@@ -31,6 +34,8 @@ import com.google.testing.junit.testparameterinjector.TestParameters
 import com.google.testing.junit.testparameterinjector.TestParameters.TestParametersValues
 import com.google.testing.junit.testparameterinjector.TestParametersValuesProvider
 import com.google.testing.junit.testparameterinjector.TestParametersValuesProvider.Context
+import java.math.BigInteger
+import java.time.Instant
 import java.time.YearMonth
 import kotlin.io.path.Path
 import kotlin.io.path.inputStream
@@ -255,4 +260,39 @@ class ExtensionTest {
       AttestationApplicationId.from(seq, inputLimits = limits)
     }
   }
+
+  @Test
+  fun asInteger_validValues_success() {
+    assertThat(UnsignedInteger(0).asInteger()).isEqualTo(0)
+    assertThat(UnsignedInteger(Int.MAX_VALUE.toLong()).asInteger()).isEqualTo(Int.MAX_VALUE)
+    assertThat(NegativeInteger(-1).asInteger()).isEqualTo(-1)
+    assertThat(NegativeInteger(Int.MIN_VALUE.toLong()).asInteger()).isEqualTo(Int.MIN_VALUE)
+  }
+
+  @Test
+  fun asInteger_outOfBounds_throws() {
+    val unsignedOutOfRange = UnsignedInteger(Int.MAX_VALUE.toLong() + 1)
+    assertThrows(CborException::class.java) { unsignedOutOfRange.asInteger() }
+
+    val negativeOutOfRange = NegativeInteger(Int.MIN_VALUE.toLong() - 1)
+    assertThrows(CborException::class.java) { negativeOutOfRange.asInteger() }
+  }
+
+  @Test
+  fun asLong_validValues_success() {
+    assertThat(UnsignedInteger(0).asLong()).isEqualTo(0L)
+    assertThat(UnsignedInteger(Long.MAX_VALUE).asLong()).isEqualTo(Long.MAX_VALUE)
+    assertThat(NegativeInteger(-1).asLong()).isEqualTo(-1L)
+    assertThat(NegativeInteger(Long.MIN_VALUE).asLong()).isEqualTo(Long.MIN_VALUE)
+  }
+
+  @Test
+  fun asLong_outOfBounds_throws() {
+    val unsignedOutOfRange = UnsignedInteger(BigInteger.valueOf(Long.MAX_VALUE).add(BigInteger.ONE))
+    assertThrows(CborException::class.java) { unsignedOutOfRange.asLong() }
+
+    val negativeOutOfRange =
+      NegativeInteger(BigInteger.valueOf(Long.MIN_VALUE).subtract(BigInteger.ONE))
+    assertThrows(CborException::class.java) { negativeOutOfRange.asLong() }
+  }
 }