Address comments

Signed-off-by: graysonwu <wgrayson@vmware.com>

go.mod

@@ -4,7 +4,7 @@ go 1.19
 
 require (
 	antrea.io/libOpenflow v0.9.2
-	antrea.io/ofnet v0.6.5
+	antrea.io/ofnet v0.6.9
 	github.com/ClickHouse/clickhouse-go v1.5.4
 	github.com/DATA-DOG/go-sqlmock v1.5.0
 	github.com/Mellanox/sriovnet v1.1.0
@@ -210,5 +210,3 @@ require (
 	sigs.k8s.io/json v0.0.0-20211208200746-9f7c6b3444d2 // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.2.1 // indirect
 )
-
-replace antrea.io/ofnet => github.com/graysonwu/ofnet v0.0.0-20230307193942-e7cecbbf5fc7

go.sum

@@ -1,5 +1,7 @@
 antrea.io/libOpenflow v0.9.2 h1:9W++nzaxxwY4NxyHHow/4bfum2UPIBJKmEOVTAG+x3o=
 antrea.io/libOpenflow v0.9.2/go.mod h1:IM9mUfHh5hUNciRRcWYIaWZTlv1TI6QBEHlml7ALdS4=
+antrea.io/ofnet v0.6.9 h1:ACoDhFhSHfNtuBKffvptspZDwKe+EQ5i35PuDUZ8svk=
+antrea.io/ofnet v0.6.9/go.mod h1:CB/Pkt+U0Yi1sM7DZ7iS215xGL+dhRRAM0EV0LTDLnY=
 bazil.org/fuse v0.0.0-20160811212531-371fbbdaa898/go.mod h1:Xbm+BRKSBEpa4q4hTSxohYNQpsxXPbPry4JJWOB3LB8=
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
@@ -637,8 +639,6 @@ github.com/gorilla/websocket v0.0.0-20170926233335-4201258b820c/go.mod h1:E7qHFY
 github.com/gorilla/websocket v1.4.0/go.mod h1:E7qHFY5m1UJ88s3WnNqhKjPHQ0heANvMoAMk2YaljkQ=
 github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
-github.com/graysonwu/ofnet v0.0.0-20230307193942-e7cecbbf5fc7 h1:iiavyk5AfEZO16zNUvRwy3KXOS0zo7hAmcgikzEV+hg=
-github.com/graysonwu/ofnet v0.0.0-20230307193942-e7cecbbf5fc7/go.mod h1:CB/Pkt+U0Yi1sM7DZ7iS215xGL+dhRRAM0EV0LTDLnY=
 github.com/gregjones/httpcache v0.0.0-20180305231024-9cad4c3443a7/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
 github.com/grpc-ecosystem/go-grpc-middleware v1.0.0/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=
 github.com/grpc-ecosystem/go-grpc-middleware v1.0.1-0.20190118093823-f849b5445de4/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=

pkg/agent/controller/networkpolicy/fqdn.go

@@ -811,6 +811,10 @@ func (f *fqdnController) handlePacketIn(pktIn *ofctrl.PacketIn) error {
 
 // sendDNSPacketout forwards the DNS response packet to the original requesting client.
 func (f *fqdnController) sendDNSPacketout(pktIn *ofctrl.PacketIn) error {
+	ethernetPkt, err := getEthernetPacket(pktIn)
+	if err != nil {
+		return err
+	}
 	inPort := f.gwPort
 	if inPort == 0 {
 		// Use the original in_port number in the packetIn message to avoid an invalid input port number. Note that,
@@ -824,5 +828,5 @@ func (f *fqdnController) sendDNSPacketout(pktIn *ofctrl.PacketIn) error {
 	mutatePacketOut := func(packetOutBuilder binding.PacketOutBuilder) binding.PacketOutBuilder {
 		return packetOutBuilder.AddLoadRegMark(openflow.CustomReasonDNSRegMark)
 	}
-	return f.ofClient.SendEthPacketOut(inPort, 0, pktIn.Data, mutatePacketOut)
+	return f.ofClient.SendEthPacketOut(inPort, 0, ethernetPkt, mutatePacketOut)
 }

pkg/agent/openflow/client.go

@@ -277,7 +277,7 @@ type Client interface {
 		udpData []byte,
 		mutatePacketOut func(builder binding.PacketOutBuilder) binding.PacketOutBuilder) error
 	// SendEthPacketOut sends ethernet packet as a packet-out to OVS.
-	SendEthPacketOut(inPort, outPort uint32, ethPkt ofutil.Message, mutatePacketOut func(builder binding.PacketOutBuilder) binding.PacketOutBuilder) error
+	SendEthPacketOut(inPort, outPort uint32, ethPkt *protocol.Ethernet, mutatePacketOut func(builder binding.PacketOutBuilder) binding.PacketOutBuilder) error
 	// NewDNSPacketInConjunction creates a policyRuleConjunction for the dns response interception flows.
 	NewDNSPacketInConjunction(id uint32) error
 	// AddAddressToDNSConjunction adds addresses to the toAddresses of the dns packetIn conjunction,
@@ -1246,7 +1246,7 @@ func (c *client) SendUDPPacketOut(
 	return c.bridge.SendPacketOut(packetOutObj)
 }
 
-func (c *client) SendEthPacketOut(inPort, outPort uint32, ethPkt ofutil.Message, mutatePacketOut func(builder binding.PacketOutBuilder) binding.PacketOutBuilder) error {
+func (c *client) SendEthPacketOut(inPort, outPort uint32, ethPkt *protocol.Ethernet, mutatePacketOut func(builder binding.PacketOutBuilder) binding.PacketOutBuilder) error {
 	packetOutBuilder := c.bridge.BuildPacketOut()
 	packetOutBuilder = packetOutBuilder.SetInport(inPort)
 	if outPort != 0 {

pkg/ovs/openflow/interfaces.go

@@ -18,6 +18,7 @@ import (
 	"net"
 	"time"
 
+	"antrea.io/libOpenflow/protocol"
 	"antrea.io/libOpenflow/util"
 	"antrea.io/ofnet/ofctrl"
 )
@@ -422,7 +423,7 @@ type PacketOutBuilder interface {
 	AddLoadRegMark(mark *RegMark) PacketOutBuilder
 	AddResubmitAction(inPort *uint16, table *uint8) PacketOutBuilder
 	SetL4Packet(packet util.Message) PacketOutBuilder
-	SetEthPacket(packet util.Message) PacketOutBuilder
+	SetEthPacket(packet *protocol.Ethernet) PacketOutBuilder
 	Done() *ofctrl.PacketOut
 }
 

pkg/ovs/openflow/ofctrl_packetin.go

@@ -19,6 +19,8 @@ import (
 	"errors"
 	"fmt"
 
+	"k8s.io/klog/v2"
+
 	"antrea.io/libOpenflow/protocol"
 	"antrea.io/libOpenflow/util"
 	"antrea.io/ofnet/ofctrl"
@@ -29,7 +31,6 @@ const (
 	icmp6EchoRequestType uint8 = 128
 	// tcpStandardHdrLen is the TCP header length without options.
 	tcpStandardHdrLen uint8 = 5
-	tcpOptionEndKind  uint8 = 0
 )
 
 func GetTCPHeaderData(ipPkt util.Message) (tcpSrcPort, tcpDstPort uint16, tcpSeqNum, tcpAckNum uint32, tcpHdrLen uint8, tcpFlags uint8, tcpWinSize uint16, err error) {
@@ -72,11 +73,16 @@ func GetTCPDNSData(tcpPkt *protocol.TCP) (data []byte, err error) {
 	// the message described by that length field, to the TCP layer at the
 	// same time (e.g., in a single "write" system call) to make it more
 	// likely that all the data will be transmitted in a single TCP segment.
-	dnsDataStartIdx := tcpOptionsLen + 2
-	if int(dnsDataStartIdx) > len(tcpPkt.Data) {
+	if int(tcpOptionsLen+2) > len(tcpPkt.Data) {
 		return nil, fmt.Errorf("no DNS data in TCP data")
 	}
-	return tcpPkt.Data[dnsDataStartIdx:], nil
+	dnsDataLen := binary.BigEndian.Uint16(tcpPkt.Data[tcpOptionsLen : tcpOptionsLen+2])
+	dnsData := tcpPkt.Data[tcpOptionsLen+2:]
+	if int(dnsDataLen) > len(dnsData) {
+		klog.Info("DNS response has been fragmented")
+		return nil, fmt.Errorf("DNS response has been fragmented")
+	}
+	return dnsData, nil
 }
 
 func GetUDPHeaderData(ipPkt util.Message) (udpSrcPort, udpDstPort uint16, err error) {

pkg/ovs/openflow/ofctrl_packetin_test.go

@@ -15,6 +15,7 @@
 package openflow
 
 import (
+	"fmt"
 	"net"
 	"testing"
 
@@ -228,22 +229,46 @@ func TestParsePacketIn(t *testing.T) {
 	}
 }
 
-func TestGetTCPDataWithoutOptions(t *testing.T) {
+func TestGetTCPDNSData(t *testing.T) {
 	type args struct {
 		tcp        protocol.TCP
+		expectErr  error
 		expectData []byte
 	}
 	tests := []struct {
 		name string
 		args args
 	}{
 		{
-			name: "GetTCPDNSData",
+			name: "GetTCPDNSDataNoDNS",
 			args: args{
 				tcp: protocol.TCP{
 					HdrLen: 6,
-					Data:   []byte{1, 2, 3, 4, 0, 0, 5},
+					Data:   []byte{1, 2, 3, 4, 0},
 				},
+				expectErr:  fmt.Errorf("no DNS data in TCP data"),
+				expectData: nil,
+			},
+		},
+		{
+			name: "GetTCPDNSDataFragmented",
+			args: args{
+				tcp: protocol.TCP{
+					HdrLen: 6,
+					Data:   []byte{1, 2, 3, 4, 0, 2, 5},
+				},
+				expectErr:  fmt.Errorf("DNS response has been fragmented"),
+				expectData: nil,
+			},
+		},
+		{
+			name: "GetTCPDNSDataSuccess",
+			args: args{
+				tcp: protocol.TCP{
+					HdrLen: 6,
+					Data:   []byte{1, 2, 3, 4, 0, 1, 5},
+				},
+				expectErr:  nil,
 				expectData: []byte{5},
 			},
 		},
@@ -252,7 +277,7 @@ func TestGetTCPDataWithoutOptions(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			tcp := tt.args.tcp
 			tcpData, err := GetTCPDNSData(&tcp)
-			require.NoError(t, err, "GetTCPDNSData() returned an error")
+			assert.Equal(t, tt.args.expectErr, err)
 			assert.Equal(t, tt.args.expectData, tcpData)
 		})
 	}

pkg/ovs/openflow/ofctrl_packetout.go

@@ -325,8 +325,8 @@ func (b *ofPacketOutBuilder) SetL4Packet(packet util.Message) PacketOutBuilder {
 	return b
 }
 
-func (b *ofPacketOutBuilder) SetEthPacket(packet util.Message) PacketOutBuilder {
-	b.pktOut.EthernetPacket = &packet
+func (b *ofPacketOutBuilder) SetEthPacket(packet *protocol.Ethernet) PacketOutBuilder {
+	b.pktOut.EthernetPacket = packet
 	return b
 }