forked from quic-go/quic-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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"))
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user