alts: Added peer attributes accessor for alts context and updated tes…

…t method (#3675)
grpc · Jun 17, 2020 · 9a46550 · 9a46550
1 parent dfc058c
commit 9a46550
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 6 deletions.
diff --git a/credentials/alts/internal/authinfo/authinfo.go b/credentials/alts/internal/authinfo/authinfo.go
@@ -48,6 +48,7 @@ func newAuthInfo(result *altspb.HandshakerResult) *altsAuthInfo {
 			PeerServiceAccount:  result.GetPeerIdentity().GetServiceAccount(),
 			LocalServiceAccount: result.GetLocalIdentity().GetServiceAccount(),
 			PeerRpcVersions:     result.GetPeerRpcVersions(),
+			PeerAttributes:      result.GetPeerIdentity().GetAttributes(),
 		},
 		CommonAuthInfo: credentials.CommonAuthInfo{SecurityLevel: credentials.PrivacyAndIntegrity},
 	}
@@ -87,3 +88,8 @@ func (s *altsAuthInfo) LocalServiceAccount() string {
 func (s *altsAuthInfo) PeerRPCVersions() *altspb.RpcProtocolVersions {
 	return s.p.GetPeerRpcVersions()
 }
+
+// PeerAttributes returns the context's peer attributes.
+func (s *altsAuthInfo) PeerAttributes() map[string]string {
+	return s.p.GetPeerAttributes()
+}
diff --git a/credentials/alts/internal/authinfo/authinfo_test.go b/credentials/alts/internal/authinfo/authinfo_test.go
@@ -35,15 +35,19 @@ func Test(t *testing.T) {
 }
 
 const (
-	testAppProtocol    = "my_app"
-	testRecordProtocol = "very_secure_protocol"
-	testPeerAccount    = "peer_service_account"
-	testLocalAccount   = "local_service_account"
-	testPeerHostname   = "peer_hostname"
-	testLocalHostname  = "local_hostname"
+	testAppProtocol             = "my_app"
+	testRecordProtocol          = "very_secure_protocol"
+	testPeerAccount             = "peer_service_account"
+	testLocalAccount            = "local_service_account"
+	testPeerHostname            = "peer_hostname"
+	testLocalHostname           = "local_hostname"
+	testLocalPeerAttributeKey   = "peer"
+	testLocalPeerAttributeValue = "attributes"
 )
 
 func (s) TestALTSAuthInfo(t *testing.T) {
+	testPeerAttributes := make(map[string]string)
+	testPeerAttributes[testLocalPeerAttributeKey] = testLocalPeerAttributeValue
 	for _, tc := range []struct {
 		result             *altspb.HandshakerResult
 		outAppProtocol     string
@@ -52,6 +56,7 @@ func (s) TestALTSAuthInfo(t *testing.T) {
 		outPeerAccount     string
 		outLocalAccount    string
 		outPeerRPCVersions *altspb.RpcProtocolVersions
+		outPeerAttributes  map[string]string
 	}{
 		{
 			&altspb.HandshakerResult{
@@ -61,6 +66,7 @@ func (s) TestALTSAuthInfo(t *testing.T) {
 					IdentityOneof: &altspb.Identity_ServiceAccount{
 						ServiceAccount: testPeerAccount,
 					},
+					Attributes: testPeerAttributes,
 				},
 				LocalIdentity: &altspb.Identity{
 					IdentityOneof: &altspb.Identity_ServiceAccount{
@@ -74,6 +80,7 @@ func (s) TestALTSAuthInfo(t *testing.T) {
 			testPeerAccount,
 			testLocalAccount,
 			nil,
+			testPeerAttributes,
 		},
 		{
 			&altspb.HandshakerResult{
@@ -83,6 +90,7 @@ func (s) TestALTSAuthInfo(t *testing.T) {
 					IdentityOneof: &altspb.Identity_Hostname{
 						Hostname: testPeerHostname,
 					},
+					Attributes: testPeerAttributes,
 				},
 				LocalIdentity: &altspb.Identity{
 					IdentityOneof: &altspb.Identity_Hostname{
@@ -115,6 +123,7 @@ func (s) TestALTSAuthInfo(t *testing.T) {
 					Minor: 11,
 				},
 			},
+			testPeerAttributes,
 		},
 	} {
 		authInfo := newAuthInfo(tc.result)
@@ -139,5 +148,9 @@ func (s) TestALTSAuthInfo(t *testing.T) {
 		if got, want := authInfo.PeerRPCVersions(), tc.outPeerRPCVersions; !reflect.DeepEqual(got, want) {
 			t.Errorf("authinfo.PeerRpcVersions()=%v, want %v", got, want)
 		}
+		if got, want := authInfo.PeerAttributes(), tc.outPeerAttributes; !reflect.DeepEqual(got, want) {
+			t.Errorf("authinfo.PeerAttributes()=%v, want %v", got, want)
+		}
+
 	}
 }