From 20263a07a68c770fe723f2406c9e0231b4e2615d Mon Sep 17 00:00:00 2001
From: Amim Knabben <amim.knabben@gmail.com>
Date: Tue, 23 Feb 2021 22:21:43 -0500
Subject: [PATCH] Added protocol specific netpol

---
 test/e2e/network/netpol/network_policy.go | 14 ++++++++++++++
 test/e2e/network/netpol/policies.go       | 22 ++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/test/e2e/network/netpol/network_policy.go b/test/e2e/network/netpol/network_policy.go
index 65c4e27fa97..3f8ee9ecbc0 100644
--- a/test/e2e/network/netpol/network_policy.go
+++ b/test/e2e/network/netpol/network_policy.go
@@ -778,6 +778,20 @@ var _ = SIGDescribeCopy("Netpol [LinuxOnly]", func() {
 			ValidateOrFail(k8s, model, &TestCase{ToPort: 80, Protocol: v1.ProtocolTCP, Reachability: reachability})
 		})
 
+		ginkgo.It("should enforce ingress policy allowing any port traffic to a server on a specific protocol [Feature:NetworkPolicy] [Feature:UDP]", func() {
+			nsX, _, _, model, k8s := getK8SModel(f)
+
+			policy := GetAllowIngressByProtocol("allow-ingress-by-proto", map[string]string{"pod": "a"}, &protocolTCP)
+			CreatePolicy(k8s, policy, nsX)
+
+			reachabilityTCP := NewReachability(model.AllPods(), true)
+			ValidateOrFail(k8s, model, &TestCase{ToPort: 80, Protocol: v1.ProtocolTCP, Reachability: reachabilityTCP})
+
+			reachabilityUDP := NewReachability(model.AllPods(), true)
+			reachabilityUDP.ExpectPeer(&Peer{}, &Peer{Namespace: nsX, Pod: "a"}, false)
+			ValidateOrFail(k8s, model, &TestCase{ToPort: 80, Protocol: v1.ProtocolUDP, Reachability: reachabilityUDP})
+		})
+
 		ginkgo.It("should enforce multiple ingress policies with ingress allow-all policy taking precedence [Feature:NetworkPolicy]", func() {
 			nsX, _, _, model, k8s := getK8SModel(f)
 			policyAllowOnlyPort80 := GetAllowIngressByPort("allow-ingress-port-80", &intstr.IntOrString{Type: intstr.Int, IntVal: 80})
diff --git a/test/e2e/network/netpol/policies.go b/test/e2e/network/netpol/policies.go
index 22b2784f1b8..f08fb84fdf1 100644
--- a/test/e2e/network/netpol/policies.go
+++ b/test/e2e/network/netpol/policies.go
@@ -322,6 +322,28 @@ func GetAllowIngressByNamespaceAndPort(name string, targetLabels map[string]stri
 	return policy
 }
 
+// GetAllowIngressByProtocol allows ingress for any ports on a specific protocol.
+func GetAllowIngressByProtocol(name string, targetLabels map[string]string, protocol *v1.Protocol) *networkingv1.NetworkPolicy {
+	policy := &networkingv1.NetworkPolicy{
+		ObjectMeta: metav1.ObjectMeta{
+			Name: name,
+		},
+		Spec: networkingv1.NetworkPolicySpec{
+			PodSelector: metav1.LabelSelector{
+				MatchLabels: targetLabels,
+			},
+			Ingress: []networkingv1.NetworkPolicyIngressRule{{
+				Ports: []networkingv1.NetworkPolicyPort{
+					{
+						Protocol: protocol,
+					},
+				},
+			}},
+		},
+	}
+	return policy
+}
+
 // GetAllowIngressByNamespaceOrPod allows ingress for pods with matching namespace OR pod labels
 func GetAllowIngressByNamespaceOrPod(name string, targetLabels map[string]string, peerNamespaceSelector *metav1.LabelSelector, peerPodSelector *metav1.LabelSelector) *networkingv1.NetworkPolicy {
 	policy := &networkingv1.NetworkPolicy{