From 480d339f9a3ca7da035e22eb0c1928673465ee3e Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Fri, 29 Apr 2016 16:49:45 +0200
Subject: [PATCH] implement public reset sending in session

fixes #14, #33
---
 packet_unpacker.go |  3 ++-
 session.go         | 20 +++++++++++++++++++-
 session_test.go    | 27 +++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/packet_unpacker.go b/packet_unpacker.go
index d51d2829d..b2de0a246 100644
--- a/packet_unpacker.go
+++ b/packet_unpacker.go
@@ -25,7 +25,8 @@ func (u *packetUnpacker) Unpack(publicHeaderBinary []byte, publicHeader *PublicH
 	ciphertext, _ := ioutil.ReadAll(r)
 	plaintext, err := u.aead.Open(publicHeader.PacketNumber, publicHeaderBinary, ciphertext)
 	if err != nil {
-		return nil, err
+		// Wrap err in quicError so that public reset is sent by session
+		return nil, protocol.NewQuicError(errorcodes.QUIC_DECRYPTION_FAILURE, err.Error())
 	}
 	r = bytes.NewReader(plaintext)
 
diff --git a/session.go b/session.go
index f2497ffc7..d385139a5 100644
--- a/session.go
+++ b/session.go
@@ -51,7 +51,8 @@ type Session struct {
 	closeChan       chan struct{}
 	closed          bool
 
-	// Used to calculate the next packet number from the truncated wire representation
+	// Used to calculate the next packet number from the truncated wire
+	// representation, and sent back in public reset packets
 	lastRcvdPacketNumber protocol.PacketNumber
 
 	rttStats   congestion.RTTStats
@@ -232,6 +233,11 @@ func (s *Session) Close(e error) error {
 		errorCode = quicError.ErrorCode
 	}
 	s.closeStreamsWithError(e)
+
+	if errorCode == errorcodes.QUIC_DECRYPTION_FAILURE {
+		return s.sendPublicReset(s.lastRcvdPacketNumber)
+	}
+
 	packet, err := s.packer.PackPacket(nil, []frames.Frame{
 		&frames.ConnectionCloseFrame{ErrorCode: errorCode, ReasonPhrase: reasonPhrase},
 	}, false)
@@ -339,3 +345,15 @@ func (s *Session) garbageCollectStreams() {
 		}
 	}
 }
+
+func (s *Session) sendPublicReset(rejectedPacketNumber protocol.PacketNumber) error {
+	fmt.Printf("Sending public reset for connection %d, packet number %d\n", s.connectionID, rejectedPacketNumber)
+	packet := &publicResetPacket{
+		connectionID:         s.connectionID,
+		rejectedPacketNumber: rejectedPacketNumber,
+		nonceProof:           0, // TODO: Currently ignored by chrome.
+	}
+	var b bytes.Buffer
+	packet.Write(&b)
+	return s.conn.write(b.Bytes())
+}
diff --git a/session_test.go b/session_test.go
index 4440e0f68..f0ccb195d 100644
--- a/session_test.go
+++ b/session_test.go
@@ -1,6 +1,7 @@
 package quic
 
 import (
+	"bytes"
 	"errors"
 	"io"
 	"os"
@@ -243,6 +244,13 @@ var _ = Describe("Session", func() {
 			Expect(conn.written[0]).To(ContainSubstring(string([]byte{0x4c, 0x2, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0})))
 			Expect(conn.written[0]).To(ContainSubstring(string("foobar")))
 		})
+
+		It("sends public reset", func() {
+			err := session.sendPublicReset(1)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(conn.written).To(HaveLen(1))
+			Expect(conn.written[0]).To(ContainSubstring(string([]byte("PRST"))))
+		})
 	})
 
 	It("closes when crypto stream errors", func() {
@@ -263,4 +271,23 @@ var _ = Describe("Session", func() {
 		_, err = s.Write([]byte{})
 		Expect(err).To(MatchError("CryptoSetup: expected CHLO"))
 	})
+
+	It("sends public reset when receiving invalid message", func() {
+		path := os.Getenv("GOPATH") + "/src/github.com/lucas-clemente/quic-go/example/"
+		signer, err := crypto.NewRSASigner(path+"cert.der", path+"key.der")
+		Expect(err).ToNot(HaveOccurred())
+		scfg := handshake.NewServerConfig(crypto.NewCurve25519KEX(), signer)
+		session = NewSession(conn, 0, 0, scfg, nil).(*Session)
+		hdr := &PublicHeader{
+			PacketNumber: 42,
+		}
+		r := bytes.NewReader([]byte("foo"))
+		err = session.handlePacket(nil, hdr, r)
+		Expect(err).To(HaveOccurred())
+		// Close() should send public reset
+		err = session.Close(err)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(conn.written).To(HaveLen(1))
+		Expect(conn.written[0]).To(ContainSubstring(string([]byte("PRST"))))
+	})
 })