Merge pull request #435 from hzxuzhonghu/fix-authStore

Fix authPolicy remove

api/v2/workloadapi/security/authorization_help.go

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2024 The Kmesh Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at:
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package security
+
+// ResourceName returns the unique key of Workload.
+func (x *Authorization) ResourceName() string {
+	return x.GetNamespace() + "/" + x.GetName()
+}

pkg/auth/policy_store.go

@@ -27,36 +27,32 @@ import (
 
 type policyStore struct {
 	// byKey maintains a mapping of ns/name to policy
-	byKey map[string]authPolicy
+	byKey map[string]*security.Authorization
 
 	// byNamespace maintains a mapping of namespace (or "" for global) to policy names
 	byNamespace map[string]sets.Set[string]
 
 	rwLock sync.RWMutex
 }
 
-func newPolicystore() *policyStore {
+func newPolicyStore() *policyStore {
 	return &policyStore{
-		byKey:       make(map[string]authPolicy),
+		byKey:       make(map[string]*security.Authorization),
 		byNamespace: make(map[string]sets.Set[string]),
 	}
 }
 
-func (ps *policyStore) updatePolicy(auth *security.Authorization) error {
-	if auth == nil {
+func (ps *policyStore) updatePolicy(authPolicy *security.Authorization) error {
+	if authPolicy == nil {
 		return nil
 	}
+	key := authPolicy.ResourceName()
 
-	authPolicy := authPolicy{
-		auth,
-	}
-	key := authPolicy.Key()
-
+	ps.rwLock.Lock()
+	defer ps.rwLock.Unlock()
 	var ns string
 	switch authPolicy.GetScope() {
 	case security.Scope_WORKLOAD_SELECTOR:
-		ps.rwLock.Lock()
-		defer ps.rwLock.Unlock()
 		// only update 'byKey' cache for Scope_WORKLOAD_SELECTOR
 		ps.byKey[key] = authPolicy
 		return nil
@@ -68,9 +64,6 @@ func (ps *policyStore) updatePolicy(auth *security.Authorization) error {
 		return fmt.Errorf("invalid scope %v of authorization policy", authPolicy.GetScope())
 	}
 
-	ps.rwLock.Lock()
-	defer ps.rwLock.Unlock()
-
 	if s, ok := ps.byNamespace[ns]; !ok {
 		ps.byNamespace[ns] = sets.New(key)
 	} else {
@@ -89,13 +82,17 @@ func (ps *policyStore) removePolicy(policyKey string) {
 		log.Warnf("Auth policy key %s does not exist in byKey", policyKey)
 		return
 	}
+	// remove authPolicy from byKey
+	delete(ps.byKey, policyKey)
 
 	var ns string
 	switch authPolicy.Scope {
 	case security.Scope_GLOBAL:
 		ns = ""
 	case security.Scope_NAMESPACE:
 		ns = authPolicy.GetNamespace()
+	default:
+		return
 	}
 
 	// remove authPolicy key from byNamespace
@@ -105,26 +102,15 @@ func (ps *policyStore) removePolicy(policyKey string) {
 			delete(ps.byNamespace, ns)
 		}
 	}
-
-	// remove authPolicy from byKey
-	delete(ps.byKey, policyKey)
 }
 
-// getByNamesapce returns a copied set of policy name in namespace, or an empty set if namespace not exists
-func (ps *policyStore) getByNamesapce(namespace string) sets.Set[string] {
+// getByNamespace returns a copied set of policy name in namespace, or an empty set if namespace not exists
+func (ps *policyStore) getByNamespace(namespace string) []string {
 	ps.rwLock.RLock()
 	defer ps.rwLock.RUnlock()
 
 	if s, ok := ps.byNamespace[namespace]; ok {
-		return s.Copy()
+		return s.UnsortedList()
 	}
-	return sets.New[string]()
-}
-
-type authPolicy struct {
-	*security.Authorization
-}
-
-func (ap *authPolicy) Key() string {
-	return fmt.Sprintf("%s/%s", ap.GetNamespace(), ap.GetName())
+	return nil
 }

pkg/auth/policy_store_test.go

@@ -77,7 +77,7 @@ func Test_policyStore_updatePolicy(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			ps := newPolicystore()
+			ps := newPolicyStore()
 			if err := ps.updatePolicy(tt.args.auth); (err != nil) != tt.wantErr {
 				t.Errorf("policyStore.updatePolicy() error = %v, wantErr %v", err, tt.wantErr)
 			}
@@ -109,7 +109,7 @@ func Test_policyStore_removePolicy(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			ps := newPolicystore()
+			ps := newPolicyStore()
 			ps.removePolicy(tt.args.policyKey)
 		})
 	}

pkg/auth/rbac.go

@@ -70,7 +70,7 @@ type rbacConnection struct {
 
 func NewRbac(workloadObj *bpf.BpfKmeshWorkload, workloadCache cache.WorkloadCache) *Rbac {
 	return &Rbac{
-		policyStore:   newPolicystore(),
+		policyStore:   newPolicyStore(),
 		workloadCache: workloadCache,
 		bpfWorkload:   workloadObj,
 	}
@@ -160,14 +160,16 @@ func (r *Rbac) RemovePolicy(policyKey string) {
 }
 
 func (r *Rbac) doRbac(conn *rbacConnection) bool {
-	var dstWorkload *workloadapi.Workload
-	if len(conn.dstIp) > 0 {
-		var networkAddress cache.NetworkAddress
-		networkAddress.Network = conn.dstNetwork
-		networkAddress.Address, _ = netip.AddrFromSlice(conn.dstIp)
-		dstWorkload = r.workloadCache.GetWorkloadByAddr(networkAddress)
+	var networkAddress cache.NetworkAddress
+	networkAddress.Network = conn.dstNetwork
+	networkAddress.Address, _ = netip.AddrFromSlice(conn.dstIp)
+	dstWorkload := r.workloadCache.GetWorkloadByAddr(networkAddress)
+	// If no workload found, deny
+	if dstWorkload == nil {
+		return false
 	}
 
+	// TODO: maybe cache them for performance issue
 	allowPolicies, denyPolicies := r.aggregate(dstWorkload)
 
 	// 1. If there is ANY deny policy, deny the request
@@ -193,17 +195,14 @@ func (r *Rbac) doRbac(conn *rbacConnection) bool {
 	return false
 }
 
-func (r *Rbac) aggregate(workload *workloadapi.Workload) (allowPolicies, denyPolicies []authPolicy) {
-	allowPolicies = make([]authPolicy, 0)
-	denyPolicies = make([]authPolicy, 0)
+func (r *Rbac) aggregate(workload *workloadapi.Workload) (allowPolicies, denyPolicies []*security.Authorization) {
+	allowPolicies = make([]*security.Authorization, 0)
+	denyPolicies = make([]*security.Authorization, 0)
 
-	// Collect policy names from workload, global namespace and namespace
-	policyNames := r.policyStore.getByNamesapce("").UnsortedList()
-	if workload != nil {
-		policyNames = append(append(policyNames,
-			r.policyStore.getByNamesapce(workload.Namespace).UnsortedList()...),
-			workload.GetAuthorizationPolicies()...)
-	}
+	// Collect policy names from workload,  namespace and global(root namespace)
+	policyNames := workload.GetAuthorizationPolicies()
+	policyNames = append(policyNames, r.policyStore.getByNamespace(workload.Namespace)...)
+	policyNames = append(policyNames, r.policyStore.getByNamespace("")...)
 
 	for _, policyName := range policyNames {
 		if policy, ok := r.policyStore.byKey[policyName]; ok {
@@ -217,7 +216,7 @@ func (r *Rbac) aggregate(workload *workloadapi.Workload) (allowPolicies, denyPol
 	return
 }
 
-func matches(conn *rbacConnection, policy authPolicy) bool {
+func matches(conn *rbacConnection, policy *security.Authorization) bool {
 	if policy.GetRules() == nil {
 		return false
 	}