From b1095d0661ac6e05405d00189e25088e28c0ccda Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 22 Feb 2018 17:23:44 +0800
Subject: [PATCH] properly close the UDP proxy used in the integration tests

---
 integrationtests/tools/proxy/proxy.go      |  7 ++++
 integrationtests/tools/proxy/proxy_test.go | 37 ++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/integrationtests/tools/proxy/proxy.go b/integrationtests/tools/proxy/proxy.go
index 5b1dc21d..0040a2d9 100644
--- a/integrationtests/tools/proxy/proxy.go
+++ b/integrationtests/tools/proxy/proxy.go
@@ -137,6 +137,13 @@ func NewQuicProxy(local string, version protocol.VersionNumber, opts *Opts) (*Qu
 
 // Close stops the UDP Proxy
 func (p *QuicProxy) Close() error {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+	for _, c := range p.clientDict {
+		if err := c.ServerConn.Close(); err != nil {
+			return err
+		}
+	}
 	return p.conn.Close()
 }
 
diff --git a/integrationtests/tools/proxy/proxy_test.go b/integrationtests/tools/proxy/proxy_test.go
index 39065062..d36e1b42 100644
--- a/integrationtests/tools/proxy/proxy_test.go
+++ b/integrationtests/tools/proxy/proxy_test.go
@@ -3,7 +3,9 @@ package quicproxy
 import (
 	"bytes"
 	"net"
+	"runtime/pprof"
 	"strconv"
+	"strings"
 	"sync/atomic"
 	"time"
 
@@ -47,9 +49,16 @@ var _ = Describe("QUIC Proxy", func() {
 		})
 
 		It("stops the UDPProxy", func() {
+			isProxyRunning := func() bool {
+				var b bytes.Buffer
+				pprof.Lookup("goroutine").WriteTo(&b, 1)
+				return strings.Contains(b.String(), "proxy.(*QuicProxy).runProxy")
+			}
+
 			proxy, err := NewQuicProxy("localhost:0", protocol.VersionWhatever, nil)
 			Expect(err).ToNot(HaveOccurred())
 			port := proxy.LocalPort()
+			Expect(isProxyRunning()).To(BeTrue())
 			err = proxy.Close()
 			Expect(err).ToNot(HaveOccurred())
 
@@ -65,6 +74,34 @@ var _ = Describe("QUIC Proxy", func() {
 				ln.Close()
 				return nil
 			}).ShouldNot(HaveOccurred())
+			Eventually(isProxyRunning).Should(BeFalse())
+		})
+
+		It("stops listening for proxied connections", func() {
+			isConnRunning := func() bool {
+				var b bytes.Buffer
+				pprof.Lookup("goroutine").WriteTo(&b, 1)
+				return strings.Contains(b.String(), "proxy.(*QuicProxy).runConnection")
+			}
+
+			serverAddr, err := net.ResolveUDPAddr("udp", "localhost:0")
+			Expect(err).ToNot(HaveOccurred())
+			serverConn, err := net.ListenUDP("udp", serverAddr)
+			Expect(err).ToNot(HaveOccurred())
+			defer serverConn.Close()
+
+			proxy, err := NewQuicProxy("localhost:0", protocol.VersionWhatever, &Opts{RemoteAddr: serverConn.LocalAddr().String()})
+			Expect(err).ToNot(HaveOccurred())
+			Expect(isConnRunning()).To(BeFalse())
+
+			// check that the proxy port is not in use anymore
+			conn, err := net.DialUDP("udp", nil, proxy.LocalAddr().(*net.UDPAddr))
+			Expect(err).ToNot(HaveOccurred())
+			_, err = conn.Write(makePacket(1, []byte("foobar")))
+			Expect(err).ToNot(HaveOccurred())
+			Eventually(isConnRunning).Should(BeTrue())
+			Expect(proxy.Close()).To(Succeed())
+			Eventually(isConnRunning).Should(BeFalse())
 		})
 
 		It("has the correct LocalAddr and LocalPort", func() {