From 8f221e6eab81b1fdf8e45225db0b7d3da65f4820 Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Tue, 26 Apr 2016 10:27:11 +0200
Subject: [PATCH] improve error handling in session

---
 errorcodes/error_codes.go |  3 +--
 server_test.go            |  1 +
 session.go                | 23 ++++++++++++-----------
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/errorcodes/error_codes.go b/errorcodes/error_codes.go
index ac0233009..30a8b2e68 100644
--- a/errorcodes/error_codes.go
+++ b/errorcodes/error_codes.go
@@ -19,8 +19,7 @@ const (
 	QUIC_ENCRYPTION_FAILURE            = protocol.ErrorCode(13)
 	QUIC_PACKET_TOO_LARGE              = protocol.ErrorCode(14)
 	// QUIC_PACKET_FOR_NONEXISTENT_STREAM= Data was sent for a stream which did not exist.
-	// QUIC_CLIENT_GOING_AWAY= The client is going away (browser close, etc.)
-	// QUIC_SERVER_GOING_AWAY= The server is going away (restart etc.)
+	QUIC_PEER_GOING_AWAY       = protocol.ErrorCode(16)
 	QUIC_INVALID_STREAM_ID     = protocol.ErrorCode(17)
 	QUIC_TOO_MANY_OPEN_STREAMS = protocol.ErrorCode(18)
 	// QUIC_CONNECTION_TIMED_OUT= We hit our pre-negotiated (or default) timeout
diff --git a/server_test.go b/server_test.go
index 91e3fd64a..4ab2e43f1 100644
--- a/server_test.go
+++ b/server_test.go
@@ -114,6 +114,7 @@ var _ = Describe("Server", func() {
 			n, _, err2 := conn.ReadFromUDP(data)
 			Expect(err2).ToNot(HaveOccurred())
 			Expect(n).ToNot(BeZero())
+			time.Sleep(10 * time.Millisecond)
 			err2 = server.Close()
 			Expect(err2).ToNot(HaveOccurred())
 		}()
diff --git a/session.go b/session.go
index 81308b0d7..d3b1719b4 100644
--- a/session.go
+++ b/session.go
@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/lucas-clemente/quic-go/ackhandler"
+	"github.com/lucas-clemente/quic-go/errorcodes"
 	"github.com/lucas-clemente/quic-go/frames"
 	"github.com/lucas-clemente/quic-go/handshake"
 	"github.com/lucas-clemente/quic-go/protocol"
@@ -76,15 +77,12 @@ func (s *Session) Run() {
 			err = s.sendPacket()
 		}
 
-		if err == ackhandler.ErrEntropy {
-			// ToDo: use appropriate QuicError here
-			fmt.Println("Session: Received ACK with incorrect entropy. Closing connection.")
-			s.Close(err)
-			break
-		}
-
-		if err != nil && err != ackhandler.ErrDuplicateOrOutOfOrderAck {
-			fmt.Printf("Error in session: %s\n", err.Error())
+		if err != nil {
+			if err == ackhandler.ErrDuplicateOrOutOfOrderAck {
+				fmt.Printf("Ignoring error in session: %s\n", err.Error())
+			} else {
+				s.Close(err)
+			}
 		}
 
 		s.garbageCollectStreams()
@@ -99,7 +97,6 @@ func (s *Session) handlePacket(addr *net.UDPAddr, publicHeader *PublicHeader, r
 
 	packet, err := s.unpacker.Unpack(publicHeader.Raw, publicHeader, r)
 	if err != nil {
-		s.Close(err)
 		return err
 	}
 
@@ -126,7 +123,6 @@ func (s *Session) handlePacket(addr *net.UDPAddr, publicHeader *PublicHeader, r
 			panic("unexpected frame type")
 		}
 		if err != nil {
-			s.Close(err)
 			return err
 		}
 	}
@@ -166,12 +162,17 @@ func (s *Session) handleStreamFrame(frame *frames.StreamFrame) error {
 
 // Close the connection by sending a ConnectionClose frame
 func (s *Session) Close(e error) error {
+	if e == nil {
+		e = protocol.NewQuicError(errorcodes.QUIC_PEER_GOING_AWAY, "peer going away")
+	}
+	fmt.Printf("Closing session with error: %s\n", e.Error())
 	errorCode := protocol.ErrorCode(1)
 	reasonPhrase := e.Error()
 	quicError, ok := e.(*protocol.QuicError)
 	if ok {
 		errorCode = quicError.ErrorCode
 	}
+	// TODO: Don't queue, but send immediately
 	return s.QueueFrame(&frames.ConnectionCloseFrame{
 		ErrorCode:    errorCode,
 		ReasonPhrase: reasonPhrase,