From 61454ac85f1209c41ffcc000213a42f3e76346e5 Mon Sep 17 00:00:00 2001 From: Marten Seemann Date: Fri, 5 Aug 2016 00:16:26 +0700 Subject: [PATCH] fix flaky proxy tests fixes #250 --- integrationtests/drop_test.go | 2 +- integrationtests/proxy/udp_proxy.go | 43 ++++++++++------ integrationtests/proxy/udp_proxy_test.go | 63 +++++++++++++++--------- 3 files changed, 69 insertions(+), 39 deletions(-) diff --git a/integrationtests/drop_test.go b/integrationtests/drop_test.go index 8a32dabe..ef4aa2f2 100644 --- a/integrationtests/drop_test.go +++ b/integrationtests/drop_test.go @@ -60,7 +60,7 @@ var _ = Describe("Drop Proxy", func() { Context(fmt.Sprintf("with quic version %d", version), func() { Context("dropping every 4th packet after the crypto handshake", func() { - dropper := func(p proxy.PacketNumber) bool { + dropper := func(p protocol.PacketNumber) bool { if p <= 5 { // don't interfere with the crypto handshake return false } diff --git a/integrationtests/proxy/udp_proxy.go b/integrationtests/proxy/udp_proxy.go index 2d5c80d1..5e99a5d6 100644 --- a/integrationtests/proxy/udp_proxy.go +++ b/integrationtests/proxy/udp_proxy.go @@ -1,10 +1,15 @@ package proxy import ( + "bytes" "net" "strconv" "sync" + "sync/atomic" "time" + + "github.com/lucas-clemente/quic-go" + "github.com/lucas-clemente/quic-go/protocol" ) // Connection is a UDP connection @@ -12,16 +17,12 @@ type connection struct { ClientAddr *net.UDPAddr // Address of the client ServerConn *net.UDPConn // UDP connection to server - incomingPacketCounter PacketNumber - outgoingPacketCounter PacketNumber + incomingPacketCounter uint64 + outgoingPacketCounter uint64 } -// A PacketNumber is a packet number for the UDP Proxy -// note that this does not necessarily correspond to the QUIC packet number -type PacketNumber uint64 - // DropCallback is a callback that determines which packet gets dropped -type DropCallback func(PacketNumber) bool +type DropCallback func(protocol.PacketNumber) bool // UDPProxy is a UDP proxy type UDPProxy struct { @@ -39,7 +40,7 @@ type UDPProxy struct { // NewUDPProxy creates a new UDP proxy func NewUDPProxy(proxyPort int, serverAddress string, serverPort int, dropIncomingPacket, dropOutgoingPacket DropCallback, rttMin time.Duration, rttMax time.Duration) (*UDPProxy, error) { - dontDrop := func(p PacketNumber) bool { + dontDrop := func(p protocol.PacketNumber) bool { return false } @@ -121,13 +122,20 @@ func (p *UDPProxy) runProxy() error { p.mutex.Unlock() } - conn.incomingPacketCounter++ + atomic.AddUint64(&conn.incomingPacketCounter, 1) - if !p.dropIncomingPacket(conn.incomingPacketCounter) { + raw := buffer[0:n] + r := bytes.NewReader(raw) + hdr, err := quic.ParsePublicHeader(r) + if err != nil { + return err + } + + if !p.dropIncomingPacket(hdr.PacketNumber) { // Relay to server go func() { time.Sleep(p.rttGen.getRTT() / 2) - conn.ServerConn.Write(buffer[0:n]) + conn.ServerConn.Write(raw) }() } } @@ -142,13 +150,20 @@ func (p *UDPProxy) runConnection(conn *connection) error { return err } - conn.outgoingPacketCounter++ + raw := buffer[0:n] + r := bytes.NewReader(raw) + hdr, err := quic.ParsePublicHeader(r) + if err != nil { + return err + } - if !p.dropOutgoingPacket(conn.outgoingPacketCounter) { + atomic.AddUint64(&conn.outgoingPacketCounter, 1) + + if !p.dropOutgoingPacket(hdr.PacketNumber) { // Relay it to client go func() { time.Sleep(p.rttGen.getRTT() / 2) - p.proxyConn.WriteToUDP(buffer[0:n], conn.ClientAddr) + p.proxyConn.WriteToUDP(raw, conn.ClientAddr) }() } } diff --git a/integrationtests/proxy/udp_proxy_test.go b/integrationtests/proxy/udp_proxy_test.go index 9db3c317..6a8a38fd 100644 --- a/integrationtests/proxy/udp_proxy_test.go +++ b/integrationtests/proxy/udp_proxy_test.go @@ -1,10 +1,13 @@ package proxy import ( + "bytes" "net" "strconv" "time" + "github.com/lucas-clemente/quic-go" + "github.com/lucas-clemente/quic-go/protocol" . "github.com/onsi/ginkgo" . "github.com/onsi/gomega" ) @@ -14,6 +17,20 @@ type packetData []byte var _ = Describe("UDP Proxy", func() { var serverPort int + makePacket := func(p protocol.PacketNumber, payload []byte) []byte { + b := &bytes.Buffer{} + hdr := quic.PublicHeader{ + PacketNumber: p, + PacketNumberLen: protocol.PacketNumberLen6, + ConnectionID: 1337, + TruncateConnectionID: false, + } + hdr.WritePublicHeader(b, protocol.Version34) + raw := b.Bytes() + raw = append(raw, payload...) + return raw + } + BeforeEach(func() { serverPort = 7331 }) @@ -107,38 +124,36 @@ var _ = Describe("UDP Proxy", func() { }) It("relays packets from the client to the server", func() { - _, err := clientConn.Write([]byte("foobar")) + _, err := clientConn.Write(makePacket(1, []byte("foobar"))) Expect(err).ToNot(HaveOccurred()) Eventually(func() map[string]*connection { return proxy.clientDict }).Should(HaveLen(1)) - var key string var conn *connection - for key, conn = range proxy.clientDict { - Expect(conn.incomingPacketCounter).To(Equal(PacketNumber(1))) + for _, conn = range proxy.clientDict { + Expect(conn.incomingPacketCounter).To(Equal(uint64(1))) } - _, err = clientConn.Write([]byte("decafbad")) + _, err = clientConn.Write(makePacket(2, []byte("decafbad"))) Expect(err).ToNot(HaveOccurred()) Expect(proxy.clientDict).To(HaveLen(1)) Eventually(func() []packetData { return serverReceivedPackets }).Should(HaveLen(2)) - Expect(proxy.clientDict[key].incomingPacketCounter).To(Equal(PacketNumber(2))) - Expect(serverReceivedPackets[0]).To(Equal(packetData("foobar"))) - Expect(serverReceivedPackets[1]).To(Equal(packetData("decafbad"))) + Expect(string(serverReceivedPackets[0])).To(ContainSubstring("foobar")) + Expect(string(serverReceivedPackets[1])).To(ContainSubstring("decafbad")) }) It("relays packets from the server to the client", func() { - _, err := clientConn.Write([]byte("foobar")) + _, err := clientConn.Write(makePacket(1, []byte("foobar"))) Expect(err).ToNot(HaveOccurred()) Eventually(func() map[string]*connection { return proxy.clientDict }).Should(HaveLen(1)) var key string var conn *connection for key, conn = range proxy.clientDict { - Expect(conn.outgoingPacketCounter).To(Equal(PacketNumber(1))) + Expect(conn.outgoingPacketCounter).To(Equal(uint64(1))) } - _, err = clientConn.Write([]byte("decafbad")) + _, err = clientConn.Write(makePacket(2, []byte("decafbad"))) Expect(err).ToNot(HaveOccurred()) Expect(proxy.clientDict).To(HaveLen(1)) - Eventually(func() PacketNumber { return proxy.clientDict[key].outgoingPacketCounter }).Should(Equal(PacketNumber(2))) + Eventually(func() uint64 { return proxy.clientDict[key].outgoingPacketCounter }).Should(Equal(uint64(2))) var clientReceivedPackets []packetData @@ -161,14 +176,14 @@ var _ = Describe("UDP Proxy", func() { Expect(serverReceivedPackets).To(HaveLen(2)) Expect(serverNumPacketsSent).To(Equal(2)) Eventually(func() []packetData { return clientReceivedPackets }).Should(HaveLen(2)) - Expect(clientReceivedPackets[0]).To(Equal(packetData("foobar"))) - Expect(clientReceivedPackets[1]).To(Equal(packetData("decafbad"))) + Expect(string(clientReceivedPackets[0])).To(ContainSubstring("foobar")) + Expect(string(clientReceivedPackets[1])).To(ContainSubstring("decafbad")) }) }) Context("Drop Callbacks", func() { It("drops incoming packets", func() { - dropper := func(p PacketNumber) bool { + dropper := func(p protocol.PacketNumber) bool { return p%2 == 0 } @@ -177,18 +192,18 @@ var _ = Describe("UDP Proxy", func() { Expect(err).ToNot(HaveOccurred()) for i := 1; i <= 6; i++ { - _, err := clientConn.Write([]byte("foobar" + strconv.Itoa(i))) + _, err := clientConn.Write(makePacket(protocol.PacketNumber(i), []byte("foobar"+strconv.Itoa(i)))) Expect(err).ToNot(HaveOccurred()) time.Sleep(time.Millisecond) } Eventually(func() []packetData { return serverReceivedPackets }).Should(HaveLen(3)) - Expect(serverReceivedPackets[0]).To(Equal(packetData("foobar1"))) - Expect(serverReceivedPackets[1]).To(Equal(packetData("foobar3"))) - Expect(serverReceivedPackets[2]).To(Equal(packetData("foobar5"))) + Expect(string(serverReceivedPackets[0])).To(ContainSubstring("foobar1")) + Expect(string(serverReceivedPackets[1])).To(ContainSubstring("foobar3")) + Expect(string(serverReceivedPackets[2])).To(ContainSubstring("foobar5")) }) It("drops outgoing packets", func() { - dropper := func(p PacketNumber) bool { + dropper := func(p protocol.PacketNumber) bool { return p%2 == 0 } @@ -197,7 +212,7 @@ var _ = Describe("UDP Proxy", func() { Expect(err).ToNot(HaveOccurred()) for i := 1; i <= 6; i++ { - _, err := clientConn.Write([]byte("foobar" + strconv.Itoa(i))) + _, err := clientConn.Write(makePacket(protocol.PacketNumber(i), []byte("foobar"+strconv.Itoa(i)))) Expect(err).ToNot(HaveOccurred()) time.Sleep(time.Millisecond) } @@ -220,9 +235,9 @@ var _ = Describe("UDP Proxy", func() { }() Eventually(func() []packetData { return clientReceivedPackets }).Should(HaveLen(3)) - Expect(clientReceivedPackets[0]).To(Equal(packetData("foobar1"))) - Expect(clientReceivedPackets[1]).To(Equal(packetData("foobar3"))) - Expect(clientReceivedPackets[2]).To(Equal(packetData("foobar5"))) + Expect(string(clientReceivedPackets[0])).To(ContainSubstring("foobar1")) + Expect(string(clientReceivedPackets[1])).To(ContainSubstring("foobar3")) + Expect(string(clientReceivedPackets[2])).To(ContainSubstring("foobar5")) }) }) })