fix flaky proxy tests

fixes #250
This commit is contained in:
Marten Seemann
2016-08-05 00:16:26 +07:00
parent 99d4e36a56
commit 61454ac85f
3 changed files with 69 additions and 39 deletions

View File

@@ -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
}

View File

@@ -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)
}()
}
}

View File

@@ -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"))
})
})
})