From 67657a37bf49dff60225bdee8eca158d07db4be1 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sat, 14 Jan 2017 18:15:50 +0700
Subject: [PATCH] only change remote address after authenticating a packet

fixes #395
---
 session.go      | 10 +++++---
 session_test.go | 63 +++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 66 insertions(+), 7 deletions(-)

diff --git a/session.go b/session.go
index f3e591e3..74679229 100644
--- a/session.go
+++ b/session.go
@@ -272,10 +272,14 @@ func (s *Session) handlePacketImpl(p *receivedPacket) error {
 		utils.Debugf("<- Reading packet 0x%x (%d bytes) for connection %x @ %s", hdr.PacketNumber, len(data)+len(hdr.Raw), hdr.ConnectionID, time.Now().Format("15:04:05.000"))
 	}
 
-	// TODO: Only do this after authenticating
-	s.conn.setCurrentRemoteAddr(p.remoteAddr)
-
 	packet, err := s.unpacker.Unpack(hdr.Raw, hdr, data)
+	// if the decryption failed, this might be a packet sent by an attacker
+	// don't update the remote address
+	if quicErr, ok := err.(*qerr.QuicError); ok && quicErr.ErrorCode == qerr.DecryptionFailure {
+		return err
+	}
+	// update the remote address, even if unpacking failed for any other reason than a decryption error
+	s.conn.setCurrentRemoteAddr(p.remoteAddr)
 	if err != nil {
 		return err
 	}
diff --git a/session_test.go b/session_test.go
index 84da7b87..6d9fdee4 100644
--- a/session_test.go
+++ b/session_test.go
@@ -25,7 +25,8 @@ import (
 )
 
 type mockConnection struct {
-	written [][]byte
+	remoteAddr net.IP
+	written    [][]byte
 }
 
 func (m *mockConnection) write(p []byte) error {
@@ -35,12 +36,21 @@ func (m *mockConnection) write(p []byte) error {
 	return nil
 }
 
-func (*mockConnection) setCurrentRemoteAddr(addr interface{}) {}
-func (*mockConnection) RemoteAddr() *net.UDPAddr              { return &net.UDPAddr{} }
+func (m *mockConnection) setCurrentRemoteAddr(addr interface{}) {
+	if ip, ok := addr.(net.IP); ok {
+		m.remoteAddr = ip
+	}
+}
+func (*mockConnection) RemoteAddr() *net.UDPAddr { return &net.UDPAddr{} }
 
-type mockUnpacker struct{}
+type mockUnpacker struct {
+	unpackErr error
+}
 
 func (m *mockUnpacker) Unpack(publicHeaderBinary []byte, hdr *PublicHeader, data []byte) (*unpackedPacket, error) {
+	if m.unpackErr != nil {
+		return nil, m.unpackErr
+	}
 	return &unpackedPacket{
 		frames: nil,
 	}, nil
@@ -613,6 +623,51 @@ var _ = Describe("Session", func() {
 			err = session.handlePacketImpl(&receivedPacket{publicHeader: hdr})
 			Expect(err).ToNot(HaveOccurred())
 		})
+
+		Context("updating the remote address", func() {
+			It("sets the remote address", func() {
+				remoteIP := net.IPv4(192, 168, 0, 100)
+				Expect(session.conn.(*mockConnection).remoteAddr).ToNot(Equal(remoteIP))
+				p := receivedPacket{
+					remoteAddr:   remoteIP,
+					publicHeader: &PublicHeader{PacketNumber: 1337},
+				}
+				err := session.handlePacketImpl(&p)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(session.conn.(*mockConnection).remoteAddr).To(Equal(remoteIP))
+			})
+
+			It("doesn't change the remote address if authenticating the packet fails", func() {
+				remoteIP := net.IPv4(192, 168, 0, 100)
+				attackerIP := net.IPv4(192, 168, 0, 102)
+				session.conn.(*mockConnection).remoteAddr = remoteIP
+				// use the real packetUnpacker here, to make sure this test fails if the error code for failed decryption changes
+				session.unpacker = &packetUnpacker{}
+				session.unpacker.(*packetUnpacker).aead = &crypto.NullAEAD{}
+				p := receivedPacket{
+					remoteAddr:   attackerIP,
+					publicHeader: &PublicHeader{PacketNumber: 1337},
+				}
+				err := session.handlePacketImpl(&p)
+				quicErr := err.(*qerr.QuicError)
+				Expect(quicErr.ErrorCode).To(Equal(qerr.DecryptionFailure))
+				Expect(session.conn.(*mockConnection).remoteAddr).To(Equal(remoteIP))
+			})
+
+			It("sets the remote address, if the packet is authenticated, but unpacking fails for another reason", func() {
+				testErr := errors.New("testErr")
+				remoteIP := net.IPv4(192, 168, 0, 100)
+				Expect(session.conn.(*mockConnection).remoteAddr).ToNot(Equal(remoteIP))
+				p := receivedPacket{
+					remoteAddr:   remoteIP,
+					publicHeader: &PublicHeader{PacketNumber: 1337},
+				}
+				session.unpacker.(*mockUnpacker).unpackErr = testErr
+				err := session.handlePacketImpl(&p)
+				Expect(err).To(MatchError(testErr))
+				Expect(session.conn.(*mockConnection).remoteAddr).To(Equal(remoteIP))
+			})
+		})
 	})
 
 	Context("sending packets", func() {