fix version negotiation

client_test.go

@@ -45,7 +45,7 @@ var _ = Describe("Client", func() {
 			session:      sess,
 			version:      protocol.SupportedVersions[0],
 			conn:         &conn{pconn: packetConn, currentAddr: addr},
-			errorChan:    make(chan struct{}),
+			versionNegotiationChan: make(chan struct{}),
 		}
 	})
 
@@ -61,9 +61,11 @@ var _ = Describe("Client", func() {
 	})
 
 	Context("Dialing", func() {
+		var acceptClientVersionPacket []byte
+
 		BeforeEach(func() {
 			newClientSession = func(
-				_ connection,
+				conn connection,
 				_ string,
 				_ protocol.VersionNumber,
 				_ protocol.ConnectionID,
@@ -71,11 +73,23 @@ var _ = Describe("Client", func() {
 				_ *Config,
 				_ []protocol.VersionNumber,
 			) (packetHandler, <-chan handshakeEvent, error) {
+				Expect(conn.Write([]byte("fake CHLO"))).To(Succeed())
+				// Expect(err).ToNot(HaveOccurred())
 				return sess, sess.handshakeChan, nil
 			}
+			// accept the QUIC version suggested by the client
+			b := &bytes.Buffer{}
+			err := (&wire.PublicHeader{
+				ConnectionID:    0x1337,
+				PacketNumber:    1,
+				PacketNumberLen: 1,
+			}).Write(b, protocol.VersionWhatever, protocol.PerspectiveServer)
+			Expect(err).ToNot(HaveOccurred())
+			acceptClientVersionPacket = b.Bytes()
 		})
 
 		It("dials non-forward-secure", func(done Done) {
+			packetConn.dataToRead = acceptClientVersionPacket
 			dialed := make(chan struct{})
 			go func() {
 				defer GinkgoRecover()
@@ -91,10 +105,27 @@ var _ = Describe("Client", func() {
 		})
 
 		It("dials a non-forward-secure address", func(done Done) {
+			serverAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
+			Expect(err).ToNot(HaveOccurred())
+			server, err := net.ListenUDP("udp", serverAddr)
+			Expect(err).ToNot(HaveOccurred())
+			defer server.Close()
+			go func() {
+				defer GinkgoRecover()
+				for {
+					_, clientAddr, err := server.ReadFromUDP(make([]byte, 200))
+					if err != nil {
+						return
+					}
+					_, err = server.WriteToUDP(acceptClientVersionPacket, clientAddr)
+					Expect(err).ToNot(HaveOccurred())
+				}
+			}()
+
 			dialed := make(chan struct{})
 			go func() {
 				defer GinkgoRecover()
-				s, err := DialAddrNonFWSecure("localhost:18901", nil, config)
+				s, err := DialAddrNonFWSecure(server.LocalAddr().String(), nil, config)
 				Expect(err).ToNot(HaveOccurred())
 				Expect(s).ToNot(BeNil())
 				close(dialed)
@@ -106,6 +137,7 @@ var _ = Describe("Client", func() {
 		})
 
 		It("Dial only returns after the handshake is complete", func(done Done) {
+			packetConn.dataToRead = acceptClientVersionPacket
 			dialed := make(chan struct{})
 			go func() {
 				defer GinkgoRecover()
@@ -173,12 +205,23 @@ var _ = Describe("Client", func() {
 			close(done)
 		})
 
-		It("returns an error that occurs while waiting for the connection to become secure", func(done Done) {
+		It("returns an error that occurs during version negotiation", func(done Done) {
 			testErr := errors.New("early handshake error")
-			var dialErr error
 			go func() {
 				defer GinkgoRecover()
-				_, dialErr = Dial(packetConn, addr, "quic.clemente.io:1337", nil, config)
+				_, dialErr := Dial(packetConn, addr, "quic.clemente.io:1337", nil, config)
+				Expect(dialErr).To(MatchError(testErr))
+				close(done)
+			}()
+			sess.Close(testErr)
+		})
+
+		It("returns an error that occurs while waiting for the connection to become secure", func(done Done) {
+			testErr := errors.New("early handshake error")
+			packetConn.dataToRead = acceptClientVersionPacket
+			go func() {
+				defer GinkgoRecover()
+				_, dialErr := Dial(packetConn, addr, "quic.clemente.io:1337", nil, config)
 				Expect(dialErr).To(MatchError(testErr))
 				close(done)
 			}()
@@ -187,9 +230,9 @@ var _ = Describe("Client", func() {
 
 		It("returns an error that occurs while waiting for the handshake to complete", func(done Done) {
 			testErr := errors.New("late handshake error")
-			var dialErr error
+			packetConn.dataToRead = acceptClientVersionPacket
 			go func() {
-				_, dialErr = Dial(packetConn, addr, "quic.clemente.io:1337", nil, config)
+				_, dialErr := Dial(packetConn, addr, "quic.clemente.io:1337", nil, config)
 				Expect(dialErr).To(MatchError(testErr))
 				close(done)
 			}()
@@ -254,10 +297,20 @@ var _ = Describe("Client", func() {
 				Expect(err).ToNot(HaveOccurred())
 				cl.handlePacket(nil, b.Bytes())
 				Expect(cl.versionNegotiated).To(BeTrue())
+				Expect(cl.versionNegotiationChan).To(BeClosed())
 			})
 
 			It("changes the version after receiving a version negotiation packet", func() {
 				var negotiatedVersions []protocol.VersionNumber
+				newVersion := protocol.VersionNumber(77)
+				Expect(newVersion).ToNot(Equal(cl.version))
+				Expect(config.Versions).To(ContainElement(newVersion))
+				packetConn.dataToRead = wire.ComposeVersionNegotiation(
+					0x1337,
+					[]protocol.VersionNumber{newVersion},
+				)
+				sessionChan := make(chan *mockSession)
+				handshakeChan := make(chan handshakeEvent)
 				newClientSession = func(
 					_ connection,
 					_ string,
@@ -268,25 +321,38 @@ var _ = Describe("Client", func() {
 					negotiatedVersionsP []protocol.VersionNumber,
 				) (packetHandler, <-chan handshakeEvent, error) {
 					negotiatedVersions = negotiatedVersionsP
-					return &mockSession{
+					// make the server accept the new version
+					if len(negotiatedVersionsP) > 0 {
+						packetConn.dataToRead = acceptClientVersionPacket
+					}
+					sess := &mockSession{
 						connectionID: connectionID,
-					}, nil, nil
+						stopRunLoop:  make(chan struct{}),
+					}
+					sessionChan <- sess
+					return sess, handshakeChan, nil
 				}
 
-				newVersion := protocol.VersionNumber(77)
-				Expect(config.Versions).To(ContainElement(newVersion))
-				Expect(newVersion).ToNot(Equal(cl.version))
-				Expect(sess.packetCount).To(BeZero())
-				cl.connectionID = 0x1337
-				cl.handlePacket(nil, wire.ComposeVersionNegotiation(0x1337, []protocol.VersionNumber{newVersion}))
-				Expect(cl.version).To(Equal(newVersion))
-				Expect(cl.versionNegotiated).To(BeTrue())
-				// it swapped the sessions
-				// Expect(cl.session).ToNot(Equal(sess))
-				Expect(cl.connectionID).ToNot(Equal(0x1337)) // it generated a new connection ID
+				established := make(chan struct{})
+				go func() {
+					defer GinkgoRecover()
+					err := cl.establishSecureConnection()
+					Expect(err).ToNot(HaveOccurred())
+					close(established)
+				}()
+				var firstSession, secondSession *mockSession
+				Eventually(sessionChan).Should(Receive(&firstSession))
+				Eventually(sessionChan).Should(Receive(&secondSession))
 				// it didn't pass the version negoation packet to the old session (since it has no payload)
-				Expect(sess.packetCount).To(BeZero())
+				Expect(firstSession.packetCount).To(BeZero())
+				Eventually(func() bool { return firstSession.closed }).Should(BeTrue())
+				Expect(firstSession.closeReason).To(Equal(errCloseSessionForNewVersion))
+				Consistently(func() bool { return secondSession.closed }).Should(BeFalse())
+				Expect(cl.connectionID).ToNot(BeEquivalentTo(0x1337))
 				Expect(negotiatedVersions).To(Equal([]protocol.VersionNumber{newVersion}))
+
+				handshakeChan <- handshakeEvent{encLevel: protocol.EncryptionSecure}
+				Eventually(established).Should(BeClosed())
 			})
 
 			It("errors if no matching version is found", func() {