fix flaky proxy tests

fixes #250

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
 					}

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
+	incomingPacketCounter uint64
-	outgoingPacketCounter PacketNumber
+	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)
 			}()
 		}
 	}

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 {
+				for _, conn = range proxy.clientDict {
-					Expect(conn.incomingPacketCounter).To(Equal(PacketNumber(1)))
+					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(string(serverReceivedPackets[0])).To(ContainSubstring("foobar"))
-				Expect(serverReceivedPackets[0]).To(Equal(packetData("foobar")))
+				Expect(string(serverReceivedPackets[1])).To(ContainSubstring("decafbad"))
-				Expect(serverReceivedPackets[1]).To(Equal(packetData("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(string(clientReceivedPackets[0])).To(ContainSubstring("foobar"))
-				Expect(clientReceivedPackets[1]).To(Equal(packetData("decafbad")))
+				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(string(serverReceivedPackets[0])).To(ContainSubstring("foobar1"))
-				Expect(serverReceivedPackets[1]).To(Equal(packetData("foobar3")))
+				Expect(string(serverReceivedPackets[1])).To(ContainSubstring("foobar3"))
-				Expect(serverReceivedPackets[2]).To(Equal(packetData("foobar5")))
+				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(string(clientReceivedPackets[0])).To(ContainSubstring("foobar1"))
-				Expect(clientReceivedPackets[1]).To(Equal(packetData("foobar3")))
+				Expect(string(clientReceivedPackets[1])).To(ContainSubstring("foobar3"))
-				Expect(clientReceivedPackets[2]).To(Equal(packetData("foobar5")))
+				Expect(string(clientReceivedPackets[2])).To(ContainSubstring("foobar5"))
 			})
 		})
 	})