add tests for the Client

client.go

@@ -93,6 +93,8 @@ func (c *Client) Listen() error {
 		err = c.handlePacket(data)
 		if err != nil {
 			utils.Errorf("error handling packet: %s", err.Error())
+			c.session.Close(err)
+			return err
 		}
 	}
 }

client_test.go

@@ -13,9 +13,12 @@ import (
 
 var _ = Describe("Client", func() {
 	var client *Client
+	var session *mockSession
 
 	BeforeEach(func() {
 		client = &Client{}
+		session = &mockSession{}
+		client.session = session
 	})
 
 	It("errors on invalid public header", func() {
@@ -33,9 +36,79 @@ var _ = Describe("Client", func() {
 		Expect(err).ToNot(HaveOccurred())
 		client.conn, err = net.ListenUDP("udp", addr)
 		Expect(err).ToNot(HaveOccurred())
-		client.session = &mockSession{}
 		err = client.Close()
 		Expect(err).ToNot(HaveOccurred())
-		Expect(client.session.(*mockSession).closed).To(BeTrue())
+		Expect(session.closed).To(BeTrue())
+		Expect(session.closeReason).To(BeNil())
+	})
+
+	Context("handling packets", func() {
+		It("errors on too large packets", func() {
+			err := client.handlePacket(bytes.Repeat([]byte{'f'}, int(protocol.MaxPacketSize+1)))
+			Expect(err).To(MatchError(qerr.PacketTooLarge))
+		})
+
+		It("handles packets", func(done Done) {
+			var err error
+			client.addr, err = net.ResolveUDPAddr("udp", "127.0.0.1:0")
+			Expect(err).ToNot(HaveOccurred())
+			client.conn, err = net.ListenUDP("udp", client.addr)
+			Expect(err).NotTo(HaveOccurred())
+			serverConn, err := net.DialUDP("udp", nil, client.conn.LocalAddr().(*net.UDPAddr))
+			Expect(err).NotTo(HaveOccurred())
+
+			go func() {
+				defer GinkgoRecover()
+				listenErr := client.Listen()
+				Expect(listenErr).ToNot(HaveOccurred())
+				close(done)
+			}()
+
+			Expect(session.packetCount).To(BeZero())
+			ph := PublicHeader{
+				PacketNumber:    1,
+				PacketNumberLen: protocol.PacketNumberLen2,
+				ConnectionID:    0x1337,
+			}
+			b := &bytes.Buffer{}
+			err = ph.Write(b, protocol.Version36, protocol.PerspectiveServer)
+			Expect(err).ToNot(HaveOccurred())
+			_, err = serverConn.Write(b.Bytes())
+			Expect(err).ToNot(HaveOccurred())
+
+			Eventually(func() int { return session.packetCount }).Should(Equal(1))
+			Expect(session.closed).To(BeFalse())
+
+			err = client.Close()
+			Expect(err).ToNot(HaveOccurred())
+		})
+
+		It("closes the session when encountering an error while handling a packet", func(done Done) {
+			var err error
+			client.addr, err = net.ResolveUDPAddr("udp", "127.0.0.1:0")
+			Expect(err).ToNot(HaveOccurred())
+			client.conn, err = net.ListenUDP("udp", client.addr)
+			Expect(err).NotTo(HaveOccurred())
+			serverConn, err := net.DialUDP("udp", nil, client.conn.LocalAddr().(*net.UDPAddr))
+			Expect(err).NotTo(HaveOccurred())
+
+			var listenErr error
+			go func() {
+				defer GinkgoRecover()
+				listenErr = client.Listen()
+				Expect(listenErr).To(HaveOccurred())
+				close(done)
+			}()
+
+			// cause a PacketTooLarge error
+			_, err = serverConn.Write(bytes.Repeat([]byte{'f'}, 100))
+			Expect(err).ToNot(HaveOccurred())
+
+			Eventually(func() bool { return session.closed }).Should(BeTrue())
+			Expect(session.closeReason).To(MatchError(listenErr))
+
+			err = client.Close()
+			Expect(err).ToNot(HaveOccurred())
+		})
 	})
 })

server_test.go

@@ -20,14 +20,19 @@ type mockSession struct {
 	connectionID protocol.ConnectionID
 	packetCount  int
 	closed       bool
+	closeReason  error
 }
 
 func (s *mockSession) handlePacket(*receivedPacket) {
 	s.packetCount++
 }
 
-func (s *mockSession) run()              {}
-func (s *mockSession) Close(error) error { s.closed = true; return nil }
+func (s *mockSession) run() {}
+func (s *mockSession) Close(e error) error {
+	s.closed = true
+	s.closeReason = e
+	return nil
+}
 
 func newMockSession(conn connection, v protocol.VersionNumber, connectionID protocol.ConnectionID, sCfg *handshake.ServerConfig, streamCallback StreamCallback, closeCallback closeCallback) (packetHandler, error) {
 	return &mockSession{