Address comments

Signed-off-by: graysonwu <[email protected]>
antrea-io · Aug 17, 2023 · 2d131de · 2d131de
1 parent a321c66
commit 2d131de
Showing 1 changed file with 29 additions and 40 deletions.
diff --git a/pkg/agent/openflow/network_policy.go b/pkg/agent/openflow/network_policy.go
@@ -708,6 +708,7 @@ func (c *client) NewDNSPacketInConjunction(id uint32) error {
 	if err := c.ofEntryOperations.AddAll(conj.actionFlows); err != nil {
 		return fmt.Errorf("error when adding action flows for the DNS conjunction: %w", err)
 	}
+
 	dnsPriority := priorityDNSIntercept
 	dnsCTState := &openflow15.CTStates{
 		// Use ct_state=+trk+rpl as matching condition.
@@ -716,6 +717,8 @@ func (c *client) NewDNSPacketInConjunction(id uint32) error {
 		Data: 0b00101000,
 		Mask: 0b00101000,
 	}
+	dnsPortMatchValue := types.BitRange{Value: uint16(dnsPort)}
+
 	conj.serviceClause = conj.newClause(1, 2, getTableByID(conj.ruleTableID), nil)
 	conj.toClause = conj.newClause(2, 2, getTableByID(conj.ruleTableID), nil)
 	c.featureNetworkPolicy.conjMatchFlowLock.Lock()
@@ -725,57 +728,43 @@ func (c *client) NewDNSPacketInConjunction(id uint32) error {
 		tcpMatch := &conjunctiveMatch{
 			tableID:  conj.serviceClause.ruleTable.GetID(),
 			priority: &dnsPriority,
-		}
-		udpMatch := &conjunctiveMatch{
-			tableID:  conj.serviceClause.ruleTable.GetID(),
-			priority: &dnsPriority,
-		}
-		if proto == binding.ProtocolIP {
-			tcpMatch.matchPairs = []matchPair{
-				{
-					matchKey:   MatchTCPSrcPort,
-					matchValue: types.BitRange{Value: uint16(dnsPort)},
-				},
+			matchPairs: []matchPair{
 				{
 					matchKey:   MatchCTState,
 					matchValue: dnsCTState,
 				},
-			}
-			udpMatch.matchPairs = []matchPair{
-				{
-					matchKey:   MatchUDPSrcPort,
-					matchValue: types.BitRange{Value: uint16(dnsPort)},
-				},
+			},
+		}
+		udpMatch := &conjunctiveMatch{
+			tableID:  conj.serviceClause.ruleTable.GetID(),
+			priority: &dnsPriority,
+			matchPairs: []matchPair{
 				// Add CTState for UDP as well to make sure only solicited DNS responses are sent
 				// to userspace.
 				{
 					matchKey:   MatchCTState,
 					matchValue: dnsCTState,
 				},
-			}
+			},
+		}
+		if proto == binding.ProtocolIP {
+			tcpMatch.matchPairs = append(tcpMatch.matchPairs, matchPair{
+				matchKey:   MatchTCPSrcPort,
+				matchValue: dnsPortMatchValue,
+			})
+			udpMatch.matchPairs = append(udpMatch.matchPairs, matchPair{
+				matchKey:   MatchUDPSrcPort,
+				matchValue: dnsPortMatchValue,
+			})
 		} else if proto == binding.ProtocolIPv6 {
-			tcpMatch.matchPairs = []matchPair{
-				{
-					matchKey:   MatchTCPv6SrcPort,
-					matchValue: types.BitRange{Value: uint16(dnsPort)},
-				},
-				{
-					matchKey:   MatchCTState,
-					matchValue: dnsCTState,
-				},
-			}
-			udpMatch.matchPairs = []matchPair{
-				{
-					matchKey:   MatchUDPv6SrcPort,
-					matchValue: types.BitRange{Value: uint16(dnsPort)},
-				},
-				// Add CTState for UDP as well to make sure only solicited DNS responses are sent
-				// to userspace.
-				{
-					matchKey:   MatchCTState,
-					matchValue: dnsCTState,
-				},
-			}
+			tcpMatch.matchPairs = append(tcpMatch.matchPairs, matchPair{
+				matchKey:   MatchTCPv6SrcPort,
+				matchValue: dnsPortMatchValue,
+			})
+			udpMatch.matchPairs = append(udpMatch.matchPairs, matchPair{
+				matchKey:   MatchUDPv6SrcPort,
+				matchValue: dnsPortMatchValue,
+			})
 		}
 		tcpCtxChange := conj.serviceClause.addConjunctiveMatchFlow(c.featureNetworkPolicy, tcpMatch, false, false)
 		udpCtxChange := conj.serviceClause.addConjunctiveMatchFlow(c.featureNetworkPolicy, udpMatch, false, false)