From 7109883a9518e2926fd1665eb30f36eeb49ca9cc Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 24 Feb 2017 09:02:33 +0700
Subject: [PATCH] call the ConnStateCallback in the client for encryption level
 changes

fixes #439
---
 client.go      | 39 ++++++++++++++++++++++++---------------
 client_test.go | 23 +++++++++++++++++++++--
 2 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/client.go b/client.go
index 6cf9e65b..70827e5e 100644
--- a/client.go
+++ b/client.go
@@ -65,6 +65,25 @@ func Dial(pconn net.PacketConn, remoteAddr net.Addr, host string, config *Config
 
 	utils.Infof("Starting new connection to %s (%s), connectionID %x, version %d", hostname, c.conn.RemoteAddr().String(), c.connectionID, c.version)
 
+	return c.establishConnection()
+}
+
+// DialAddr establishes a new QUIC connection to a server
+func DialAddr(hostname string, config *Config) (Session, error) {
+	udpAddr, err := net.ResolveUDPAddr("udp", hostname)
+	if err != nil {
+		return nil, err
+	}
+
+	udpConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
+	if err != nil {
+		return nil, err
+	}
+
+	return Dial(udpConn, udpAddr, hostname, config)
+}
+
+func (c *client) establishConnection() (Session, error) {
 	go c.listen()
 
 	c.mutex.Lock()
@@ -86,21 +105,6 @@ func Dial(pconn net.PacketConn, remoteAddr net.Addr, host string, config *Config
 	return c.session, nil
 }
 
-// DialAddr establishes a new QUIC connection to a server
-func DialAddr(hostname string, config *Config) (Session, error) {
-	udpAddr, err := net.ResolveUDPAddr("udp", hostname)
-	if err != nil {
-		return nil, err
-	}
-
-	udpConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
-	if err != nil {
-		return nil, err
-	}
-
-	return Dial(udpConn, udpAddr, hostname, config)
-}
-
 // Listen listens
 func (c *client) listen() {
 	var err error
@@ -228,6 +232,11 @@ func (c *client) cryptoChangeCallback(_ Session, isForwardSecure bool) {
 		state = ConnStateSecure
 	}
 
+	c.mutex.Lock()
+	c.connState = state
+	c.connStateChangeOrErrCond.Signal()
+	c.mutex.Unlock()
+
 	if c.config.ConnState != nil {
 		go c.config.ConnState(c.session, state)
 	}
diff --git a/client_test.go b/client_test.go
index f6c6a2d9..295bc653 100644
--- a/client_test.go
+++ b/client_test.go
@@ -72,9 +72,8 @@ var _ = Describe("Client", func() {
 			Expect(err).To(MatchError(testErr))
 		})
 
-		// TODO: actually test this
 		// now we're only testing that Dial doesn't return directly after version negotiation
-		It("only returns once a forward-secure connection is established if no ConnState is defined", func() {
+		It("doesn't return after version negotiation is established if no ConnState is defined", func() {
 			packetConn.dataToRead = []byte{0x0, 0x1, 0x0}
 			config.ConnState = nil
 			var dialReturned bool
@@ -86,6 +85,26 @@ var _ = Describe("Client", func() {
 			}()
 			Consistently(func() bool { return dialReturned }).Should(BeFalse())
 		})
+
+		It("only establishes a connection once it is forward-secure if no ConnState is defined", func() {
+			config.ConnState = nil
+			client := &client{conn: &conn{pconn: packetConn, currentAddr: addr}, config: config}
+			client.connStateChangeOrErrCond.L = &client.mutex
+			var returned bool
+			go func() {
+				defer GinkgoRecover()
+				_, err := client.establishConnection()
+				Expect(err).ToNot(HaveOccurred())
+				returned = true
+			}()
+			Consistently(func() bool { return returned }).Should(BeFalse())
+			// switch to a secure connection
+			client.cryptoChangeCallback(nil, false)
+			Consistently(func() bool { return returned }).Should(BeFalse())
+			// switch to a forward-secure connection
+			client.cryptoChangeCallback(nil, true)
+			Eventually(func() bool { return returned }).Should(BeFalse())
+		})
 	})
 
 	It("errors on invalid public header", func() {