Build authorizationMap using a stream

google · Dec 5, 2023 · 738c750 · 738c750
1 parent 61d8232
commit 738c750
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 14 deletions.
diff --git a/server/src/main/java/com/google/android/attestation/AuthorizationList.java b/server/src/main/java/com/google/android/attestation/AuthorizationList.java
@@ -60,6 +60,7 @@
 import static com.google.android.attestation.Constants.KM_TAG_USER_AUTH_TYPE;
 import static com.google.android.attestation.Constants.KM_TAG_VENDOR_PATCH_LEVEL;
 import static com.google.android.attestation.Constants.UINT32_MAX;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Streams.stream;
 
@@ -73,9 +74,8 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Stream;
@@ -551,22 +551,28 @@ static AuthorizationList createAuthorizationList(
   }
 
   private static ParsedAuthorizationMap getAuthorizationMap(ASN1Encodable[] authorizationList) {
-    Map<Integer, ASN1Object> authorizationMap = new HashMap<>();
+    // authorizationMap must retain the order of authorizationList, otherwise
+    // the code searching for out of order tags below will break. Helpfully
+    // a ImmutableMap preserves insertion order.
+    //
+    // https://guava.dev/releases/23.0/api/docs/com/google/common/collect/ImmutableCollection.html
+    ImmutableMap<Integer, ASN1Object> authorizationMap =
+        Arrays.stream(authorizationList)
+            .map(ASN1TaggedObject::getInstance)
+            .collect(
+                toImmutableMap(
+                    ASN1TaggedObject::getTagNo,
+                    obj -> ASN1Util.getExplicitContextBaseObject(obj, obj.getTagNo())));
+
     List<Integer> unorderedTags = new ArrayList<>();
-    int currentTag = 0;
     int previousTag = 0;
-    for (ASN1Encodable entry : authorizationList) {
-      ASN1TaggedObject taggedEntry = ASN1TaggedObject.getInstance(entry);
-      previousTag = currentTag;
-      currentTag = taggedEntry.getTagNo();
+    for (int currentTag : authorizationMap.keySet()) {
       if (previousTag > currentTag) {
         unorderedTags.add(previousTag);
       }
-      authorizationMap.put(
-          currentTag, ASN1Util.getExplicitContextBaseObject(taggedEntry, taggedEntry.getTagNo()));
+      previousTag = currentTag;
     }
-    return new ParsedAuthorizationMap(
-        ImmutableMap.copyOf(authorizationMap), ImmutableList.copyOf(unorderedTags));
+    return new ParsedAuthorizationMap(authorizationMap, ImmutableList.copyOf(unorderedTags));
   }
 
   @VisibleForTesting

diff --git a/server/src/test/java/com/google/android/attestation/AuthorizationListTest.java b/server/src/test/java/com/google/android/attestation/AuthorizationListTest.java
@@ -41,6 +41,9 @@
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
 import java.io.IOException;
 import java.time.Instant;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import org.bouncycastle.asn1.ASN1Encodable;
 import org.bouncycastle.asn1.ASN1Sequence;
 import org.bouncycastle.asn1.ASN1TaggedObject;
@@ -98,7 +101,7 @@ public class AuthorizationListTest {
   private static ASN1Encodable[] getEncodableAuthorizationList(String extensionData)
       throws IOException {
     byte[] extensionDataBytes = Base64.decode(extensionData);
-    return ((ASN1Sequence) ASN1Sequence.fromByteArray(extensionDataBytes)).toArray();
+    return ASN1Sequence.getInstance(extensionDataBytes).toArray();
   }
 
   @Test
@@ -208,7 +211,7 @@ public void testCreateAndParse() throws IOException {
   }
 
   @Test
-  public void testCreateWithUnorderedTagsAndParse() throws IOException {
+  public void testCreateWithUnorderedTagAndParse() throws IOException {
     // Create a encodable auth list from valid attestation extension data.
     ASN1Encodable[] encodableAuthList =
         getEncodableAuthorizationList(EXTENTION_DATA_WITH_INDIVIDUAL_ATTESTATION);
@@ -224,6 +227,19 @@ public void testCreateWithUnorderedTagsAndParse() throws IOException {
     assertThat(authorizationList.unorderedTags).containsExactly(taggedEntry.getTagNo());
   }
 
+  @Test
+  public void testCreateWithUnorderedTagsAndParse() throws IOException {
+    List<ASN1Encodable> encodableAuthList =
+        Arrays.asList(getEncodableAuthorizationList(EXTENTION_DATA_WITH_INDIVIDUAL_ATTESTATION));
+    ASN1TaggedObject taggedEntry = ASN1TaggedObject.getInstance(encodableAuthList.get(0));
+    Collections.swap(encodableAuthList, 3, 5);
+    Collections.swap(encodableAuthList, 12, 13);
+    AuthorizationList authorizationList =
+        AuthorizationList.createAuthorizationList(
+            encodableAuthList.stream().toArray(ASN1Encodable[]::new), ATTESTATION_VERSION);
+    assertThat(authorizationList.unorderedTags).containsExactly(6, 200, 712);
+  }
+
   @Test
   public void testPaddingModeMap(@TestParameter AuthorizationList.PaddingMode paddingMode) {
     assertThat(