From e9ea70935b91e69a5ebc21765c06a895dae6f523 Mon Sep 17 00:00:00 2001
From: d-reidenbach <66145057+d-reidenbach@users.noreply.github.com>
Date: Wed, 24 Jun 2020 17:24:10 +0000
Subject: [PATCH] alts: add peer attribute access for the ALTS context (#7108)

* Added Peer Attribute getter for ALTS plumbing
---
 .../io/grpc/alts/internal/AltsAuthContext.java    | 11 +++++++++++
 .../grpc/alts/internal/AltsAuthContextTest.java   | 15 ++++++++++++++-
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/alts/src/main/java/io/grpc/alts/internal/AltsAuthContext.java b/alts/src/main/java/io/grpc/alts/internal/AltsAuthContext.java
index ac0cd36c81..b35d4faf98 100644
--- a/alts/src/main/java/io/grpc/alts/internal/AltsAuthContext.java
+++ b/alts/src/main/java/io/grpc/alts/internal/AltsAuthContext.java
@@ -17,6 +17,7 @@
 package io.grpc.alts.internal;
 
 import com.google.common.annotations.VisibleForTesting;
+import java.util.Map;
 
 /** AltsAuthContext contains security-related context information about an ALTs connection. */
 public final class AltsAuthContext {
@@ -33,6 +34,7 @@ public final class AltsAuthContext {
             .setPeerServiceAccount(result.getPeerIdentity().getServiceAccount())
             .setLocalServiceAccount(result.getLocalIdentity().getServiceAccount())
             .setPeerRpcVersions(result.getPeerRpcVersions())
+            .putAllPeerAttributes(result.getPeerIdentity().getAttributes())
             .build();
   }
 
@@ -94,4 +96,13 @@ public final class AltsAuthContext {
   public RpcProtocolVersions getPeerRpcVersions() {
     return context.getPeerRpcVersions();
   }
+
+  /**
+   * Get peer attributes.
+   *
+   * @return the context's peer attributes.
+   */
+  public Map<String, String> getPeerAttributes() {
+    return context.getPeerAttributes();
+  }
 }
diff --git a/alts/src/test/java/io/grpc/alts/internal/AltsAuthContextTest.java b/alts/src/test/java/io/grpc/alts/internal/AltsAuthContextTest.java
index 2fc2c43464..c73a145a4c 100644
--- a/alts/src/test/java/io/grpc/alts/internal/AltsAuthContextTest.java
+++ b/alts/src/test/java/io/grpc/alts/internal/AltsAuthContextTest.java
@@ -18,6 +18,8 @@ package io.grpc.alts.internal;
 
 import static org.junit.Assert.assertEquals;
 
+import java.util.HashMap;
+import java.util.Map;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -35,12 +37,17 @@ public final class AltsAuthContextTest {
   private static final String TEST_LOCAL_SERVICE_ACCOUNT = "local@gserviceaccount.com";
   private static final String TEST_PEER_SERVICE_ACCOUNT = "peer@gserviceaccount.com";
   private static final String TEST_RECORD_PROTOCOL = "ALTSRP_GCM_AES128";
+  private static final String TEST_PEER_ATTRIBUTES_KEY = "peer";
+  private static final String TEST_PEER_ATTRIBUTES_VALUE = "attributes";
 
+  private Map<String, String> testPeerAttributes;
   private HandshakerResult handshakerResult;
   private RpcProtocolVersions rpcVersions;
 
   @Before
   public void setUp() {
+    testPeerAttributes = new HashMap<String, String>();
+    testPeerAttributes.put(TEST_PEER_ATTRIBUTES_KEY, TEST_PEER_ATTRIBUTES_VALUE);
     rpcVersions =
         RpcProtocolVersions.newBuilder()
             .setMaxRpcVersion(
@@ -54,11 +61,14 @@ public final class AltsAuthContextTest {
                     .setMinor(TEST_MIN_RPC_VERSION_MINOR)
                     .build())
             .build();
+    Identity.Builder peerIdentity = Identity.newBuilder()
+        .setServiceAccount(TEST_PEER_SERVICE_ACCOUNT);
+    peerIdentity.putAllAttributes(testPeerAttributes);
     handshakerResult =
         HandshakerResult.newBuilder()
             .setApplicationProtocol(TEST_APPLICATION_PROTOCOL)
             .setRecordProtocol(TEST_RECORD_PROTOCOL)
-            .setPeerIdentity(Identity.newBuilder().setServiceAccount(TEST_PEER_SERVICE_ACCOUNT))
+            .setPeerIdentity(peerIdentity)
             .setLocalIdentity(Identity.newBuilder().setServiceAccount(TEST_LOCAL_SERVICE_ACCOUNT))
             .setPeerRpcVersions(rpcVersions)
             .build();
@@ -73,5 +83,8 @@ public final class AltsAuthContextTest {
     assertEquals(TEST_PEER_SERVICE_ACCOUNT, authContext.getPeerServiceAccount());
     assertEquals(TEST_LOCAL_SERVICE_ACCOUNT, authContext.getLocalServiceAccount());
     assertEquals(rpcVersions, authContext.getPeerRpcVersions());
+    assertEquals(testPeerAttributes, authContext.getPeerAttributes());
+    assertEquals(TEST_PEER_ATTRIBUTES_VALUE, authContext.getPeerAttributes()
+        .get(TEST_PEER_ATTRIBUTES_KEY));
   }
 }