From a438b25f762b279c3a3f7d8c308e2e4f04b2b84a Mon Sep 17 00:00:00 2001
From: Zhan Jiang <zhanjiang@google.com>
Date: Fri, 16 Aug 2024 16:08:45 -0700
Subject: [PATCH] Allow attributes to be populated through S2A

---
 s2a_options.go | 71 ++++++++++++++++++++++++++++++++++++++------------
 1 file changed, 54 insertions(+), 17 deletions(-)

diff --git a/s2a_options.go b/s2a_options.go
index d4252a9..45558a9 100644
--- a/s2a_options.go
+++ b/s2a_options.go
@@ -28,7 +28,7 @@ import (
 	"github.com/google/s2a-go/stream"
 	"google.golang.org/grpc/credentials"
 
-	s2apbv1 "github.com/google/s2a-go/internal/proto/common_go_proto"
+	s2av1pb "github.com/google/s2a-go/internal/proto/common_go_proto"
 	s2apb "github.com/google/s2a-go/internal/proto/v2/common_go_proto"
 )
 
@@ -36,6 +36,17 @@ import (
 type Identity interface {
 	// Name returns the name of the identity.
 	Name() string
+	Attributes() map[string]string
+}
+
+type UnspecifiedID struct {
+	Attr map[string]string
+}
+
+func (u *UnspecifiedID) Name() string { return "" }
+
+func (u *UnspecifiedID) Attributes() map[string]string {
+	return u.Attr
 }
 
 type spiffeID struct {
@@ -44,10 +55,10 @@ type spiffeID struct {
 
 func (s *spiffeID) Name() string { return s.spiffeID }
 
+func (spiffeID) Attributes() map[string]string { return nil }
+
 // NewSpiffeID creates a SPIFFE ID from id.
-func NewSpiffeID(id string) Identity {
-	return &spiffeID{spiffeID: id}
-}
+func NewSpiffeID(id string) Identity { return &spiffeID{spiffeID: id} }
 
 type hostname struct {
 	hostname string
@@ -55,10 +66,10 @@ type hostname struct {
 
 func (h *hostname) Name() string { return h.hostname }
 
+func (hostname) Attributes() map[string]string { return nil }
+
 // NewHostname creates a hostname from name.
-func NewHostname(name string) Identity {
-	return &hostname{hostname: name}
-}
+func NewHostname(name string) Identity { return &hostname{hostname: name} }
 
 type uid struct {
 	uid string
@@ -66,10 +77,10 @@ type uid struct {
 
 func (h *uid) Name() string { return h.uid }
 
+func (uid) Attributes() map[string]string { return nil }
+
 // NewUID creates a UID from name.
-func NewUID(name string) Identity {
-	return &uid{uid: name}
-}
+func NewUID(name string) Identity { return &uid{uid: name} }
 
 // VerificationModeType specifies the mode that S2A must use to verify the peer
 // certificate chain.
@@ -203,17 +214,30 @@ func DefaultServerOptions(s2aAddress string) *ServerOptions {
 	}
 }
 
-func toProtoIdentity(identity Identity) (*s2apbv1.Identity, error) {
+func toProtoIdentity(identity Identity) (*s2av1pb.Identity, error) {
 	if identity == nil {
 		return nil, nil
 	}
 	switch id := identity.(type) {
 	case *spiffeID:
-		return &s2apbv1.Identity{IdentityOneof: &s2apbv1.Identity_SpiffeId{SpiffeId: id.Name()}}, nil
+		return &s2av1pb.Identity{
+			IdentityOneof: &s2av1pb.Identity_SpiffeId{SpiffeId: id.Name()},
+			Attributes:    id.Attributes(),
+		}, nil
 	case *hostname:
-		return &s2apbv1.Identity{IdentityOneof: &s2apbv1.Identity_Hostname{Hostname: id.Name()}}, nil
+		return &s2av1pb.Identity{
+			IdentityOneof: &s2av1pb.Identity_Hostname{Hostname: id.Name()},
+			Attributes:    id.Attributes(),
+		}, nil
 	case *uid:
-		return &s2apbv1.Identity{IdentityOneof: &s2apbv1.Identity_Uid{Uid: id.Name()}}, nil
+		return &s2av1pb.Identity{
+			IdentityOneof: &s2av1pb.Identity_Uid{Uid: id.Name()},
+			Attributes:    id.Attributes(),
+		}, nil
+	case *UnspecifiedID:
+		return &s2av1pb.Identity{
+			Attributes: id.Attributes(),
+		}, nil
 	default:
 		return nil, errors.New("unrecognized identity type")
 	}
@@ -225,11 +249,24 @@ func toV2ProtoIdentity(identity Identity) (*s2apb.Identity, error) {
 	}
 	switch id := identity.(type) {
 	case *spiffeID:
-		return &s2apb.Identity{IdentityOneof: &s2apb.Identity_SpiffeId{SpiffeId: id.Name()}}, nil
+		return &s2apb.Identity{
+			IdentityOneof: &s2apb.Identity_SpiffeId{SpiffeId: id.Name()},
+			Attributes:    id.Attributes(),
+		}, nil
 	case *hostname:
-		return &s2apb.Identity{IdentityOneof: &s2apb.Identity_Hostname{Hostname: id.Name()}}, nil
+		return &s2apb.Identity{
+			IdentityOneof: &s2apb.Identity_Hostname{Hostname: id.Name()},
+			Attributes:    id.Attributes(),
+		}, nil
 	case *uid:
-		return &s2apb.Identity{IdentityOneof: &s2apb.Identity_Uid{Uid: id.Name()}}, nil
+		return &s2apb.Identity{
+			IdentityOneof: &s2apb.Identity_Uid{Uid: id.Name()},
+			Attributes:    id.Attributes(),
+		}, nil
+	case *UnspecifiedID:
+		return &s2apb.Identity{
+			Attributes: id.Attributes(),
+		}, nil
 	default:
 		return nil, errors.New("unrecognized identity type")
 	}