correctly handle connection errors in the client

client.go

@@ -83,7 +83,6 @@ 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)
 
-	// TODO: handle errors
 	go c.Listen()
 
 	c.mutex.Lock()
@@ -111,17 +110,17 @@ func DialAddr(hostname string, config *Config) (Session, error) {
 }
 
 // Listen listens
-func (c *client) Listen() error {
+func (c *client) Listen() {
 	for {
 		data := getPacketBuffer()
 		data = data[:protocol.MaxPacketSize]
 
 		n, addr, err := c.conn.Read(data)
 		if err != nil {
-			if strings.HasSuffix(err.Error(), "use of closed network connection") {
+			if !strings.HasSuffix(err.Error(), "use of closed network connection") {
-				return nil
+				c.session.Close(err)
 			}
-			return err
+			return
 		}
 		data = data[:n]
 
@@ -129,7 +128,7 @@ func (c *client) Listen() error {
 		if err != nil {
 			utils.Errorf("error handling packet: %s", err.Error())
 			c.session.Close(err)
-			return err
+			return
 		}
 	}
 }

client_test.go

@@ -74,9 +74,7 @@ var _ = Describe("Client", func() {
 
 		var stoppedListening bool
 		go func() {
-			defer GinkgoRecover()
+			cl.Listen()
-			err := cl.Listen()
-			Expect(err).ToNot(HaveOccurred())
 			stoppedListening = true
 		}()
 
@@ -131,7 +129,7 @@ var _ = Describe("Client", func() {
 			Expect(sess.packetCount).To(BeZero())
 			var stoppedListening bool
 			go func() {
-				_ = cl.Listen()
+				cl.Listen()
 				// it should continue listening when receiving valid packets
 				stoppedListening = true
 			}()
@@ -142,11 +140,19 @@ var _ = Describe("Client", func() {
 		})
 
 		It("closes the session when encountering an error while handling a packet", func() {
+			Expect(sess.closeReason).ToNot(HaveOccurred())
 			packetConn.dataToRead = bytes.Repeat([]byte{0xff}, 100)
-			listenErr := cl.Listen()
+			cl.Listen()
-			Expect(listenErr).To(HaveOccurred())
 			Expect(sess.closed).To(BeTrue())
-			Expect(sess.closeReason).To(MatchError(listenErr))
+			Expect(sess.closeReason).To(HaveOccurred())
+		})
+
+		It("closes the session when encountering an error while reading from the connection", func() {
+			testErr := errors.New("test error")
+			packetConn.readErr = testErr
+			cl.Listen()
+			Expect(sess.closed).To(BeTrue())
+			Expect(sess.closeReason).To(MatchError(testErr))
 		})
 	})
 

conn_test.go

@@ -14,12 +14,16 @@ type mockPacketConn struct {
 	addr          net.Addr
 	dataToRead    []byte
 	dataReadFrom  net.Addr
+	readErr       error
 	dataWritten   bytes.Buffer
 	dataWrittenTo net.Addr
 	closed        bool
 }
 
 func (c *mockPacketConn) ReadFrom(b []byte) (int, net.Addr, error) {
+	if c.readErr != nil {
+		return 0, nil, c.readErr
+	}
 	if c.dataToRead == nil { // block if there's no data
 		time.Sleep(time.Hour)
 		return 0, nil, io.EOF