From 8c007923dc99d2344cd31ace9f6bd4ccc0443126 Mon Sep 17 00:00:00 2001
From: Qiyue Yao <yaoq@vmware.com>
Date: Fri, 8 Dec 2023 16:17:58 -0800
Subject: [PATCH] Fix endpoint querier rule index

The current endpoint querier rule index shows the index
of the rule among all matched rules in the policy for
this endpoint, which is not super useful for the users.
This change updates the rule index to show the rule index
among all rules in the policy.

Fixes #5782

Signed-off-by: Qiyue Yao <yaoq@vmware.com>
---
 .../networkpolicy/endpoint_querier.go         | 14 +++--
 .../networkpolicy/endpoint_querier_test.go    | 60 +++++++++++++++++--
 2 files changed, 66 insertions(+), 8 deletions(-)

diff --git a/pkg/controller/networkpolicy/endpoint_querier.go b/pkg/controller/networkpolicy/endpoint_querier.go
index b8b39c46fa8..b22234dc2f4 100644
--- a/pkg/controller/networkpolicy/endpoint_querier.go
+++ b/pkg/controller/networkpolicy/endpoint_querier.go
@@ -20,6 +20,7 @@ package networkpolicy
 import (
 	"k8s.io/apimachinery/pkg/types"
 
+	"antrea.io/antrea/pkg/apis/controlplane"
 	cpv1beta "antrea.io/antrea/pkg/apis/controlplane/v1beta2"
 	"antrea.io/antrea/pkg/controller/networkpolicy/store"
 	antreatypes "antrea.io/antrea/pkg/controller/types"
@@ -127,13 +128,11 @@ func (eq *endpointQuerier) QueryNetworkPolicies(namespace string, podName string
 			return nil, err
 		}
 		for _, policy := range policies {
-			egressIndex := 0
-			ingressIndex := 0
+			egressIndex, ingressIndex := 0, 0
 			for _, rule := range policy.(*antreatypes.NetworkPolicy).Rules {
 				for _, addressGroupTrial := range rule.To.AddressGroups {
 					if addressGroupTrial == string(addressGroup.(*antreatypes.AddressGroup).UID) {
 						egress = append(egress, &ruleTemp{policy: policy.(*antreatypes.NetworkPolicy), index: egressIndex})
-						egressIndex++
 						// an AddressGroup can only be referenced in a rule once
 						break
 					}
@@ -141,11 +140,18 @@ func (eq *endpointQuerier) QueryNetworkPolicies(namespace string, podName string
 				for _, addressGroupTrial := range rule.From.AddressGroups {
 					if addressGroupTrial == string(addressGroup.(*antreatypes.AddressGroup).UID) {
 						ingress = append(ingress, &ruleTemp{policy: policy.(*antreatypes.NetworkPolicy), index: ingressIndex})
-						ingressIndex++
 						// an AddressGroup can only be referenced in a rule once
 						break
 					}
 				}
+				// IngressIndex/egressIndex indicates the rule indexes among this policy's original ingress/egress
+				// rules. The calculation accounts for policy rules not referencing this pod, and guarantees that
+				// users can reference the rules from configuration without accessing the internal policies.
+				if rule.Direction == controlplane.DirectionIn {
+					ingressIndex++
+				} else {
+					egressIndex++
+				}
 			}
 		}
 	}
diff --git a/pkg/controller/networkpolicy/endpoint_querier_test.go b/pkg/controller/networkpolicy/endpoint_querier_test.go
index 4f41120813d..c9a0a0a80dc 100644
--- a/pkg/controller/networkpolicy/endpoint_querier_test.go
+++ b/pkg/controller/networkpolicy/endpoint_querier_test.go
@@ -98,8 +98,7 @@ var policies = []*networkingv1.NetworkPolicy{
 					From: []networkingv1.NetworkPolicyPeer{
 						{
 							PodSelector: &metav1.LabelSelector{
-								MatchLabels:      map[string]string{"foo": "bar"},
-								MatchExpressions: nil,
+								MatchLabels: map[string]string{"foo": "bar"},
 							},
 						},
 					},
@@ -110,14 +109,14 @@ var policies = []*networkingv1.NetworkPolicy{
 					To: []networkingv1.NetworkPolicyPeer{
 						{
 							PodSelector: &metav1.LabelSelector{
-								MatchLabels:      map[string]string{"foo": "bar"},
-								MatchExpressions: nil,
+								MatchLabels: map[string]string{"foo": "bar"},
 							},
 						},
 					},
 				},
 			},
 			PolicyTypes: []networkingv1.PolicyType{
+				networkingv1.PolicyTypeIngress,
 				networkingv1.PolicyTypeEgress,
 			},
 		},
@@ -137,6 +136,38 @@ var policies = []*networkingv1.NetworkPolicy{
 			},
 		},
 	},
+	{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-multiple-ingress-rules",
+			Namespace: "testNamespace",
+			UID:       types.UID("uid-3"),
+		},
+		Spec: networkingv1.NetworkPolicySpec{
+			PodSelector: metav1.LabelSelector{
+				MatchLabels: map[string]string{"foo": "bar"},
+			},
+			Ingress: []networkingv1.NetworkPolicyIngressRule{
+				{
+					From: []networkingv1.NetworkPolicyPeer{
+						{
+							PodSelector: &metav1.LabelSelector{
+								MatchLabels: map[string]string{"foo": "baz"},
+							},
+						},
+					},
+				},
+				{
+					From: []networkingv1.NetworkPolicyPeer{
+						{
+							PodSelector: &metav1.LabelSelector{
+								MatchLabels: map[string]string{"foo": "bar"},
+							},
+						},
+					},
+				},
+			},
+		},
+	},
 }
 
 var namespaces = []*corev1.Namespace{
@@ -184,6 +215,7 @@ func makeControllerAndEndpointQuerier(objects ...runtime.Object) *endpointQuerie
 func TestEndpointQuery(t *testing.T) {
 	policyRef0 := PolicyRef{policies[0].Namespace, policies[0].Name, policies[0].UID}
 	policyRef1 := PolicyRef{policies[1].Namespace, policies[1].Name, policies[1].UID}
+	policyRef2 := PolicyRef{policies[2].Namespace, policies[2].Name, policies[2].UID}
 
 	testCases := []struct {
 		name             string
@@ -251,6 +283,26 @@ func TestEndpointQuery(t *testing.T) {
 				},
 			},
 		},
+		{
+			"MultipleRule", // Pod is selected by policy with multiple rules
+			[]runtime.Object{namespaces[0], pods[0], policies[2]},
+			"testNamespace",
+			"podA",
+			&EndpointQueryResponse{
+				[]Endpoint{
+					{
+						Namespace: "testNamespace",
+						Name:      "podA",
+						Policies: []Policy{
+							{policyRef2},
+						},
+						Rules: []Rule{
+							{policyRef2, v1beta2.DirectionIn, 1},
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, tc := range testCases {