diff --git a/mock_packer_test.go b/mock_packer_test.go index 1123390f..d628c2f9 100644 --- a/mock_packer_test.go +++ b/mock_packer_test.go @@ -108,3 +108,13 @@ func (m *MockPacker) PackRetransmission(arg0 *ackhandler.Packet) ([]*packedPacke func (mr *MockPackerMockRecorder) PackRetransmission(arg0 interface{}) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PackRetransmission", reflect.TypeOf((*MockPacker)(nil).PackRetransmission), arg0) } + +// SetToken mocks base method +func (m *MockPacker) SetToken(arg0 []byte) { + m.ctrl.Call(m, "SetToken", arg0) +} + +// SetToken indicates an expected call of SetToken +func (mr *MockPackerMockRecorder) SetToken(arg0 interface{}) *gomock.Call { + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetToken", reflect.TypeOf((*MockPacker)(nil).SetToken), arg0) +} diff --git a/packet_packer.go b/packet_packer.go index 93a3bc3b..8a9b6637 100644 --- a/packet_packer.go +++ b/packet_packer.go @@ -21,6 +21,7 @@ type packer interface { PackConnectionClose(*wire.ConnectionCloseFrame) (*packedPacket, error) HandleTransportParameters(*handshake.TransportParameters) + SetToken([]byte) ChangeDestConnectionID(protocol.ConnectionID) } @@ -123,7 +124,6 @@ func newPacketPacker( handshakeStream cryptoStream, packetNumberManager packetNumberManager, remoteAddr net.Addr, // only used for determining the max packet size - token []byte, cryptoSetup sealingManager, framer frameSource, acks ackFrameSource, @@ -132,7 +132,6 @@ func newPacketPacker( ) *packetPacker { return &packetPacker{ cryptoSetup: cryptoSetup, - token: token, destConnID: destConnID, srcConnID: srcConnID, initialStream: initialStream, @@ -476,6 +475,10 @@ func (p *packetPacker) ChangeDestConnectionID(connID protocol.ConnectionID) { p.destConnID = connID } +func (p *packetPacker) SetToken(token []byte) { + p.token = token +} + func (p *packetPacker) HandleTransportParameters(params *handshake.TransportParameters) { if params.MaxPacketSize != 0 { p.maxPacketSize = utils.MinByteCount(p.maxPacketSize, params.MaxPacketSize) diff --git a/packet_packer_test.go b/packet_packer_test.go index dda434bd..21c9354d 100644 --- a/packet_packer_test.go +++ b/packet_packer_test.go @@ -26,7 +26,6 @@ var _ = Describe("Packet packer", func() { handshakeStream *MockCryptoStream sealingManager *MockSealingManager pnManager *mockackhandler.MockSentPacketHandler - token []byte ) checkLength := func(data []byte) { @@ -57,7 +56,6 @@ var _ = Describe("Packet packer", func() { BeforeEach(func() { rand.Seed(GinkgoRandomSeed()) version := protocol.VersionTLS - token = []byte("initial token") mockSender := NewMockStreamSender(mockCtrl) mockSender.EXPECT().onHasStreamData(gomock.Any()).AnyTimes() initialStream = NewMockCryptoStream(mockCtrl) @@ -74,7 +72,6 @@ var _ = Describe("Packet packer", func() { handshakeStream, pnManager, &net.TCPAddr{}, - token, // token sealingManager, framer, ackFramer, @@ -772,6 +769,8 @@ var _ = Describe("Packet packer", func() { }) It("pads Initial packets to the required minimum packet size", func() { + token := []byte("initial token") + packer.SetToken(token) f := &wire.CryptoFrame{Data: []byte("foobar")} pnManager.EXPECT().PeekPacketNumber().Return(protocol.PacketNumber(0x42), protocol.PacketNumberLen2) pnManager.EXPECT().PopPacketNumber().Return(protocol.PacketNumber(0x42)) @@ -856,7 +855,6 @@ var _ = Describe("Packet packer", func() { packer.perspective = protocol.PerspectiveClient packet, err := packer.PackPacket() Expect(err).ToNot(HaveOccurred()) - Expect(packet.header.Token).To(Equal(token)) Expect(packet.raw).To(HaveLen(protocol.MinInitialPacketSize)) Expect(packet.frames).To(HaveLen(2)) Expect(packet.frames[0]).To(Equal(ack)) @@ -884,6 +882,8 @@ var _ = Describe("Packet packer", func() { }) It("packs a retransmission for an Initial packet", func() { + token := []byte("initial token") + packer.SetToken(token) pnManager.EXPECT().PeekPacketNumber().Return(protocol.PacketNumber(0x42), protocol.PacketNumberLen2) pnManager.EXPECT().PopPacketNumber().Return(protocol.PacketNumber(0x42)) sealingManager.EXPECT().GetSealerWithEncryptionLevel(protocol.EncryptionInitial).Return(sealer, nil) diff --git a/session.go b/session.go index 90a168e1..ab8d4ec5 100644 --- a/session.go +++ b/session.go @@ -202,7 +202,6 @@ var newSession = func( handshakeStream, s.sentPacketHandler, s.RemoteAddr(), - nil, // no token cs, s.framer, s.receivedPacketHandler, @@ -287,13 +286,13 @@ var newClientSession = func( handshakeStream, s.sentPacketHandler, s.RemoteAddr(), - token, cs, s.framer, s.receivedPacketHandler, s.perspective, s.version, ) + s.packer.SetToken(token) return s, s.postSetup() }