Repository: kubernetes/kubernetes
File: test/e2e/network/netpol/network_policy.go (Go)

Code before:
		namespaces[0]: {},
		namespaces[1]: {},
		namespaces[2]: {},
	}
	seenPodNamesByNamespace := map[string]map[string]struct{}{
		namespaces[0]: {},
		namespaces[1]: {},
		namespaces[2]: {},
	}

	for _, modelPod := range modelPods {
		namespaceSuffix, podName, found := strings.Cut(modelPod, "/")
		if !found || namespaceSuffix == "" || podName == "" {
			return nil, fmt.Errorf("invalid model pod %q; expected <namespace>/<pod> like x/a", modelPod)
		}

		namespaceName, found := namespaceBySuffix[namespaceSuffix]
		if !found {
			return nil, fmt.Errorf("invalid model pod %q; namespace suffix %q must be one of x, y, z", modelPod, namespaceSuffix)
		}

		if _, seen := seenPodNamesByNamespace[namespaceName][podName]; seen {
			continue
		}
		seenPodNamesByNamespace[namespaceName][podName] = struct{}{}
		podNamesByNamespace[namespaceName] = append(podNamesByNamespace[namespaceName], podName)
	}

	return newModelWithPerNamespacePodNames(namespaces, podNamesByNamespace, ports, protocols), nil
}

// getK8sNamespaces returns the 3 actual namespace names.
func getK8sNamespaces(k8s *kubeManager) (string, string, string) {
	ns := k8s.NamespaceNames()
	return ns[0], ns[1], ns[2]
}

func initializeCluster(ctx context.Context, f *framework.Framework, protocols []v1.Protocol, ports []int32, modelPods ...string) (*kubeManager, error) {
	dnsDomain := framework.TestContext.ClusterDNSDomain
	framework.Logf("dns domain: %s", dnsDomain)

	k8s := newKubeManager(f, dnsDomain)
	rootNs := f.BaseName
	namespaceNames := getNamespaceNames(rootNs)

	model, err := modelFromPodStrings(namespaceNames, modelPods, protocols, ports)
	if err != nil {
		return nil, err
	}

	framework.Logf("initializing cluster: ensuring namespaces, pods and services exist and are ready")

Review comment (suggestion):
Don't do a `map[X]struct{}` by hand, we have `sets.Set` for that.

But in this case, you can just use a set for `podNamesByNamespace` and then you don't need a separate `seenPodNamesByNamespace`. Have `podNamesByNamespace` be a `map[string]sets.Set[string]`, and then here you can just do `podNamesByNamespace[namespaceName].Insert(podName)` and the set will deal with de-dup-ing for you.

Code after:
	var k8s *kubeManager
	ginkgo.BeforeEach(func() {
		// Windows does not support SCTP
		e2eskipper.SkipIfNodeOSDistroIs("windows")
	})

	ginkgo.Context("NetworkPolicy between server and client using SCTP", func() {

		f.It("should support a 'default-deny-ingress' policy", feature.NetworkPolicy, func(ctx context.Context) {
			protocols := []v1.Protocol{protocolSCTP}
			ports := []int32{80}
			// SCTP: namespace X has a default-deny-ingress policy, so we need 2 pods in X
			// to verify X-to-X ingress is blocked, and 2 pods in Y to verify Y-to-X is also
			// blocked while X-to-Y and Y-to-Y traffic remain allowed.
			k8s = initializeResources(ctx, f, protocols, ports, "x/a", "x/b", "y/a", "y/b")
			nsX, _, _ := getK8sNamespaces(k8s)

			policy := GenNetworkPolicyWithNameAndPodSelector("deny-all", metav1.LabelSelector{}, SetSpecIngressRules())
			CreatePolicy(ctx, k8s, policy, nsX)

			reachability := NewReachability(k8s.AllPodStrings(), true)
			reachability.ExpectPeer(&Peer{}, &Peer{Namespace: nsX}, false)

			ValidateOrFail(k8s, &TestCase{ToPort: 80, Protocol: v1.ProtocolSCTP, Reachability: reachability})
		})

		f.It("should enforce policy based on Ports", feature.NetworkPolicy, func(ctx context.Context) {
			protocols := []v1.Protocol{protocolSCTP}
			ports := []int32{81}
			// SCTP port test: namespace X should allow ingress to x/a on port 81 from
			// namespace Y only. We include x/b as a same-namespace source (denied) and z/a
			// as an unrelated namespace (denied). One pod in Y is enough because the rule
			// is namespace-based.
			k8s = initializeResources(ctx, f, protocols, ports, "x/a", "x/b", "y/a", "z/a")
			nsX, nsY, nsZ := getK8sNamespaces(k8s)

			ginkgo.By("Creating a network allowPort81Policy which only allows allow listed namespaces (y) to connect on exactly one port (81)")
			allowedLabels := &metav1.LabelSelector{
				MatchLabels: map[string]string{
					namespaceLabelKey: nsY,
				},
			}
			ingressRule := networkingv1.NetworkPolicyIngressRule{}
			ingressRule.From = append(ingressRule.From, networkingv1.NetworkPolicyPeer{NamespaceSelector: allowedLabels})
			ingressRule.Ports = append(ingressRule.Ports, networkingv1.NetworkPolicyPort{Port: &intstr.IntOrString{IntVal: 81}, Protocol: &protocolSCTP})
			allowPort81Policy := GenNetworkPolicyWithNameAndPodMatchLabel("allow-ingress-on-port-81-ns-x", map[string]string{"pod": "a"}, SetSpecIngressRules(ingressRule))
			CreatePolicy(ctx, k8s, allowPort81Policy, nsX)

			reachability := NewReachability(k8s.AllPodStrings(), true)
			reachability.ExpectPeer(&Peer{Namespace: nsX}, &Peer{Namespace: nsX, Pod: "a"}, false)
			reachability.ExpectPeer(&Peer{Namespace: nsZ}, &Peer{Namespace: nsX, Pod: "a"}, false)