feat(sdk): Lets SDK use PublicKeys from policy service

examples/cmd/attributes.go

@@ -21,10 +21,10 @@ import (
 var (
 	ns         string
 	attr       string
-	kas        string
+	kases      []string
 	longformat bool
 	rule       string
-	values     string
+	values     []string
 	unsafeBool bool
 )
 
@@ -43,7 +43,7 @@ func init() {
 	}
 	add.Flags().StringVarP(&attr, "attr", "a", "", "attribute prefix, e.g. http://name.space/attr/name")
 	add.Flags().StringVarP(&rule, "rule", "", "allof", "attribute type, either allof, anyof, or hierarchy")
-	add.Flags().StringVarP(&values, "values", "v", "", "list of attribute values")
+	add.Flags().StringSliceVarP(&values, "values", "v", []string{}, "list of attribute values")
 	attributes.AddCommand(add)
 
 	assign := &cobra.Command{
@@ -54,8 +54,8 @@ func init() {
 		},
 	}
 	assign.Flags().StringVarP(&attr, "attr", "a", "", "attribute prefix, e.g. https://name.space/attr/name")
-	assign.Flags().StringVarP(&kas, "kas", "k", "", "which kas to assign")
-	assign.Flags().StringVarP(&values, "values", "v", "", "any attribute values to include; if empty, applies to all")
+	assign.Flags().StringSliceVarP(&kases, "kas", "k", []string{}, "which kas to assign")
+	assign.Flags().StringSliceVarP(&values, "values", "v", []string{}, "any attribute values to include; if empty, applies to all")
 	attributes.AddCommand(assign)
 
 	list := &cobra.Command{
@@ -80,7 +80,7 @@ func init() {
 		},
 	}
 	remove.Flags().StringVarP(&attr, "attr", "a", "", "attribute prefix, e.g. http://name.space/attr/name")
-	remove.Flags().StringVarP(&values, "values", "v", "", "list of attribute values to remove; if absent, removes all")
+	remove.Flags().StringSliceVarP(&values, "values", "v", []string{}, "list of attribute values to remove; if absent, removes all")
 	remove.Flags().BoolVarP(&unsafeBool, "unsafe", "f", false, "delete for real; otherwise deactivate (soft delete)")
 	attributes.AddCommand(remove)
 
@@ -92,8 +92,8 @@ func init() {
 		},
 	}
 	unassign.Flags().StringVarP(&attr, "attr", "a", "", "attribute prefix, e.g. https://name.space/attr/name")
-	unassign.Flags().StringVarP(&kas, "kas", "k", "", "which kas to assign")
-	unassign.Flags().StringVarP(&values, "values", "v", "", "any attribute values to include; if empty, applies to all")
+	unassign.Flags().StringSliceVarP(&kases, "kas", "k", []string{}, "which kases to assign")
+	unassign.Flags().StringSliceVarP(&values, "values", "v", []string{}, "any attribute values to include; if empty, applies to all")
 	attributes.AddCommand(unassign)
 
 	ExamplesCmd.AddCommand(attributes)
@@ -245,7 +245,7 @@ func addAttribute(cmd *cobra.Command) error {
 		slog.Error("url.PathUnescape(attr)", "err", err, "attr", m[2])
 		return err
 	}
-	aid, err := upsertAttr(cmd.Context(), s, nsu, attr, strings.Split(values, " "))
+	aid, err := upsertAttr(cmd.Context(), s, nsu, attr, values)
 	if err != nil {
 		return err
 	}
@@ -275,7 +275,7 @@ func removeAttribute(cmd *cobra.Command) error {
 	if err != nil {
 		return err
 	}
-	if values == "" {
+	if len(values) == 0 {
 		if unsafeBool {
 			resp, err := s.Unsafe.UnsafeDeleteAttribute(cmd.Context(), &unsafe.UnsafeDeleteAttributeRequest{
 				Id:  auuid,
@@ -298,7 +298,7 @@ func removeAttribute(cmd *cobra.Command) error {
 		slog.Info("deactivated attribute", "attr", attr, "resp", resp)
 		return nil
 	} else {
-		for _, v := range strings.Split(values, " ") {
+		for _, v := range values {
 			avu, err := avuuid(cmd.Context(), s, auuid, v)
 			if err != nil {
 				return err
@@ -359,16 +359,18 @@ func assignAttribute(cmd *cobra.Command, assign bool) error {
 
 	var kasids []string
 	switch {
-	case kas != "":
-		kasid, err := upsertKasRegistration(cmd.Context(), s, kas, nil)
-		if err != nil {
-			return err
+	case len(kases) != 0:
+		for _, kas := range kases {
+			kasid, err := upsertKasRegistration(cmd.Context(), s, kas, nil)
+			if err != nil {
+				return err
+			}
+			kasids = append(kasids, kasid)
+			kasById[kasid] = kas
 		}
-		kasids = append(kasids, kasid)
-		kasById[kasid] = kas
 	case assign:
 		return fmt.Errorf("assign must take a `--kas` parameter")
-	case values == "":
+	case len(values) == 0:
 		// look up all kasids associated with the attribute
 		ar, err := s.Attributes.GetAttribute(cmd.Context(), &attributes.GetAttributeRequest{Id: auuid})
 		if err != nil {
@@ -378,12 +380,11 @@ func assignAttribute(cmd *cobra.Command, assign bool) error {
 			kasids = append(kasids, b.GetId())
 			kasById[b.GetId()] = b.GetUri()
 		}
+	case len(values) > 1:
+		return fmt.Errorf("TODO: unassign from multiple values at a time")
 	default:
 		// look up all kasids associated with the value
-		if strings.Index(values, " ") >= 0 {
-			return fmt.Errorf("TODO: unassign from multiple values at a time")
-		}
-		avu, err := avuuid(cmd.Context(), s, auuid, values)
+		avu, err := avuuid(cmd.Context(), s, auuid, values[0])
 		if err != nil {
 			return err
 		}
@@ -398,7 +399,7 @@ func assignAttribute(cmd *cobra.Command, assign bool) error {
 	}
 
 	for _, kasid := range kasids {
-		if values == "" {
+		if len(values) == 0 {
 			if assign {
 				r, err := s.Attributes.AssignKeyAccessServerToAttribute(cmd.Context(), &attributes.AssignKeyAccessServerToAttributeRequest{
 					AttributeKeyAccessServer: &attributes.AttributeKeyAccessServer{
@@ -423,7 +424,7 @@ func assignAttribute(cmd *cobra.Command, assign bool) error {
 				cmd.Printf("successfully unassigned [%s] from [%s] (binding %v)\n", attr, kasById[kasid], *r.GetAttributeKeyAccessServer())
 			}
 		} else {
-			for _, v := range strings.Split(values, " ") {
+			for _, v := range values {
 				avu, err := avuuid(cmd.Context(), s, auuid, v)
 				if err != nil {
 					return err

examples/cmd/kas.go

@@ -12,7 +12,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
-var key string
+var algorithm, kas, key, keyIdentifier string
 
 func init() {
 	kasc := &cobra.Command{
@@ -28,8 +28,12 @@ func init() {
 			return updateKas(cmd)
 		},
 	}
+	// Note we currently only store one pk at a time. must be fixed for nano tests
+	update.Flags().StringVarP(&algorithm, "algorithm", "", "", "algorithm used with the public key")
 	update.Flags().StringVarP(&kas, "kas", "k", "", "kas uri")
 	update.Flags().StringVarP(&key, "public-key", "", "", "public key value, e.g. $(<my-key.pem)")
+	update.Flags().StringVarP(&keyIdentifier, "kid", "", "", "key identifier used to uniquely identify a key across rotations")
+
 	kasc.AddCommand(update)
 
 	list := &cobra.Command{
@@ -146,7 +150,25 @@ func updateKas(cmd *cobra.Command) error {
 	defer s.Close()
 
 	var pk *policy.PublicKey
-	if key != "" {
+	switch {
+	case keyIdentifier != "":
+		if key == "" || algorithm == "" {
+			err := fmt.Errorf("if --kid is found, --public-key and --algorithm must also be specified")
+			return err
+		}
+		pk = new(policy.PublicKey)
+		pk.PublicKey = &policy.PublicKey_Cached{
+			Cached: &policy.KasPublicKeySet{
+				Keys: []*policy.KasPublicKey{
+					{
+						Pem: key,
+						Kid: keyIdentifier,
+						Alg: algorithm,
+					},
+				},
+			},
+		}
+	case key != "":
 		pk = new(policy.PublicKey)
 		pk.PublicKey = &policy.PublicKey_Local{
 			Local: key,

sdk/internal/autoconfigure/granter.go → sdk/granter.go

@@ -1,4 +1,4 @@
-package autoconfigure
+package sdk
 
 import (
 	"context"
@@ -28,7 +28,7 @@ const (
 )
 
 // Represents a which KAS a split with the associated ID should shared with.
-type SplitStep struct {
+type keySplitStep struct {
 	KAS, SplitID string
 }
 
@@ -163,7 +163,7 @@ func (a AttributeValueFQN) Name() string {
 }
 
 // Structure capable of generating a split plan from a given set of data tags.
-type Granter struct {
+type granter struct {
 	policy []AttributeValueFQN
 	grants map[string]*keyAccessGrant
 }
@@ -173,15 +173,15 @@ type keyAccessGrant struct {
 	kases []string
 }
 
-func (r Granter) addGrant(fqn AttributeValueFQN, kas string, attr *policy.Attribute) {
+func (r granter) addGrant(fqn AttributeValueFQN, kas string, attr *policy.Attribute) {
 	if _, ok := r.grants[fqn.key]; ok {
 		r.grants[fqn.key].kases = append(r.grants[fqn.key].kases, kas)
 	} else {
 		r.grants[fqn.key] = &keyAccessGrant{attr, []string{kas}}
 	}
 }
 
-func (r Granter) addAllGrants(fqn AttributeValueFQN, gs []*policy.KeyAccessServer, attr *policy.Attribute) {
+func (r granter) addAllGrants(fqn AttributeValueFQN, gs []*policy.KeyAccessServer, attr *policy.Attribute) {
 	for _, g := range gs {
 		if g != nil {
 			r.addGrant(fqn, g.GetUri(), attr)
@@ -194,12 +194,12 @@ func (r Granter) addAllGrants(fqn AttributeValueFQN, gs []*policy.KeyAccessServe
 	}
 }
 
-func (r Granter) byAttribute(fqn AttributeValueFQN) *keyAccessGrant {
+func (r granter) byAttribute(fqn AttributeValueFQN) *keyAccessGrant {
 	return r.grants[fqn.key]
 }
 
 // Gets a list of directory of KAS grants for a list of attribute FQNs
-func NewGranterFromService(ctx context.Context, as attributes.AttributesServiceClient, fqns ...AttributeValueFQN) (Granter, error) {
+func newGranterFromService(ctx context.Context, keyCache *kasKeyCache, as attributes.AttributesServiceClient, fqns ...AttributeValueFQN) (granter, error) {
 	fqnsStr := make([]string, len(fqns))
 	for i, v := range fqns {
 		fqnsStr[i] = v.String()
@@ -213,10 +213,10 @@ func NewGranterFromService(ctx context.Context, as attributes.AttributesServiceC
 			},
 		})
 	if err != nil {
-		return Granter{}, err
+		return granter{}, err
 	}
 
-	grants := Granter{
+	grants := granter{
 		policy: fqns,
 		grants: make(map[string]*keyAccessGrant),
 	}
@@ -228,21 +228,40 @@ func NewGranterFromService(ctx context.Context, as attributes.AttributesServiceC
 		def := pair.GetAttribute()
 		if def != nil {
 			grants.addAllGrants(fqn, def.GetGrants(), def)
+			storeKeysToCache(def.GetGrants(), keyCache)
 		}
 		v := pair.GetValue()
 		if v != nil {
 			grants.addAllGrants(fqn, v.GetGrants(), def)
+			storeKeysToCache(v.GetGrants(), keyCache)
 		}
 	}
 
 	return grants, nil
 }
 
+func storeKeysToCache(kases []*policy.KeyAccessServer, c *kasKeyCache) {
+	for _, kas := range kases {
+		if kas.GetPublicKey() == nil || kas.GetPublicKey().GetCached() == nil || kas.GetPublicKey().GetCached().GetKeys() == nil || len(kas.GetPublicKey().GetCached().GetKeys()) == 0 {
+			slog.Debug("no cached key in policy service", "kas", kas.GetUri())
+			continue
+		}
+		for _, ki := range kas.GetPublicKey().GetCached().GetKeys() {
+			c.store(KASInfo{
+				URL:       kas.GetUri(),
+				KID:       ki.GetKid(),
+				Algorithm: ki.GetAlg(),
+				PublicKey: ki.GetPem(),
+			})
+		}
+	}
+}
+
 // Given a policy (list of data attributes or tags),
 // get a set of grants from attribute values to KASes.
 // Unlike `NewGranterFromService`, this works offline.
-func NewGranterFromAttributes(attrs ...*policy.Value) (Granter, error) {
-	grants := Granter{
+func newGranterFromAttributes(attrs ...*policy.Value) (granter, error) {
+	grants := granter{
 		grants: make(map[string]*keyAccessGrant),
 		policy: make([]AttributeValueFQN, len(attrs)),
 	}
@@ -254,7 +273,7 @@ func NewGranterFromAttributes(attrs ...*policy.Value) (Granter, error) {
 		grants.policy[i] = fqn
 		def := v.GetAttribute()
 		if def == nil {
-			return Granter{}, fmt.Errorf("no associated definition with value [%s]", fqn)
+			return granter{}, fmt.Errorf("no associated definition with value [%s]", fqn)
 		}
 		grants.addAllGrants(fqn, def.GetGrants(), def)
 		grants.addAllGrants(fqn, v.GetGrants(), def)
@@ -263,34 +282,6 @@ func NewGranterFromAttributes(attrs ...*policy.Value) (Granter, error) {
 	return grants, nil
 }
 
-type AttributeService struct {
-	dict map[AttributeNameFQN]*policy.Attribute
-}
-
-func (s *AttributeService) Put(ad *policy.Attribute) error {
-	if s.dict == nil {
-		s.dict = make(map[AttributeNameFQN]*policy.Attribute)
-	}
-	prefix, err := NewAttributeNameFQN(ad.GetFqn())
-	if err != nil {
-		return err
-	}
-	if _, exists := s.dict[prefix]; exists {
-		return fmt.Errorf("ad prefix already found [%s]", prefix)
-	}
-	s.dict[prefix] = ad
-	return nil
-}
-
-// Given an attribute without a value (everything before /value/...), get the definition
-func (s *AttributeService) Get(prefix AttributeNameFQN) (*policy.Attribute, error) {
-	ad, exists := s.dict[prefix]
-	if !exists {
-		return nil, fmt.Errorf("[404] Unknown attribute type: [%s], not in [%v]", prefix, s.dict)
-	}
-	return ad, nil
-}
-
 type singleAttributeClause struct {
 	def    *policy.Attribute
 	values []AttributeValueFQN
@@ -329,7 +320,7 @@ func (e attributeBooleanExpression) String() string {
 	return sb.String()
 }
 
-func (r Granter) Plan(defaultKas []string, genSplitID func() string) ([]SplitStep, error) {
+func (r granter) plan(defaultKas []string, genSplitID func() string) ([]keySplitStep, error) {
 	b := r.constructAttributeBoolean()
 	k, err := r.insertKeysForAttribute(*b)
 	if err != nil {
@@ -344,29 +335,29 @@ func (r Granter) Plan(defaultKas []string, genSplitID func() string) ([]SplitSte
 		case 0:
 			return nil, fmt.Errorf("no default KAS specified; required for grantless plans")
 		case 1:
-			return []SplitStep{{KAS: defaultKas[0]}}, nil
+			return []keySplitStep{{KAS: defaultKas[0]}}, nil
 		default:
-			p := make([]SplitStep, 0, len(defaultKas))
+			p := make([]keySplitStep, 0, len(defaultKas))
 			for _, kas := range defaultKas {
-				p = append(p, SplitStep{KAS: kas, SplitID: genSplitID()})
+				p = append(p, keySplitStep{KAS: kas, SplitID: genSplitID()})
 			}
 			return p, nil
 		}
 	}
-	p := make([]SplitStep, 0, l)
+	p := make([]keySplitStep, 0, l)
 	for _, v := range k.values {
 		splitID := ""
 		if l > 1 {
 			splitID = genSplitID()
 		}
 		for _, o := range v.values {
-			p = append(p, SplitStep{KAS: o.kas, SplitID: splitID})
+			p = append(p, keySplitStep{KAS: o.kas, SplitID: splitID})
 		}
 	}
 	return p, nil
 }
 
-func (r Granter) constructAttributeBoolean() *attributeBooleanExpression {
+func (r granter) constructAttributeBoolean() *attributeBooleanExpression {
 	prefixes := make(map[string]*singleAttributeClause)
 	sortedPrefixes := make([]string, 0)
 	for _, aP := range r.policy {
@@ -447,7 +438,7 @@ func ruleToOperator(e policy.AttributeRuleTypeEnum) string {
 	return ""
 }
 
-func (r *Granter) insertKeysForAttribute(e attributeBooleanExpression) (booleanKeyExpression, error) {
+func (r *granter) insertKeysForAttribute(e attributeBooleanExpression) (booleanKeyExpression, error) {
 	kcs := make([]keyClause, 0, len(e.must))
 	for _, clause := range e.must {
 		kcv := make([]publicKeyInfo, 0, len(clause.values))