move listening from the multiplexer to the packet handler map

packet_handler_map_test.go

@@ -1,18 +1,26 @@
 package quic
 
 import (
+	"bytes"
 	"time"
 
+	"github.com/golang/mock/gomock"
 	"github.com/lucas-clemente/quic-go/internal/protocol"
+	"github.com/lucas-clemente/quic-go/internal/utils"
+	"github.com/lucas-clemente/quic-go/internal/wire"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 )
 
 var _ = Describe("Packet Handler Map", func() {
-	var handler *packetHandlerMap
+	var (
+		handler *packetHandlerMap
+		conn    *mockPacketConn
+	)
 
 	BeforeEach(func() {
-		handler = newPacketHandlerMap().(*packetHandlerMap)
+		conn = newMockPacketConn()
+		handler = newPacketHandlerMap(conn, 5, utils.DefaultLogger, true).(*packetHandlerMap)
 	})
 
 	It("adds and gets", func() {
@@ -53,4 +61,124 @@ var _ = Describe("Packet Handler Map", func() {
 		handler.Add(protocol.ConnectionID{2, 2, 2, 2}, sess2)
 		handler.Close()
 	})
+
+	Context("handling packets", func() {
+		getPacket := func(connID protocol.ConnectionID) []byte {
+			buf := &bytes.Buffer{}
+			err := (&wire.Header{
+				DestConnectionID: connID,
+				PacketNumberLen:  protocol.PacketNumberLen1,
+			}).Write(buf, protocol.PerspectiveServer, protocol.VersionWhatever)
+			Expect(err).ToNot(HaveOccurred())
+			return buf.Bytes()
+		}
+
+		It("handles packets for different packet handlers on the same packet conn", func() {
+			connID1 := protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8}
+			connID2 := protocol.ConnectionID{8, 7, 6, 5, 4, 3, 2, 1}
+			packetHandler1 := NewMockQuicSession(mockCtrl)
+			packetHandler2 := NewMockQuicSession(mockCtrl)
+			handledPacket1 := make(chan struct{})
+			handledPacket2 := make(chan struct{})
+			packetHandler1.EXPECT().handlePacket(gomock.Any()).Do(func(p *receivedPacket) {
+				Expect(p.header.DestConnectionID).To(Equal(connID1))
+				close(handledPacket1)
+			})
+			packetHandler1.EXPECT().GetVersion()
+			packetHandler2.EXPECT().handlePacket(gomock.Any()).Do(func(p *receivedPacket) {
+				Expect(p.header.DestConnectionID).To(Equal(connID2))
+				close(handledPacket2)
+			})
+			packetHandler2.EXPECT().GetVersion()
+			handler.Add(connID1, packetHandler1)
+			handler.Add(connID2, packetHandler2)
+
+			conn.dataToRead <- getPacket(connID1)
+			conn.dataToRead <- getPacket(connID2)
+			Eventually(handledPacket1).Should(BeClosed())
+			Eventually(handledPacket2).Should(BeClosed())
+
+			// makes the listen go routine return
+			packetHandler1.EXPECT().Close().AnyTimes()
+			packetHandler2.EXPECT().Close().AnyTimes()
+			close(conn.dataToRead)
+		})
+
+		It("drops unparseable packets", func() {
+			err := handler.handlePacket(nil, []byte("invalid"))
+			Expect(err).To(HaveOccurred())
+			Expect(err.Error()).To(ContainSubstring("error parsing invariant header:"))
+		})
+
+		It("ignores packets arriving late for closed sessions", func() {
+			handler.deleteClosedSessionsAfter = time.Hour
+			connID := protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8}
+			handler.Add(connID, NewMockQuicSession(mockCtrl))
+			handler.Remove(connID)
+			err := handler.handlePacket(nil, getPacket(connID))
+			Expect(err).ToNot(HaveOccurred())
+		})
+
+		It("drops packets for unknown receivers", func() {
+			connID := protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8}
+			err := handler.handlePacket(nil, getPacket(connID))
+			Expect(err).To(MatchError("received a packet with an unexpected connection ID 0x0102030405060708"))
+		})
+
+		It("errors on packets that are smaller than the Payload Length in the packet header", func() {
+			connID := protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8}
+			packetHandler := NewMockQuicSession(mockCtrl)
+			packetHandler.EXPECT().GetVersion().Return(versionIETFFrames)
+			handler.Add(connID, packetHandler)
+			hdr := &wire.Header{
+				IsLongHeader:     true,
+				Type:             protocol.PacketTypeHandshake,
+				PayloadLen:       1000,
+				DestConnectionID: connID,
+				PacketNumberLen:  protocol.PacketNumberLen1,
+				Version:          versionIETFFrames,
+			}
+			buf := &bytes.Buffer{}
+			Expect(hdr.Write(buf, protocol.PerspectiveServer, versionIETFFrames)).To(Succeed())
+			buf.Write(bytes.Repeat([]byte{0}, 500))
+
+			err := handler.handlePacket(nil, buf.Bytes())
+			Expect(err).To(MatchError("packet payload (500 bytes) is smaller than the expected payload length (1000 bytes)"))
+		})
+
+		It("cuts packets at the Payload Length", func() {
+			connID := protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8}
+			packetHandler := NewMockQuicSession(mockCtrl)
+			packetHandler.EXPECT().GetVersion().Return(versionIETFFrames)
+			handler.Add(connID, packetHandler)
+			packetHandler.EXPECT().handlePacket(gomock.Any()).Do(func(p *receivedPacket) {
+				Expect(p.data).To(HaveLen(456))
+			})
+
+			hdr := &wire.Header{
+				IsLongHeader:     true,
+				Type:             protocol.PacketTypeHandshake,
+				PayloadLen:       456,
+				DestConnectionID: connID,
+				PacketNumberLen:  protocol.PacketNumberLen1,
+				Version:          versionIETFFrames,
+			}
+			buf := &bytes.Buffer{}
+			Expect(hdr.Write(buf, protocol.PerspectiveServer, versionIETFFrames)).To(Succeed())
+			buf.Write(bytes.Repeat([]byte{0}, 500))
+			err := handler.handlePacket(nil, buf.Bytes())
+			Expect(err).ToNot(HaveOccurred())
+		})
+
+		It("closes the packet handlers when reading from the conn fails", func() {
+			done := make(chan struct{})
+			packetHandler := NewMockQuicSession(mockCtrl)
+			packetHandler.EXPECT().Close().Do(func() {
+				close(done)
+			})
+			handler.Add(protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8}, packetHandler)
+			conn.Close()
+			Eventually(done).Should(BeClosed())
+		})
+	})
 })