Repository: quic-go/quic-go
Pull request #4691 (2024-10-10 13:17:30)

http3: send GOAWAY when server is shutting down
Continuation of #4669. Part of #153.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/quic-go/quic-go/pull/4691?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go) Report
Attention: Patch coverage is `59.09091%` with `45 lines` in your changes missing coverage. Please review.
> Project coverage is 84.49%. Comparing base [(`e5693d0`)](https://app.codecov.io/gh/quic-go/quic-go/commit/e5693d0ad7ddc8a3971808379773fd4940260cc0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go) to head [(`016d0e9`)](https://app.codecov.io/gh/quic-go/quic-go/commit/016d0e955bcc76b11b8dc611c1ab90b88a96ea3f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).

| [Files with missing lines](https://app.codecov.io/gh/quic-go/quic-go/pull/4691?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go) | Patch % | Lines |
|---|---|---|
| [http3/server.go](https://app.codecov.io/gh/quic-go/quic-go/pull/4691?src=pr&el=tree&filepath=http3%2Fserver.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go#diff-aHR0cDMvc2VydmVyLmdv) | 57.14% | [33 Missing and 6 partials :warning: ](https://app.codecov.io/gh/quic-go/quic-go/pull/4691?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go) |
| [http3/frames.go](https://app.codecov.io/gh/quic-go/quic-go/pull/4691?src=pr&el=tree&filepath=http3%2Fframes.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go#diff-aHR0cDMvZnJhbWVzLmdv) | 68.42% | [4 Missing and 2 partials :warning: ](https://app.codecov.io/gh/quic-go/quic-go/pull/4691?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #4691      +/-   ##
==========================================
- Coverage   84.79%   84.49%   -0.30%     
==========================================
  Files         150      147       -3     
  Lines       14929    14741     -188     
==========================================
- Hits        12658    12454     -204     
- Misses       1742     1768      +26     
+ Partials      529      519      -10     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/quic-go/quic-go/pull/4691?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).

marten-seemann: This should be ready to go now.

Patch:
diff --git a/http3/frames.go b/http3/frames.go
index 66cba68c..b54afb31 100644
--- a/http3/frames.go
+++ b/http3/frames.go
@@ -66,7 +66,8 @@ func (p *frameParser) ParseNext() (frame, error) {
 			return parseSettingsFrame(p.r, l)
 		case 0x3: // CANCEL_PUSH
 		case 0x5: // PUSH_PROMISE
-		case 0x7: // GOAWAY
+		case 0x7:
+			return parseGoAwayFrame(qr, l)
 		case 0xd: // MAX_PUSH_ID
 		case 0x2, 0x6, 0x8, 0x9:
 			p.conn.CloseWithError(quic.ApplicationErrorCode(ErrCodeFrameUnexpected), "")
@@ -194,3 +195,27 @@ func (f *settingsFrame) Append(b []byte) []byte {
 	}
 	return b
 }
+
+type goAwayFrame struct {
+	StreamID quic.StreamID
+}
+
+func parseGoAwayFrame(r io.ByteReader, l uint64) (*goAwayFrame, error) {
+	frame := &goAwayFrame{}
+	cbr := countingByteReader{ByteReader: r}
+	id, err := quicvarint.Read(&cbr)
+	if err != nil {
+		return nil, err
+	}
+	if cbr.Read != int(l) {
+		return nil, errors.New("GOAWAY frame: inconsistent length")
+	}
+	frame.StreamID = quic.StreamID(id)
+	return frame, nil
+}
+
+func (f *goAwayFrame) Append(b []byte) []byte {
+	b = quicvarint.Append(b, 0x7)
+	b = quicvarint.Append(b, uint64(quicvarint.Len(uint64(f.StreamID))))
+	return quicvarint.Append(b, uint64(f.StreamID))
+}
diff --git a/http3/server.go b/http3/server.go
index dd4970bf..6f8f5e47 100644
--- a/http3/server.go
+++ b/http3/server.go
@@ -13,6 +13,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/quic-go/quic-go"
@@ -45,6 +46,8 @@ const (
 	streamTypeQPACKDecoderStream = 3
 )
 
+const goawayTimeout = 5 * time.Second
+
 // A QUICEarlyListener listens for incoming QUIC connections.
 type QUICEarlyListener interface {
 	Accept(context.Context) (quic.EarlyConnection, error)
@@ -213,7 +216,13 @@ type Server struct {
 	mutex     sync.RWMutex
 	listeners map[*QUICEarlyListener]listenerInfo
 
-	closed bool
+	closed           bool
+	closeCtx         context.Context    // canceled when the server is closed
+	closeCancel      context.CancelFunc // cancels the closeCtx
+	graceCtx         context.Context    // canceled when the server is closed or gracefully closed
+	graceCancel      context.CancelFunc // cancels the graceCtx
+	connCount        atomic.Int64
+	connHandlingDone chan struct{}
 
 	altSvcHeader string
 }
@@ -265,11 +274,32 @@ func (s *Server) Serve(conn net.PacketConn) error {
 	return s.serveListener(ln)
 }
 
+// init initializes the contexts used for shutting down the server.
+// It must be called with the mutex held.
+func (s *Server) init() {
+	if s.closeCtx == nil {
+		s.closeCtx, s.closeCancel = context.WithCancel(context.Background())
+		s.graceCtx, s.graceCancel = context.WithCancel(s.closeCtx)
+	}
+	s.connHandlingDone = make(chan struct{}, 1)
+}
+
+func (s *Server) decreaseConnCount() {
+	if s.connCount.Add(-1) == 0 {
+		close(s.connHandlingDone)
+	}
+}
+
 // ServeQUICConn serves a single QUIC connection.
-// It is the caller's responsibility to close the connection.
-// Specifically, closing the server does not close the connection.
 func (s *Server) ServeQUICConn(conn quic.Connection) error {
-	return s.handleConn(context.Background(), conn)
+	s.mutex.Lock()
+	s.init()
+	s.mutex.Unlock()
+
+	s.connCount.Add(1)
+	defer s.decreaseConnCount()
+
+	return s.handleConn(conn)
 }
 
 // ServeListener serves an existing QUIC listener.
@@ -290,19 +320,19 @@ func (s *Server) ServeListener(ln QUICEarlyListener) error {
 }
 
 func (s *Server) serveListener(ln QUICEarlyListener) error {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
 	for {
-		conn, err := ln.Accept(context.Background())
-		if errors.Is(err, quic.ErrServerClosed) {
+		conn, err := ln.Accept(s.graceCtx)
+		// server closed
+		if errors.Is(err, quic.ErrServerClosed) || s.graceCtx.Err() != nil {
 			return http.ErrServerClosed
 		}
 		if err != nil {
 			return err
 		}
+		s.connCount.Add(1)
 		go func() {
-			if err := s.handleConn(ctx, conn); err != nil {
+			defer s.decreaseConnCount()
+			if err := s.handleConn(conn); err != nil {
 				if s.Logger != nil {
 					s.Logger.Debug("handling connection failed", "error", err)
 				}
@@ -435,6 +465,7 @@ func (s *Server) addListener(l *QUICEarlyListener) error {
 	if s.listeners == nil {
 		s.listeners = make(map[*QUICEarlyListener]listenerInfo)
 	}
+	s.init()
 
 	laddr := (*l).Addr()
 	if port, err := extractPort(laddr.String()); err == nil {
@@ -458,9 +489,12 @@ func (s *Server) removeListener(l *QUICEarlyListener) {
 	s.generateAltSvcHeader()
 }
 
-func (s *Server) handleConn(serverCtx context.Context, conn quic.Connection) error {
-	// send a SETTINGS frame
-	str, err := conn.OpenUniStream()
+// handleConn handles the HTTP/3 exchange on a QUIC connection.
+// It blocks until all HTTP handlers for all streams have returned.
+func (s *Server) handleConn(conn quic.Connection) error {
+	// open the control stream and send a SETTINGS frame, it's also used to send a GOAWAY frame later
+	// when the server is gracefully closed
+	ctrlStr, err := conn.OpenUniStream()
 	if err != nil {
 		return fmt.Errorf("opening the control stream failed: %w", err)
 	}
@@ -471,7 +505,7 @@ func (s *Server) handleConn(serverCtx context.Context, conn quic.Connection) err
 		ExtendedConnect: true,
 		Other:           s.AdditionalSettings,
 	}).Append(b)
-	str.Write(b)
+	ctrlStr.Write(b)
 
 	ctx := conn.Context()
 	ctx = context.WithValue(ctx, ServerContextKey, s)
@@ -494,23 +528,58 @@ func (s *Server) handleConn(serverCtx context.Context, conn quic.Connection) err
 	)
 	go hconn.HandleUnidirectionalStreams(s.UniStreamHijacker)
 
+	var nextStreamID quic.StreamID
+	var wg sync.WaitGroup
+	var handleErr error
 	// Process all requests immediately.
 	// It's the client's responsibility to decide which requests are eligible for 0-RTT.
 	for {
-		str, datagrams, err := hconn.acceptStream(serverCtx)
+		str, datagrams, err := hconn.acceptStream(s.graceCtx)
 		if err != nil {
-			// close the connection if the server was closed
-			if errors.Is(err, context.Canceled) {
+			// server (not gracefully) closed, close the connection immediately
+			if s.closeCtx.Err() != nil {
 				conn.CloseWithError(quic.ApplicationErrorCode(ErrCodeNoError), "")
+				handleErr = http.ErrServerClosed
+				break
 			}
+
+			// gracefully closed, send GOAWAY frame and wait for requests to complete or grace period to end
+			// new requests will be rejected and shouldn't be sent
+			if s.graceCtx.Err() != nil {
+				b = (&goAwayFrame{StreamID: nextStreamID}).Append(b[:0])
+				// set a deadline to send the GOAWAY frame
+				ctrlStr.SetWriteDeadline(time.Now().Add(goawayTimeout))
+				ctrlStr.Write(b)
+
+				select {
+				case <-hconn.Context().Done():
+					// we expect the client to eventually close the connection after receiving the GOAWAY
+				case <-s.closeCtx.Done():
+					// close the connection after graceful period
+					conn.CloseWithError(quic.ApplicationErrorCode(ErrCodeNoError), "")
+				}
+				handleErr = http.ErrServerClosed
+				break
+			}
+
 			var appErr *quic.ApplicationError
-			if errors.As(err, &appErr) && appErr.ErrorCode == quic.ApplicationErrorCode(ErrCodeNoError) {
-				return nil
+			if !errors.As(err, &appErr) || appErr.ErrorCode != quic.ApplicationErrorCode(ErrCodeNoError) {
+				handleErr = fmt.Errorf("accepting stream failed: %w", err)
 			}
-			return fmt.Errorf("accepting stream failed: %w", err)
+			break
 		}
-		go s.handleRequest(hconn, str, datagrams, hconn.decoder)
+
+		nextStreamID = str.StreamID() + 4
+		wg.Add(1)
+		go func() {
+			// handleRequest will return once the request has been handled,
+			// or the underlying connection is closed
+			defer wg.Done()
+			s.handleRequest(hconn, str, datagrams, hconn.decoder)
+		}()
 	}
+	wg.Wait()
+	return handleErr
 }
 
 func (s *Server) maxHeaderBytes() uint64 {
@@ -652,11 +721,17 @@ func (s *Server) handleRequest(conn *connection, str quic.Stream, datagrams *dat
 
 // Close the server immediately, aborting requests and sending CONNECTION_CLOSE frames to connected clients.
 // Close in combination with ListenAndServe() (instead of Serve()) may race if it is called before a UDP socket is established.
+// It is the caller's responsibility to close any connection passed to ServeQUICConn.
 func (s *Server) Close() error {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
 	s.closed = true
+	// server is never used
+	if s.closeCtx == nil {
+		return nil
+	}
+	s.closeCancel()
 
 	var err error
 	for ln := range s.listeners {
@@ -664,14 +739,44 @@ func (s *Server) Close() error {
 			err = cerr
 		}
 	}
+	if s.connCount.Load() == 0 {
+		return err
+	}
+	// wait for all connections to be closed
+	<-s.connHandlingDone
 	return err
 }
 
-// CloseGracefully shuts down the server gracefully. The server sends a GOAWAY frame first, then waits for either timeout to trigger, or for all running requests to complete.
+// CloseGracefully shuts down the server gracefully.
+// The server sends a GOAWAY frame first, then or for all running requests to complete.
 // CloseGracefully in combination with ListenAndServe() (instead of Serve()) may race if it is called before a UDP socket is established.
-func (s *Server) CloseGracefully(timeout time.Duration) error {
-	// TODO: implement
-	return nil
+func (s *Server) CloseGracefully(ctx context.Context) error {
+	s.mutex.Lock()
+	s.closed = true
+	// server is never used
+	if s.closeCtx == nil {
+		s.mutex.Unlock()
+		return nil
+	}
+	s.graceCancel()
+	s.mutex.Unlock()
+
+	if s.connCount.Load() == 0 {
+		return nil
+	}
+	select {
+	case <-s.connHandlingDone: // all connections were closed
+		// When receiving a GOAWAY frame, HTTP/3 clients are expected to close the connection
+		// once all requests were successfully handled...
+		return s.Close()
+	case <-ctx.Done():
+		// ... however, clients handling long-lived requests (and misbehaving clients),
+		// might not do so before the context is cancelled.
+		// In this case, we close the server, which closes all existing connections
+		// (expect those passed to ServeQUICConn).
+		_ = s.Close()
+		return ctx.Err()
+	}
 }
 
 // ErrNoAltSvcPort is the error returned by SetQUICHeaders when no port was found
diff --git a/internal/handshake/interface.go b/internal/handshake/interface.go
index 6e0eec77..c3a59fcd 100644
--- a/internal/handshake/interface.go
+++ b/internal/handshake/interface.go
@@ -83,6 +83,29 @@ const (
 	EventHandshakeComplete
 )
 
+func (k EventKind) String() string {
+	switch k {
+	case EventNoEvent:
+		return "EventNoEvent"
+	case EventWriteInitialData:
+		return "EventWriteInitialData"
+	case EventWriteHandshakeData:
+		return "EventWriteHandshakeData"
+	case EventReceivedReadKeys:
+		return "EventReceivedReadKeys"
+	case EventDiscard0RTTKeys:
+		return "EventDiscard0RTTKeys"
+	case EventReceivedTransportParameters:
+		return "EventReceivedTransportParameters"
+	case EventRestoredTransportParameters:
+		return "EventRestoredTransportParameters"
+	case EventHandshakeComplete:
+		return "EventHandshakeComplete"
+	default:
+		return "Unknown EventKind"
+	}
+}
+
 // Event is a handshake event.
 type Event struct {
 	Kind                EventKind


Test patch:
diff --git a/http3/frames_test.go b/http3/frames_test.go
index fc198496..418876ca 100644
--- a/http3/frames_test.go
+++ b/http3/frames_test.go
@@ -252,6 +252,49 @@ var _ = Describe("Frames", func() {
 		})
 	})
 
+	Context("GOAWAY frames", func() {
+		It("parses", func() {
+			data := quicvarint.Append(nil, 0x7) // type byte
+			data = quicvarint.Append(data, uint64(quicvarint.Len(100)))
+			data = quicvarint.Append(data, 100)
+			fp := frameParser{r: bytes.NewReader(data)}
+			frame, err := fp.ParseNext()
+			Expect(err).ToNot(HaveOccurred())
+			Expect(frame).To(BeAssignableToTypeOf(&goAwayFrame{}))
+			Expect(frame.(*goAwayFrame).StreamID).To(Equal(quic.StreamID(100)))
+		})
+
+		It("errors on inconsistent lengths", func() {
+			data := quicvarint.Append(nil, 0x7) // type byte
+			data = quicvarint.Append(data, uint64(quicvarint.Len(100))+1)
+			data = quicvarint.Append(data, 100)
+			fp := frameParser{r: bytes.NewReader(data)}
+			_, err := fp.ParseNext()
+			Expect(err).To(MatchError("GOAWAY frame: inconsistent length"))
+		})
+
+		It("writes", func() {
+			data := (&goAwayFrame{StreamID: 200}).Append(nil)
+			fp := frameParser{r: bytes.NewReader(data)}
+			frame, err := fp.ParseNext()
+			Expect(err).ToNot(HaveOccurred())
+			Expect(frame).To(BeAssignableToTypeOf(&goAwayFrame{}))
+			Expect(frame.(*goAwayFrame).StreamID).To(Equal(quic.StreamID(200)))
+		})
+
+		It("errors on EOF", func() {
+			data := (&goAwayFrame{StreamID: 1337}).Append(nil)
+			fp := frameParser{r: bytes.NewReader(data)}
+			_, err := fp.ParseNext()
+			Expect(err).ToNot(HaveOccurred())
+			for i := range data {
+				fp := frameParser{r: bytes.NewReader(data[:i])}
+				_, err := fp.ParseNext()
+				Expect(err).To(MatchError(io.EOF))
+			}
+		})
+	})
+
 	Context("hijacking", func() {
 		It("reads a frame without hijacking the stream", func() {
 			buf := bytes.NewBuffer(quicvarint.Append(nil, 1337))
diff --git a/http3/server_test.go b/http3/server_test.go
index 848fe700..49206d15 100644
--- a/http3/server_test.go
+++ b/http3/server_test.go
@@ -75,6 +75,8 @@ var _ = Describe("Server", func() {
 				return context.WithValue(ctx, testConnContextKey("test"), c)
 			},
 		}
+		s.closeCtx, s.closeCancel = context.WithCancel(context.Background())
+		s.graceCtx, s.graceCancel = context.WithCancel(s.closeCtx)
 		origQuicListenAddr = quicListenAddr
 	})
 
@@ -356,7 +358,7 @@ var _ = Describe("Server", func() {
 				})
 				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, id)
 				conn.EXPECT().Context().Return(ctx).AnyTimes()
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(frameTypeChan).Should(Receive(BeEquivalentTo(0x41)))
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
 			})
@@ -384,7 +386,7 @@ var _ = Describe("Server", func() {
 				})
 				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
 				conn.EXPECT().Context().Return(ctx).AnyTimes()
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(frameTypeChan).Should(Receive(BeEquivalentTo(0x41)))
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
 			})
@@ -412,7 +414,7 @@ var _ = Describe("Server", func() {
 				})
 				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
 				conn.EXPECT().Context().Return(ctx).AnyTimes()
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(frameTypeChan).Should(Receive(BeEquivalentTo(0x41)))
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
 			})
@@ -440,7 +442,7 @@ var _ = Describe("Server", func() {
 				})
 				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
 				conn.EXPECT().Context().Return(ctx).AnyTimes()
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(done).Should(BeClosed())
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
 			})
@@ -485,7 +487,7 @@ var _ = Describe("Server", func() {
 				})
 				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, id)
 				conn.EXPECT().Context().Return(ctx).AnyTimes()
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(streamTypeChan).Should(Receive(BeEquivalentTo(0x54)))
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
 			})
@@ -510,7 +512,7 @@ var _ = Describe("Server", func() {
 				})
 				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
 				conn.EXPECT().Context().Return(ctx).AnyTimes()
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(done).Should(BeClosed())
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
 			})
@@ -537,7 +539,7 @@ var _ = Describe("Server", func() {
 				})
 				ctx := context.WithValue(context.Background(), quic.ConnectionTracingKey, quic.ConnectionTracingID(1234))
 				conn.EXPECT().Context().Return(ctx).AnyTimes()
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(streamTypeChan).Should(Receive(BeEquivalentTo(0x54)))
 				time.Sleep(scaleDuration(20 * time.Millisecond)) // don't EXPECT any calls to conn.CloseWithError
 			})
@@ -585,7 +587,7 @@ var _ = Describe("Server", func() {
 				str.EXPECT().CancelRead(quic.StreamErrorCode(ErrCodeNoError))
 				str.EXPECT().Close().Do(func() error { close(done); return nil })
 
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(done).Should(BeClosed())
 				hfs := decodeHeader(responseBuf)
 				Expect(hfs).To(HaveKeyWithValue(":status", []string{"200"}))
@@ -607,7 +609,7 @@ var _ = Describe("Server", func() {
 				var buf bytes.Buffer
 				str.EXPECT().Write(gomock.Any()).DoAndReturn(buf.Write).AnyTimes()
 
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(handlerCalled).Should(BeClosed())
 
 				// The buffer is expected to contain:
@@ -643,7 +645,7 @@ var _ = Describe("Server", func() {
 				str.EXPECT().CancelRead(quic.StreamErrorCode(ErrCodeFrameError))
 				str.EXPECT().CancelWrite(quic.StreamErrorCode(ErrCodeFrameError)).Do(func(quic.StreamErrorCode) { close(done) })
 
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(done).Should(BeClosed())
 			})
 
@@ -659,7 +661,7 @@ var _ = Describe("Server", func() {
 				str.EXPECT().CancelRead(quic.StreamErrorCode(ErrCodeRequestIncomplete))
 				str.EXPECT().CancelWrite(quic.StreamErrorCode(ErrCodeRequestIncomplete)).Do(func(quic.StreamErrorCode) { close(done) })
 
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Consistently(handlerCalled).ShouldNot(BeClosed())
 			})
 
@@ -680,7 +682,7 @@ var _ = Describe("Server", func() {
 					close(done)
 					return nil
 				})
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(done).Should(BeClosed())
 			})
 
@@ -703,7 +705,7 @@ var _ = Describe("Server", func() {
 				str.EXPECT().CancelRead(quic.StreamErrorCode(ErrCodeFrameError))
 				str.EXPECT().CancelWrite(quic.StreamErrorCode(ErrCodeFrameError)).Do(func(quic.StreamErrorCode) { close(done) })
 
-				s.handleConn(context.Background(), conn)
+				s.handleConn(conn)
 				Eventually(done).Should(BeClosed())
 			})
 		})
@@ -1198,7 +1200,7 @@ var _ = Describe("Server", func() {
 	})
 
 	It("closes gracefully", func() {
-		Expect(s.CloseGracefully(0)).To(Succeed())
+		Expect(s.CloseGracefully(context.Background())).To(Succeed())
 	})
 
 	It("errors when listening fails", func() {
diff --git a/integrationtests/self/hotswap_test.go b/integrationtests/self/hotswap_test.go
index 4091cda6..ef33a932 100644
--- a/integrationtests/self/hotswap_test.go
+++ b/integrationtests/self/hotswap_test.go
@@ -46,7 +46,6 @@ type fakeClosingListener struct {
 }
 
 func (ln *fakeClosingListener) Accept(ctx context.Context) (quic.EarlyConnection, error) {
-	Expect(ctx).To(Equal(context.Background()))
 	return ln.listenerWrapper.Accept(ln.ctx)
 }
 
diff --git a/integrationtests/self/http_test.go b/integrationtests/self/http_test.go
index 69309a9c..65d1c847 100644
--- a/integrationtests/self/http_test.go
+++ b/integrationtests/self/http_test.go
@@ -1071,8 +1071,11 @@ var _ = Describe("HTTP tests", func() {
 
 	It("aborts requests on shutdown", func() {
 		mux.HandleFunc("/shutdown", func(w http.ResponseWriter, r *http.Request) {
-			defer GinkgoRecover()
-			Expect(server.Close()).To(Succeed())
+			go func() {
+				defer GinkgoRecover()
+				Expect(server.Close()).To(Succeed())
+			}()
+			time.Sleep(scaleDuration(50 * time.Millisecond)) // make sure the server started shutting down
 		})
 
 		_, err := client.Get(fmt.Sprintf("https://localhost:%d/shutdown", port))
@@ -1081,4 +1084,79 @@ var _ = Describe("HTTP tests", func() {
 		Expect(errors.As(err, &appErr)).To(BeTrue())
 		Expect(appErr.ErrorCode).To(Equal(http3.ErrCodeNoError))
 	})
+
+	It("allows existing requests to complete on graceful shutdown", func() {
+		delay := scaleDuration(100 * time.Millisecond)
+		done := make(chan struct{})
+
+		mux.HandleFunc("/shutdown", func(w http.ResponseWriter, r *http.Request) {
+			go func() {
+				defer GinkgoRecover()
+				defer close(done)
+				Expect(server.CloseGracefully(context.Background())).To(Succeed())
+				fmt.Println("close gracefully done")
+			}()
+			time.Sleep(delay)
+			w.Write([]byte("shutdown"))
+		})
+
+		ctx, cancel := context.WithTimeout(context.Background(), 3*delay)
+		defer cancel()
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("https://localhost:%d/shutdown", port), nil)
+		Expect(err).ToNot(HaveOccurred())
+		resp, err := client.Do(req)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(resp.StatusCode).To(Equal(http.StatusOK))
+		body, err := io.ReadAll(resp.Body)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(body).To(Equal([]byte("shutdown")))
+		// manually close the client, since we don't support
+		client.Transport.(*http3.RoundTripper).Close()
+
+		// make sure that CloseGracefully returned
+		Eventually(done).Should(BeClosed())
+	})
+
+	It("aborts long-lived requests on graceful shutdown", func() {
+		delay := scaleDuration(100 * time.Millisecond)
+		shutdownDone := make(chan struct{})
+		requestChan := make(chan time.Duration, 1)
+
+		mux.HandleFunc("/shutdown", func(w http.ResponseWriter, r *http.Request) {
+			start := time.Now()
+			w.WriteHeader(http.StatusOK)
+			w.(http.Flusher).Flush()
+			go func() {
+				defer GinkgoRecover()
+				ctx, cancel := context.WithTimeout(context.Background(), delay)
+				defer cancel()
+				defer close(shutdownDone)
+				Expect(server.CloseGracefully(ctx)).To(MatchError(context.DeadlineExceeded))
+			}()
+			for t := range time.NewTicker(delay / 10).C {
+				if _, err := w.Write([]byte(t.String())); err != nil {
+					requestChan <- time.Since(start)
+					return
+				}
+			}
+		})
+
+		start := time.Now()
+		resp, err := client.Get(fmt.Sprintf("https://localhost:%d/shutdown", port))
+		Expect(err).ToNot(HaveOccurred())
+		Expect(resp.StatusCode).To(Equal(http.StatusOK))
+		_, err = io.Copy(io.Discard, resp.Body)
+		Expect(err).To(HaveOccurred())
+		var h3Err *http3.Error
+		Expect(errors.As(err, &h3Err)).To(BeTrue())
+		Expect(h3Err.ErrorCode).To(Equal(http3.ErrCodeNoError))
+		took := time.Since(start)
+		Expect(took).To(BeNumerically("~", delay, delay/2))
+		var requestDuration time.Duration
+		Eventually(requestChan).Should(Receive(&requestDuration))
+		Expect(requestDuration).To(BeNumerically("~", delay, delay/2))
+
+		// make sure that CloseGracefully returned
+		Eventually(shutdownDone).Should(BeClosed())
+	})
 })
diff --git a/internal/handshake/aead_test.go b/internal/handshake/aead_test.go
index d2d58a02..693ca1e5 100644
--- a/internal/handshake/aead_test.go
+++ b/internal/handshake/aead_test.go
@@ -1,137 +1,114 @@
 package handshake
 
 import (
-	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/rand"
 	"crypto/tls"
 	"fmt"
+	"testing"
 
 	"github.com/quic-go/quic-go/internal/protocol"
 
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/require"
 )
 
-var _ = Describe("Long Header AEAD", func() {
-	for _, ver := range []protocol.Version{protocol.Version1, protocol.Version2} {
-		v := ver
-
-		Context(fmt.Sprintf("using version %s", v), func() {
-			for i := range cipherSuites {
-				cs := cipherSuites[i]
-
-				Context(fmt.Sprintf("using %s", tls.CipherSuiteName(cs.ID)), func() {
-					getSealerAndOpener := func() (LongHeaderSealer, LongHeaderOpener) {
-						key := make([]byte, 16)
-						hpKey := make([]byte, 16)
-						rand.Read(key)
-						rand.Read(hpKey)
-						block, err := aes.NewCipher(key)
-						Expect(err).ToNot(HaveOccurred())
-						aead, err := cipher.NewGCM(block)
-						Expect(err).ToNot(HaveOccurred())
-
-						return newLongHeaderSealer(&xorNonceAEAD{aead: aead}, newHeaderProtector(cs, hpKey, true, v)),
-							newLongHeaderOpener(&xorNonceAEAD{aead: aead}, newHeaderProtector(cs, hpKey, true, v))
-					}
-
-					Context("message encryption", func() {
-						msg := []byte("Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.")
-						ad := []byte("Donec in velit neque.")
-
-						It("encrypts and decrypts a message", func() {
-							sealer, opener := getSealerAndOpener()
-							encrypted := sealer.Seal(nil, msg, 0x1337, ad)
-							opened, err := opener.Open(nil, encrypted, 0x1337, ad)
-							Expect(err).ToNot(HaveOccurred())
-							Expect(opened).To(Equal(msg))
-						})
-
-						It("fails to open a message if the associated data is not the same", func() {
-							sealer, opener := getSealerAndOpener()
-							encrypted := sealer.Seal(nil, msg, 0x1337, ad)
-							_, err := opener.Open(nil, encrypted, 0x1337, []byte("wrong ad"))
-							Expect(err).To(MatchError(ErrDecryptionFailed))
-						})
-
-						It("fails to open a message if the packet number is not the same", func() {
-							sealer, opener := getSealerAndOpener()
-							encrypted := sealer.Seal(nil, msg, 0x1337, ad)
-							_, err := opener.Open(nil, encrypted, 0x42, ad)
-							Expect(err).To(MatchError(ErrDecryptionFailed))
-						})
-
-						It("decodes the packet number", func() {
-							sealer, opener := getSealerAndOpener()
-							encrypted := sealer.Seal(nil, msg, 0x1337, ad)
-							_, err := opener.Open(nil, encrypted, 0x1337, ad)
-							Expect(err).ToNot(HaveOccurred())
-							Expect(opener.DecodePacketNumber(0x38, protocol.PacketNumberLen1)).To(BeEquivalentTo(0x1338))
-						})
-
-						It("ignores packets it can't decrypt for packet number derivation", func() {
-							sealer, opener := getSealerAndOpener()
-							encrypted := sealer.Seal(nil, msg, 0x1337, ad)
-							_, err := opener.Open(nil, encrypted[:len(encrypted)-1], 0x1337, ad)
-							Expect(err).To(HaveOccurred())
-							Expect(opener.DecodePacketNumber(0x38, protocol.PacketNumberLen1)).To(BeEquivalentTo(0x38))
-						})
-					})
-
-					Context("header encryption", func() {
-						It("encrypts and encrypts the header", func() {
-							sealer, opener := getSealerAndOpener()
-							var lastFourBitsDifferent int
-							for i := 0; i < 100; i++ {
-								sample := make([]byte, 16)
-								rand.Read(sample)
-								header := []byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}
-								sealer.EncryptHeader(sample, &header[0], header[9:13])
-								if header[0]&0xf != 0xb5&0xf {
-									lastFourBitsDifferent++
-								}
-								Expect(header[0] & 0xf0).To(Equal(byte(0xb5 & 0xf0)))
-								Expect(header[1:9]).To(Equal([]byte{1, 2, 3, 4, 5, 6, 7, 8}))
-								Expect(header[9:13]).ToNot(Equal([]byte{0xde, 0xad, 0xbe, 0xef}))
-								opener.DecryptHeader(sample, &header[0], header[9:13])
-								Expect(header).To(Equal([]byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}))
-							}
-							Expect(lastFourBitsDifferent).To(BeNumerically(">", 75))
-						})
-
-						It("encrypts and encrypts the header, for a 0xfff..fff sample", func() {
-							sealer, opener := getSealerAndOpener()
-							var lastFourBitsDifferent int
-							for i := 0; i < 100; i++ {
-								sample := bytes.Repeat([]byte{0xff}, 16)
-								header := []byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}
-								sealer.EncryptHeader(sample, &header[0], header[9:13])
-								if header[0]&0xf != 0xb5&0xf {
-									lastFourBitsDifferent++
-								}
-								Expect(header[0] & 0xf0).To(Equal(byte(0xb5 & 0xf0)))
-								Expect(header[1:9]).To(Equal([]byte{1, 2, 3, 4, 5, 6, 7, 8}))
-								Expect(header[9:13]).ToNot(Equal([]byte{0xde, 0xad, 0xbe, 0xef}))
-								opener.DecryptHeader(sample, &header[0], header[9:13])
-								Expect(header).To(Equal([]byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}))
-							}
-						})
-
-						It("fails to decrypt the header when using a different sample", func() {
-							sealer, opener := getSealerAndOpener()
-							header := []byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}
-							sample := make([]byte, 16)
-							rand.Read(sample)
-							sealer.EncryptHeader(sample, &header[0], header[9:13])
-							rand.Read(sample) // use a different sample
-							opener.DecryptHeader(sample, &header[0], header[9:13])
-							Expect(header).ToNot(Equal([]byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}))
-						})
-					})
+func getSealerAndOpener(t *testing.T, cs *cipherSuite, v protocol.Version) (LongHeaderSealer, LongHeaderOpener) {
+	t.Helper()
+	key := make([]byte, 16)
+	hpKey := make([]byte, 16)
+	rand.Read(key)
+	rand.Read(hpKey)
+	block, err := aes.NewCipher(key)
+	require.NoError(t, err)
+	aead, err := cipher.NewGCM(block)
+	require.NoError(t, err)
+
+	return newLongHeaderSealer(&xorNonceAEAD{aead: aead}, newHeaderProtector(cs, hpKey, true, v)),
+		newLongHeaderOpener(&xorNonceAEAD{aead: aead}, newHeaderProtector(cs, hpKey, true, v))
+}
+
+func TestEncryptAndDecryptMessage(t *testing.T) {
+	for _, v := range []protocol.Version{protocol.Version1, protocol.Version2} {
+		for _, cs := range cipherSuites {
+			t.Run(fmt.Sprintf("QUIC %s/%s", v, tls.CipherSuiteName(cs.ID)), func(t *testing.T) {
+				sealer, opener := getSealerAndOpener(t, cs, v)
+				msg := []byte("Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.")
+				ad := []byte("Donec in velit neque.")
+
+				encrypted := sealer.Seal(nil, msg, 0x1337, ad)
+
+				opened, err := opener.Open(nil, encrypted, 0x1337, ad)
+				require.NoError(t, err)
+				require.Equal(t, msg, opened)
+
+				// incorrect associated data
+				_, err = opener.Open(nil, encrypted, 0x1337, []byte("wrong ad"))
+				require.Equal(t, ErrDecryptionFailed, err)
+
+				// incorrect packet number
+				_, err = opener.Open(nil, encrypted, 0x42, ad)
+				require.Equal(t, ErrDecryptionFailed, err)
+			})
+		}
+	}
+}
+
+func TestDecodePacketNumber(t *testing.T) {
+	msg := []byte("Lorem ipsum dolor sit amet")
+	ad := []byte("Donec in velit neque.")
+
+	sealer, opener := getSealerAndOpener(t, getCipherSuite(tls.TLS_AES_128_GCM_SHA256), protocol.Version1)
+	encrypted := sealer.Seal(nil, msg, 0x1337, ad)
+
+	// can't decode the packet number if encryption failed
+	_, err := opener.Open(nil, encrypted[:len(encrypted)-1], 0x1337, ad)
+	require.Error(t, err)
+	require.Equal(t, protocol.PacketNumber(0x38), opener.DecodePacketNumber(0x38, protocol.PacketNumberLen1))
+
+	_, err = opener.Open(nil, encrypted, 0x1337, ad)
+	require.NoError(t, err)
+	require.Equal(t, protocol.PacketNumber(0x1338), opener.DecodePacketNumber(0x38, protocol.PacketNumberLen1))
+}
+
+func TestEncryptAndDecryptHeader(t *testing.T) {
+	for _, v := range []protocol.Version{protocol.Version1, protocol.Version2} {
+		t.Run("QUIC "+v.String(), func(t *testing.T) {
+			for _, cs := range cipherSuites {
+				t.Run(tls.CipherSuiteName(cs.ID), func(t *testing.T) {
+					testEncryptAndDecryptHeader(t, cs, v)
 				})
 			}
 		})
 	}
-})
+}
+
+func testEncryptAndDecryptHeader(t *testing.T, cs *cipherSuite, v protocol.Version) {
+	sealer, opener := getSealerAndOpener(t, cs, v)
+	var lastFourBitsDifferent int
+
+	for i := 0; i < 100; i++ {
+		sample := make([]byte, 16)
+		rand.Read(sample)
+		header := []byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}
+		sealer.EncryptHeader(sample, &header[0], header[9:13])
+		if header[0]&0xf != 0xb5&0xf {
+			lastFourBitsDifferent++
+		}
+		require.Equal(t, byte(0xb5&0xf0), header[0]&0xf0)
+		require.Equal(t, []byte{1, 2, 3, 4, 5, 6, 7, 8}, header[1:9])
+		require.NotEqual(t, []byte{0xde, 0xad, 0xbe, 0xef}, header[9:13])
+		opener.DecryptHeader(sample, &header[0], header[9:13])
+		require.Equal(t, []byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}, header)
+	}
+	require.Greater(t, lastFourBitsDifferent, 75)
+
+	// decryption failure with different sample
+	header := []byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}
+	sample := make([]byte, 16)
+	rand.Read(sample)
+	sealer.EncryptHeader(sample, &header[0], header[9:13])
+	rand.Read(sample) // use a different sample
+	opener.DecryptHeader(sample, &header[0], header[9:13])
+	require.NotEqual(t, []byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}, header)
+}
diff --git a/internal/handshake/crypto_setup_test.go b/internal/handshake/crypto_setup_test.go
index 18edceb1..64525cb4 100644
--- a/internal/handshake/crypto_setup_test.go
+++ b/internal/handshake/crypto_setup_test.go
@@ -10,6 +10,7 @@ import (
 	"errors"
 	"math/big"
 	"net"
+	"testing"
 	"time"
 
 	"github.com/quic-go/quic-go/internal/protocol"
@@ -18,8 +19,7 @@ import (
 	"github.com/quic-go/quic-go/internal/utils"
 	"github.com/quic-go/quic-go/internal/wire"
 
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/require"
 )
 
 const (
@@ -50,502 +50,544 @@ func (m *mockClientSessionCache) Put(sessionKey string, cs *tls.ClientSessionSta
 	m.cache.Put(sessionKey, cs)
 }
 
-var _ = Describe("Crypto Setup TLS", func() {
-	generateCert := func() tls.Certificate {
-		priv, err := rsa.GenerateKey(rand.Reader, 2048)
-		Expect(err).ToNot(HaveOccurred())
-		tmpl := &x509.Certificate{
-			SerialNumber:          big.NewInt(1),
-			Subject:               pkix.Name{},
-			SignatureAlgorithm:    x509.SHA256WithRSA,
-			NotBefore:             time.Now(),
-			NotAfter:              time.Now().Add(time.Hour), // valid for an hour
-			BasicConstraintsValid: true,
-		}
-		certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, priv.Public(), priv)
-		Expect(err).ToNot(HaveOccurred())
-		return tls.Certificate{
-			PrivateKey:  priv,
-			Certificate: [][]byte{certDER},
-		}
+func getTLSConfigs() (clientConf, serverConf *tls.Config) {
+	clientConf = &tls.Config{
+		ServerName: "localhost",
+		RootCAs:    testdata.GetRootCA(),
+		NextProtos: []string{"crypto-setup"},
 	}
+	serverConf = testdata.GetTLSConfig()
+	serverConf.NextProtos = []string{"crypto-setup"}
+	return clientConf, serverConf
+}
 
-	var clientConf, serverConf *tls.Config
+func TestErrorBeforeClientHelloGeneration(t *testing.T) {
+	tlsConf := testdata.GetTLSConfig()
+	tlsConf.InsecureSkipVerify = true
+	tlsConf.NextProtos = []string{""}
+	cl := NewCryptoSetupClient(
+		protocol.ConnectionID{},
+		&wire.TransportParameters{},
+		tlsConf,
+		false,
+		&utils.RTTStats{},
+		nil,
+		utils.DefaultLogger.WithPrefix("client"),
+		protocol.Version1,
+	)
+
+	var terr *qerr.TransportError
+	err := cl.StartHandshake(context.Background())
+	require.True(t, errors.As(err, &terr))
+	require.Equal(t, uint64(0x100+0x50), uint64(terr.ErrorCode))
+	require.Contains(t, err.Error(), "tls: invalid NextProtos value")
+}
 
-	BeforeEach(func() {
-		serverConf = testdata.GetTLSConfig()
-		serverConf.NextProtos = []string{"crypto-setup"}
-		clientConf = &tls.Config{
-			ServerName: "localhost",
-			RootCAs:    testdata.GetRootCA(),
-			NextProtos: []string{"crypto-setup"},
-		}
-	})
-
-	It("handles qtls errors occurring before during ClientHello generation", func() {
-		tlsConf := testdata.GetTLSConfig()
-		tlsConf.InsecureSkipVerify = true
-		tlsConf.NextProtos = []string{""}
-		cl := NewCryptoSetupClient(
-			protocol.ConnectionID{},
-			&wire.TransportParameters{},
-			tlsConf,
-			false,
-			&utils.RTTStats{},
-			nil,
-			utils.DefaultLogger.WithPrefix("client"),
-			protocol.Version1,
-		)
-
-		var terr *qerr.TransportError
-		err := cl.StartHandshake(context.Background())
-		Expect(errors.As(err, &terr)).To(BeTrue())
-		Expect(terr.ErrorCode).To(BeEquivalentTo(0x100 + 0x50))
-		Expect(err.Error()).To(ContainSubstring("tls: invalid NextProtos value"))
-	})
-
-	It("errors when a message is received at the wrong encryption level", func() {
-		var token protocol.StatelessResetToken
-		server := NewCryptoSetupServer(
-			protocol.ConnectionID{},
-			&net.UDPAddr{IP: net.IPv6loopback, Port: 1234},
-			&net.UDPAddr{IP: net.IPv6loopback, Port: 4321},
-			&wire.TransportParameters{StatelessResetToken: &token},
-			testdata.GetTLSConfig(),
-			false,
-			&utils.RTTStats{},
-			nil,
-			utils.DefaultLogger.WithPrefix("server"),
-			protocol.Version1,
-		)
-
-		Expect(server.StartHandshake(context.Background())).To(Succeed())
-
-		fakeCH := append([]byte{typeClientHello, 0, 0, 6}, []byte("foobar")...)
-		// wrong encryption level
-		err := server.HandleMessage(fakeCH, protocol.EncryptionHandshake)
-		Expect(err).To(HaveOccurred())
-		Expect(err.Error()).To(ContainSubstring("tls: handshake data received at wrong level"))
-	})
-
-	Context("doing the handshake", func() {
-		newRTTStatsWithRTT := func(rtt time.Duration) *utils.RTTStats {
-			rttStats := &utils.RTTStats{}
-			rttStats.UpdateRTT(rtt, 0, time.Now())
-			ExpectWithOffset(1, rttStats.SmoothedRTT()).To(Equal(rtt))
-			return rttStats
-		}
+func TestMessageReceivedAtWrongEncryptionLevel(t *testing.T) {
+	var token protocol.StatelessResetToken
+	server := NewCryptoSetupServer(
+		protocol.ConnectionID{},
+		&net.UDPAddr{IP: net.IPv6loopback, Port: 1234},
+		&net.UDPAddr{IP: net.IPv6loopback, Port: 4321},
+		&wire.TransportParameters{StatelessResetToken: &token},
+		testdata.GetTLSConfig(),
+		false,
+		&utils.RTTStats{},
+		nil,
+		utils.DefaultLogger.WithPrefix("server"),
+		protocol.Version1,
+	)
+
+	require.NoError(t, server.StartHandshake(context.Background()))
+
+	fakeCH := append([]byte{typeClientHello, 0, 0, 6}, []byte("foobar")...)
+	// wrong encryption level
+	err := server.HandleMessage(fakeCH, protocol.EncryptionHandshake)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "tls: handshake data received at wrong level")
+}
 
-		// The clientEvents and serverEvents contain all events that were not processed by the function,
-		// i.e. not EventWriteInitialData, EventWriteHandshakeData, EventHandshakeComplete.
-		handshake := func(client, server CryptoSetup) (clientEvents []Event, clientErr error, serverEvents []Event, serverErr error) {
-			Expect(client.StartHandshake(context.Background())).To(Succeed())
-			Expect(server.StartHandshake(context.Background())).To(Succeed())
-
-			var clientHandshakeComplete, serverHandshakeComplete bool
-
-			for {
-			clientLoop:
-				for {
-					ev := client.NextEvent()
-					switch ev.Kind {
-					case EventNoEvent:
-						break clientLoop
-					case EventWriteInitialData:
-						if err := server.HandleMessage(ev.Data, protocol.EncryptionInitial); err != nil {
-							serverErr = err
-							return
-						}
-					case EventWriteHandshakeData:
-						if err := server.HandleMessage(ev.Data, protocol.EncryptionHandshake); err != nil {
-							serverErr = err
-							return
-						}
-					case EventHandshakeComplete:
-						clientHandshakeComplete = true
-					default:
-						clientEvents = append(clientEvents, ev)
-					}
-				}
+func newRTTStatsWithRTT(t *testing.T, rtt time.Duration) *utils.RTTStats {
+	t.Helper()
+	rttStats := &utils.RTTStats{}
+	rttStats.UpdateRTT(rtt, 0, time.Now())
+	require.Equal(t, rtt, rttStats.SmoothedRTT())
+	return rttStats
+}
 
-			serverLoop:
-				for {
-					ev := server.NextEvent()
-					switch ev.Kind {
-					case EventNoEvent:
-						break serverLoop
-					case EventWriteInitialData:
-						if err := client.HandleMessage(ev.Data, protocol.EncryptionInitial); err != nil {
-							clientErr = err
-							return
-						}
-					case EventWriteHandshakeData:
-						if err := client.HandleMessage(ev.Data, protocol.EncryptionHandshake); err != nil {
-							clientErr = err
-							return
-						}
-					case EventHandshakeComplete:
-						serverHandshakeComplete = true
-						ticket, err := server.GetSessionTicket()
-						Expect(err).ToNot(HaveOccurred())
-						if ticket != nil {
-							Expect(client.HandleMessage(ticket, protocol.Encryption1RTT)).To(Succeed())
-						}
-					default:
-						serverEvents = append(serverEvents, ev)
-					}
+// The clientEvents and serverEvents contain all events that were not processed by the function,
+// i.e. not EventWriteInitialData, EventWriteHandshakeData, EventHandshakeComplete.
+func handshake(t *testing.T, client, server CryptoSetup) (clientEvents []Event, clientErr error, serverEvents []Event, serverErr error) {
+	t.Helper()
+	require.NoError(t, client.StartHandshake(context.Background()))
+	require.NoError(t, server.StartHandshake(context.Background()))
+
+	var clientHandshakeComplete, serverHandshakeComplete bool
+
+	for {
+	clientLoop:
+		for {
+			ev := client.NextEvent()
+			switch ev.Kind {
+			case EventNoEvent:
+				break clientLoop
+			case EventWriteInitialData:
+				serverErr = server.HandleMessage(ev.Data, protocol.EncryptionInitial)
+				if serverErr != nil {
+					return
 				}
-
-				if clientHandshakeComplete && serverHandshakeComplete {
-					break
+			case EventWriteHandshakeData:
+				serverErr = server.HandleMessage(ev.Data, protocol.EncryptionHandshake)
+				if serverErr != nil {
+					return
 				}
+			case EventHandshakeComplete:
+				clientHandshakeComplete = true
+			default:
+				clientEvents = append(clientEvents, ev)
 			}
-			return
 		}
 
-		handshakeWithTLSConf := func(
-			clientConf, serverConf *tls.Config,
-			clientRTTStats, serverRTTStats *utils.RTTStats,
-			clientTransportParameters, serverTransportParameters *wire.TransportParameters,
-			enable0RTT bool,
-		) (CryptoSetup /* client */, []Event /* more client events */, error, /* client error */
-			CryptoSetup /* server */, []Event /* more server events */, error, /* server error */
-		) {
-			client := NewCryptoSetupClient(
-				protocol.ConnectionID{},
-				clientTransportParameters,
-				clientConf,
-				enable0RTT,
-				clientRTTStats,
-				nil,
-				utils.DefaultLogger.WithPrefix("client"),
-				protocol.Version1,
-			)
-
-			if serverTransportParameters.StatelessResetToken == nil {
-				var token protocol.StatelessResetToken
-				serverTransportParameters.StatelessResetToken = &token
+	serverLoop:
+		for {
+			ev := server.NextEvent()
+			switch ev.Kind {
+			case EventNoEvent:
+				break serverLoop
+			case EventWriteInitialData:
+				clientErr = client.HandleMessage(ev.Data, protocol.EncryptionInitial)
+				if clientErr != nil {
+					return
+				}
+			case EventWriteHandshakeData:
+				clientErr = client.HandleMessage(ev.Data, protocol.EncryptionHandshake)
+				if clientErr != nil {
+					return
+				}
+			case EventHandshakeComplete:
+				serverHandshakeComplete = true
+				ticket, err := server.GetSessionTicket()
+				require.NoError(t, err)
+				if ticket != nil {
+					require.NoError(t, client.HandleMessage(ticket, protocol.Encryption1RTT))
+				}
+			default:
+				serverEvents = append(serverEvents, ev)
 			}
-			server := NewCryptoSetupServer(
-				protocol.ConnectionID{},
-				&net.UDPAddr{IP: net.IPv6loopback, Port: 1234},
-				&net.UDPAddr{IP: net.IPv6loopback, Port: 4321},
-				serverTransportParameters,
-				serverConf,
-				enable0RTT,
-				serverRTTStats,
-				nil,
-				utils.DefaultLogger.WithPrefix("server"),
-				protocol.Version1,
-			)
-			cEvents, cErr, sEvents, sErr := handshake(client, server)
-			return client, cEvents, cErr, server, sEvents, sErr
 		}
 
-		It("handshakes", func() {
-			_, _, clientErr, _, _, serverErr := handshakeWithTLSConf(
-				clientConf, serverConf,
-				&utils.RTTStats{}, &utils.RTTStats{},
-				&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
-				false,
-			)
-			Expect(clientErr).ToNot(HaveOccurred())
-			Expect(serverErr).ToNot(HaveOccurred())
-		})
-
-		It("performs a HelloRetryRequst", func() {
-			serverConf.CurvePreferences = []tls.CurveID{tls.CurveP384}
-			_, _, clientErr, _, _, serverErr := handshakeWithTLSConf(
-				clientConf, serverConf,
-				&utils.RTTStats{}, &utils.RTTStats{},
-				&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
-				false,
-			)
-			Expect(clientErr).ToNot(HaveOccurred())
-			Expect(serverErr).ToNot(HaveOccurred())
-		})
-
-		It("handshakes with client auth", func() {
-			clientConf.Certificates = []tls.Certificate{generateCert()}
-			serverConf.ClientAuth = tls.RequireAnyClientCert
-			_, _, clientErr, _, _, serverErr := handshakeWithTLSConf(
-				clientConf, serverConf,
-				&utils.RTTStats{}, &utils.RTTStats{},
-				&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
-				false,
-			)
-			Expect(clientErr).ToNot(HaveOccurred())
-			Expect(serverErr).ToNot(HaveOccurred())
-		})
-
-		It("receives transport parameters", func() {
-			cTransportParameters := &wire.TransportParameters{ActiveConnectionIDLimit: 2, MaxIdleTimeout: 42 * time.Second}
-			client := NewCryptoSetupClient(
-				protocol.ConnectionID{},
-				cTransportParameters,
-				clientConf,
-				false,
-				&utils.RTTStats{},
-				nil,
-				utils.DefaultLogger.WithPrefix("client"),
-				protocol.Version1,
-			)
-
-			var token protocol.StatelessResetToken
-			sTransportParameters := &wire.TransportParameters{
-				MaxIdleTimeout:          1337 * time.Second,
-				StatelessResetToken:     &token,
-				ActiveConnectionIDLimit: 2,
-			}
-			server := NewCryptoSetupServer(
-				protocol.ConnectionID{},
-				&net.UDPAddr{IP: net.IPv6loopback, Port: 1234},
-				&net.UDPAddr{IP: net.IPv6loopback, Port: 4321},
-				sTransportParameters,
-				serverConf,
-				false,
-				&utils.RTTStats{},
-				nil,
-				utils.DefaultLogger.WithPrefix("server"),
-				protocol.Version1,
-			)
-
-			clientEvents, cErr, serverEvents, sErr := handshake(client, server)
-			Expect(cErr).ToNot(HaveOccurred())
-			Expect(sErr).ToNot(HaveOccurred())
-			var clientReceivedTransportParameters *wire.TransportParameters
-			for _, ev := range clientEvents {
-				if ev.Kind == EventReceivedTransportParameters {
-					clientReceivedTransportParameters = ev.TransportParameters
-				}
-			}
-			Expect(clientReceivedTransportParameters).ToNot(BeNil())
-			Expect(clientReceivedTransportParameters.MaxIdleTimeout).To(Equal(1337 * time.Second))
+		if clientHandshakeComplete && serverHandshakeComplete {
+			break
+		}
+	}
+	return
+}
 
-			var serverReceivedTransportParameters *wire.TransportParameters
-			for _, ev := range serverEvents {
-				if ev.Kind == EventReceivedTransportParameters {
-					serverReceivedTransportParameters = ev.TransportParameters
-				}
-			}
-			Expect(serverReceivedTransportParameters).ToNot(BeNil())
-			Expect(serverReceivedTransportParameters.MaxIdleTimeout).To(Equal(42 * time.Second))
-		})
-
-		Context("with session tickets", func() {
-			It("errors when the NewSessionTicket is sent at the wrong encryption level", func() {
-				client, _, clientErr, _, _, serverErr := handshakeWithTLSConf(
-					clientConf, serverConf,
-					&utils.RTTStats{}, &utils.RTTStats{},
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
-					false,
-				)
-				Expect(clientErr).ToNot(HaveOccurred())
-				Expect(serverErr).ToNot(HaveOccurred())
-
-				// inject an invalid session ticket
-				b := append([]byte{uint8(typeNewSessionTicket), 0, 0, 6}, []byte("foobar")...)
-				err := client.HandleMessage(b, protocol.EncryptionHandshake)
-				Expect(err).To(HaveOccurred())
-				Expect(err.Error()).To(ContainSubstring("tls: handshake data received at wrong level"))
-			})
-
-			It("errors when handling the NewSessionTicket fails", func() {
-				client, _, clientErr, _, _, serverErr := handshakeWithTLSConf(
-					clientConf, serverConf,
-					&utils.RTTStats{}, &utils.RTTStats{},
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
-					false,
-				)
-				Expect(clientErr).ToNot(HaveOccurred())
-				Expect(serverErr).ToNot(HaveOccurred())
-
-				// inject an invalid session ticket
-				b := append([]byte{uint8(typeNewSessionTicket), 0, 0, 6}, []byte("foobar")...)
-				err := client.HandleMessage(b, protocol.Encryption1RTT)
-				Expect(err).To(BeAssignableToTypeOf(&qerr.TransportError{}))
-				Expect(err.(*qerr.TransportError).ErrorCode.IsCryptoError()).To(BeTrue())
-			})
-
-			It("uses session resumption", func() {
-				csc := newMockClientSessionCache()
-				clientConf.ClientSessionCache = csc
-				const serverRTT = 25 * time.Millisecond // RTT as measured by the server. Should be restored.
-				const clientRTT = 30 * time.Millisecond // RTT as measured by the client. Should be restored.
-				serverOrigRTTStats := newRTTStatsWithRTT(serverRTT)
-				clientOrigRTTStats := newRTTStatsWithRTT(clientRTT)
-				client, _, clientErr, server, _, serverErr := handshakeWithTLSConf(
-					clientConf, serverConf,
-					clientOrigRTTStats, serverOrigRTTStats,
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
-					false,
-				)
-				Expect(clientErr).ToNot(HaveOccurred())
-				Expect(serverErr).ToNot(HaveOccurred())
-				Eventually(csc.puts).Should(Receive())
-				Expect(server.ConnectionState().DidResume).To(BeFalse())
-				Expect(client.ConnectionState().DidResume).To(BeFalse())
-
-				clientRTTStats := &utils.RTTStats{}
-				serverRTTStats := &utils.RTTStats{}
-				client, _, clientErr, server, _, serverErr = handshakeWithTLSConf(
-					clientConf, serverConf,
-					clientRTTStats, serverRTTStats,
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
-					false,
-				)
-				Expect(clientErr).ToNot(HaveOccurred())
-				Expect(serverErr).ToNot(HaveOccurred())
-				Eventually(csc.puts).Should(Receive())
-				Expect(server.ConnectionState().DidResume).To(BeTrue())
-				Expect(client.ConnectionState().DidResume).To(BeTrue())
-				Expect(clientRTTStats.SmoothedRTT()).To(Equal(clientRTT))
-				Expect(serverRTTStats.SmoothedRTT()).To(Equal(serverRTT))
-			})
-
-			It("doesn't use session resumption if the server disabled it", func() {
-				csc := newMockClientSessionCache()
-				clientConf.ClientSessionCache = csc
-				client, _, clientErr, server, _, serverErr := handshakeWithTLSConf(
-					clientConf, serverConf,
-					&utils.RTTStats{}, &utils.RTTStats{},
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
-					false,
-				)
-				Expect(clientErr).ToNot(HaveOccurred())
-				Expect(serverErr).ToNot(HaveOccurred())
-				Eventually(csc.puts).Should(Receive())
-				Expect(server.ConnectionState().DidResume).To(BeFalse())
-				Expect(client.ConnectionState().DidResume).To(BeFalse())
-
-				serverConf.SessionTicketsDisabled = true
-				client, _, clientErr, server, _, serverErr = handshakeWithTLSConf(
-					clientConf, serverConf,
-					&utils.RTTStats{}, &utils.RTTStats{},
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
-					false,
-				)
-				Expect(clientErr).ToNot(HaveOccurred())
-				Expect(serverErr).ToNot(HaveOccurred())
-				Consistently(csc.puts, 25*time.Millisecond).ShouldNot(Receive())
-				Expect(server.ConnectionState().DidResume).To(BeFalse())
-				Expect(client.ConnectionState().DidResume).To(BeFalse())
-			})
-
-			It("uses 0-RTT", func() {
-				csc := newMockClientSessionCache()
-				clientConf.ClientSessionCache = csc
-				const serverRTT = 25 * time.Millisecond // RTT as measured by the server. Should be restored.
-				const clientRTT = 30 * time.Millisecond // RTT as measured by the client. Should be restored.
-				serverOrigRTTStats := newRTTStatsWithRTT(serverRTT)
-				clientOrigRTTStats := newRTTStatsWithRTT(clientRTT)
-				const initialMaxData protocol.ByteCount = 1337
-				client, _, clientErr, server, _, serverErr := handshakeWithTLSConf(
-					clientConf, serverConf,
-					clientOrigRTTStats, serverOrigRTTStats,
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2},
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2, InitialMaxData: initialMaxData},
-					true,
-				)
-				Expect(clientErr).ToNot(HaveOccurred())
-				Expect(serverErr).ToNot(HaveOccurred())
-				Eventually(csc.puts).Should(Receive())
-				Expect(server.ConnectionState().DidResume).To(BeFalse())
-				Expect(client.ConnectionState().DidResume).To(BeFalse())
-
-				clientRTTStats := &utils.RTTStats{}
-				serverRTTStats := &utils.RTTStats{}
-				client, clientEvents, clientErr, server, serverEvents, serverErr := handshakeWithTLSConf(
-					clientConf, serverConf,
-					clientRTTStats, serverRTTStats,
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2},
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2, InitialMaxData: initialMaxData},
-					true,
-				)
-				Expect(clientErr).ToNot(HaveOccurred())
-				Expect(serverErr).ToNot(HaveOccurred())
-				Expect(clientRTTStats.SmoothedRTT()).To(Equal(clientRTT))
-				Expect(serverRTTStats.SmoothedRTT()).To(Equal(serverRTT))
-
-				var tp *wire.TransportParameters
-				var clientReceived0RTTKeys bool
-				for _, ev := range clientEvents {
-					switch ev.Kind {
-					case EventRestoredTransportParameters:
-						tp = ev.TransportParameters
-					case EventReceivedReadKeys:
-						clientReceived0RTTKeys = true
-					}
-				}
-				Expect(clientReceived0RTTKeys).To(BeTrue())
-				Expect(tp).ToNot(BeNil())
-				Expect(tp.InitialMaxData).To(Equal(initialMaxData))
-
-				var serverReceived0RTTKeys bool
-				for _, ev := range serverEvents {
-					switch ev.Kind {
-					case EventReceivedReadKeys:
-						serverReceived0RTTKeys = true
-					}
-				}
-				Expect(serverReceived0RTTKeys).To(BeTrue())
-
-				Expect(server.ConnectionState().DidResume).To(BeTrue())
-				Expect(client.ConnectionState().DidResume).To(BeTrue())
-				Expect(server.ConnectionState().Used0RTT).To(BeTrue())
-				Expect(client.ConnectionState().Used0RTT).To(BeTrue())
-			})
-
-			It("rejects 0-RTT, when the transport parameters changed", func() {
-				csc := newMockClientSessionCache()
-				clientConf.ClientSessionCache = csc
-				const clientRTT = 30 * time.Millisecond // RTT as measured by the client. Should be restored.
-				clientOrigRTTStats := newRTTStatsWithRTT(clientRTT)
-				const initialMaxData protocol.ByteCount = 1337
-				client, _, clientErr, server, _, serverErr := handshakeWithTLSConf(
-					clientConf, serverConf,
-					clientOrigRTTStats, &utils.RTTStats{},
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2},
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2, InitialMaxData: initialMaxData},
-					true,
-				)
-				Expect(clientErr).ToNot(HaveOccurred())
-				Expect(serverErr).ToNot(HaveOccurred())
-				Eventually(csc.puts).Should(Receive())
-				Expect(server.ConnectionState().DidResume).To(BeFalse())
-				Expect(client.ConnectionState().DidResume).To(BeFalse())
-
-				clientRTTStats := &utils.RTTStats{}
-				client, clientEvents, clientErr, server, _, serverErr := handshakeWithTLSConf(
-					clientConf, serverConf,
-					clientRTTStats, &utils.RTTStats{},
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2},
-					&wire.TransportParameters{ActiveConnectionIDLimit: 2, InitialMaxData: initialMaxData - 1},
-					true,
-				)
-				Expect(clientErr).ToNot(HaveOccurred())
-				Expect(serverErr).ToNot(HaveOccurred())
-				Expect(clientRTTStats.SmoothedRTT()).To(Equal(clientRTT))
-
-				var tp *wire.TransportParameters
-				var clientReceived0RTTKeys bool
-				for _, ev := range clientEvents {
-					switch ev.Kind {
-					case EventRestoredTransportParameters:
-						tp = ev.TransportParameters
-					case EventReceivedReadKeys:
-						clientReceived0RTTKeys = true
-					}
-				}
-				Expect(clientReceived0RTTKeys).To(BeTrue())
-				Expect(tp).ToNot(BeNil())
-				Expect(tp.InitialMaxData).To(Equal(initialMaxData))
-
-				Expect(server.ConnectionState().DidResume).To(BeTrue())
-				Expect(client.ConnectionState().DidResume).To(BeTrue())
-				Expect(server.ConnectionState().Used0RTT).To(BeFalse())
-				Expect(client.ConnectionState().Used0RTT).To(BeFalse())
-			})
-		})
-	})
-})
+func handshakeWithTLSConf(
+	t *testing.T,
+	clientConf, serverConf *tls.Config,
+	clientRTTStats, serverRTTStats *utils.RTTStats,
+	clientTransportParameters, serverTransportParameters *wire.TransportParameters,
+	enable0RTT bool,
+) (CryptoSetup /* client */, []Event /* more client events */, error, /* client error */
+	CryptoSetup /* server */, []Event /* more server events */, error, /* server error */
+) {
+	t.Helper()
+	client := NewCryptoSetupClient(
+		protocol.ConnectionID{},
+		clientTransportParameters,
+		clientConf,
+		enable0RTT,
+		clientRTTStats,
+		nil,
+		utils.DefaultLogger.WithPrefix("client"),
+		protocol.Version1,
+	)
+
+	if serverTransportParameters.StatelessResetToken == nil {
+		var token protocol.StatelessResetToken
+		serverTransportParameters.StatelessResetToken = &token
+	}
+	server := NewCryptoSetupServer(
+		protocol.ConnectionID{},
+		&net.UDPAddr{IP: net.IPv6loopback, Port: 1234},
+		&net.UDPAddr{IP: net.IPv6loopback, Port: 4321},
+		serverTransportParameters,
+		serverConf,
+		enable0RTT,
+		serverRTTStats,
+		nil,
+		utils.DefaultLogger.WithPrefix("server"),
+		protocol.Version1,
+	)
+	cEvents, cErr, sEvents, sErr := handshake(t, client, server)
+	return client, cEvents, cErr, server, sEvents, sErr
+}
+
+func TestHandshake(t *testing.T) {
+	clientConf, serverConf := getTLSConfigs()
+	_, _, clientErr, _, _, serverErr := handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		&utils.RTTStats{}, &utils.RTTStats{},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		false,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+}
+
+func TestHelloRetryRequest(t *testing.T) {
+	clientConf, serverConf := getTLSConfigs()
+	serverConf.CurvePreferences = []tls.CurveID{tls.CurveP384}
+	_, _, clientErr, _, _, serverErr := handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		&utils.RTTStats{}, &utils.RTTStats{},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		false,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+}
+
+func TestWithClientAuth(t *testing.T) {
+	priv, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(t, err)
+	tmpl := &x509.Certificate{
+		SerialNumber:          big.NewInt(1),
+		Subject:               pkix.Name{},
+		SignatureAlgorithm:    x509.SHA256WithRSA,
+		NotBefore:             time.Now(),
+		NotAfter:              time.Now().Add(time.Hour),
+		BasicConstraintsValid: true,
+	}
+	certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, priv.Public(), priv)
+	require.NoError(t, err)
+	clientCert := tls.Certificate{
+		PrivateKey:  priv,
+		Certificate: [][]byte{certDER},
+	}
+
+	clientConf, serverConf := getTLSConfigs()
+	clientConf.Certificates = []tls.Certificate{clientCert}
+	serverConf.ClientAuth = tls.RequireAnyClientCert
+	_, _, clientErr, _, _, serverErr := handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		&utils.RTTStats{}, &utils.RTTStats{},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		false,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+}
+
+func TestTransportParameters(t *testing.T) {
+	clientConf, serverConf := getTLSConfigs()
+	cTransportParameters := &wire.TransportParameters{ActiveConnectionIDLimit: 2, MaxIdleTimeout: 42 * time.Second}
+	client := NewCryptoSetupClient(
+		protocol.ConnectionID{},
+		cTransportParameters,
+		clientConf,
+		false,
+		&utils.RTTStats{},
+		nil,
+		utils.DefaultLogger.WithPrefix("client"),
+		protocol.Version1,
+	)
+
+	var token protocol.StatelessResetToken
+	sTransportParameters := &wire.TransportParameters{
+		MaxIdleTimeout:          1337 * time.Second,
+		StatelessResetToken:     &token,
+		ActiveConnectionIDLimit: 2,
+	}
+	server := NewCryptoSetupServer(
+		protocol.ConnectionID{},
+		&net.UDPAddr{IP: net.IPv6loopback, Port: 1234},
+		&net.UDPAddr{IP: net.IPv6loopback, Port: 4321},
+		sTransportParameters,
+		serverConf,
+		false,
+		&utils.RTTStats{},
+		nil,
+		utils.DefaultLogger.WithPrefix("server"),
+		protocol.Version1,
+	)
+
+	clientEvents, cErr, serverEvents, sErr := handshake(t, client, server)
+	require.NoError(t, cErr)
+	require.NoError(t, sErr)
+	var clientReceivedTransportParameters *wire.TransportParameters
+	for _, ev := range clientEvents {
+		if ev.Kind == EventReceivedTransportParameters {
+			clientReceivedTransportParameters = ev.TransportParameters
+		}
+	}
+	require.NotNil(t, clientReceivedTransportParameters)
+	require.Equal(t, 1337*time.Second, clientReceivedTransportParameters.MaxIdleTimeout)
+
+	var serverReceivedTransportParameters *wire.TransportParameters
+	for _, ev := range serverEvents {
+		if ev.Kind == EventReceivedTransportParameters {
+			serverReceivedTransportParameters = ev.TransportParameters
+		}
+	}
+	require.NotNil(t, serverReceivedTransportParameters)
+	require.Equal(t, 42*time.Second, serverReceivedTransportParameters.MaxIdleTimeout)
+}
+
+func TestNewSessionTicketAtWrongEncryptionLevel(t *testing.T) {
+	clientConf, serverConf := getTLSConfigs()
+	client, _, clientErr, _, _, serverErr := handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		&utils.RTTStats{}, &utils.RTTStats{},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		false,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+
+	// inject an invalid session ticket
+	b := append([]byte{uint8(typeNewSessionTicket), 0, 0, 6}, []byte("foobar")...)
+	err := client.HandleMessage(b, protocol.EncryptionHandshake)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "tls: handshake data received at wrong level")
+}
+
+func TestHandlingNewSessionTicketFails(t *testing.T) {
+	clientConf, serverConf := getTLSConfigs()
+	client, _, clientErr, _, _, serverErr := handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		&utils.RTTStats{}, &utils.RTTStats{},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		false,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+
+	// inject an invalid session ticket
+	b := append([]byte{uint8(typeNewSessionTicket), 0, 0, 6}, []byte("foobar")...)
+	err := client.HandleMessage(b, protocol.Encryption1RTT)
+	require.IsType(t, &qerr.TransportError{}, err)
+	require.True(t, err.(*qerr.TransportError).ErrorCode.IsCryptoError())
+}
+
+func TestSessionResumption(t *testing.T) {
+	clientConf, serverConf := getTLSConfigs()
+	csc := newMockClientSessionCache()
+	clientConf.ClientSessionCache = csc
+	const serverRTT = 25 * time.Millisecond // RTT as measured by the server. Should be restored.
+	const clientRTT = 30 * time.Millisecond // RTT as measured by the client. Should be restored.
+	serverOrigRTTStats := newRTTStatsWithRTT(t, serverRTT)
+	clientOrigRTTStats := newRTTStatsWithRTT(t, clientRTT)
+	client, _, clientErr, server, _, serverErr := handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		clientOrigRTTStats, serverOrigRTTStats,
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		false,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+	select {
+	case <-csc.puts:
+	case <-time.After(time.Second):
+		t.Fatal("didn't receive a session ticket")
+	}
+	require.False(t, server.ConnectionState().DidResume)
+	require.False(t, client.ConnectionState().DidResume)
+
+	clientRTTStats := &utils.RTTStats{}
+	serverRTTStats := &utils.RTTStats{}
+	client, _, clientErr, server, _, serverErr = handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		clientRTTStats, serverRTTStats,
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		false,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+	select {
+	case <-csc.puts:
+	case <-time.After(time.Second):
+		t.Fatal("didn't receive a session ticket")
+	}
+	require.True(t, server.ConnectionState().DidResume)
+	require.True(t, client.ConnectionState().DidResume)
+	require.Equal(t, clientRTT, clientRTTStats.SmoothedRTT())
+	require.Equal(t, serverRTT, serverRTTStats.SmoothedRTT())
+}
+
+func TestSessionResumptionDisabled(t *testing.T) {
+	clientConf, serverConf := getTLSConfigs()
+	csc := newMockClientSessionCache()
+	clientConf.ClientSessionCache = csc
+	client, _, clientErr, server, _, serverErr := handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		&utils.RTTStats{}, &utils.RTTStats{},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		false,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+	select {
+	case <-csc.puts:
+	case <-time.After(time.Second):
+		t.Fatal("didn't receive a session ticket")
+	}
+	require.False(t, server.ConnectionState().DidResume)
+	require.False(t, client.ConnectionState().DidResume)
+
+	serverConf.SessionTicketsDisabled = true
+	client, _, clientErr, server, _, serverErr = handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		&utils.RTTStats{}, &utils.RTTStats{},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2}, &wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		false,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+	select {
+	case <-csc.puts:
+		t.Fatal("didn't expect to receive a session ticket")
+	case <-time.After(25 * time.Millisecond):
+	}
+	require.False(t, server.ConnectionState().DidResume)
+	require.False(t, client.ConnectionState().DidResume)
+}
+
+func Test0RTT(t *testing.T) {
+	clientConf, serverConf := getTLSConfigs()
+	csc := newMockClientSessionCache()
+	clientConf.ClientSessionCache = csc
+	const serverRTT = 25 * time.Millisecond // RTT as measured by the server. Should be restored.
+	const clientRTT = 30 * time.Millisecond // RTT as measured by the client. Should be restored.
+	serverOrigRTTStats := newRTTStatsWithRTT(t, serverRTT)
+	clientOrigRTTStats := newRTTStatsWithRTT(t, clientRTT)
+	const initialMaxData protocol.ByteCount = 1337
+	client, _, clientErr, server, _, serverErr := handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		clientOrigRTTStats, serverOrigRTTStats,
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2, InitialMaxData: initialMaxData},
+		true,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+	select {
+	case <-csc.puts:
+	case <-time.After(time.Second):
+		t.Fatal("didn't receive a session ticket")
+	}
+	require.False(t, server.ConnectionState().DidResume)
+	require.False(t, client.ConnectionState().DidResume)
+
+	clientRTTStats := &utils.RTTStats{}
+	serverRTTStats := &utils.RTTStats{}
+	client, clientEvents, clientErr, server, serverEvents, serverErr := handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		clientRTTStats, serverRTTStats,
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2, InitialMaxData: initialMaxData},
+		true,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+	require.Equal(t, clientRTT, clientRTTStats.SmoothedRTT())
+	require.Equal(t, serverRTT, serverRTTStats.SmoothedRTT())
+
+	var tp *wire.TransportParameters
+	var clientReceived0RTTKeys bool
+	for _, ev := range clientEvents {
+		switch ev.Kind {
+		case EventRestoredTransportParameters:
+			tp = ev.TransportParameters
+		case EventReceivedReadKeys:
+			clientReceived0RTTKeys = true
+		}
+	}
+	require.True(t, clientReceived0RTTKeys)
+	require.NotNil(t, tp)
+	require.Equal(t, initialMaxData, tp.InitialMaxData)
+
+	var serverReceived0RTTKeys bool
+	for _, ev := range serverEvents {
+		switch ev.Kind {
+		case EventReceivedReadKeys:
+			serverReceived0RTTKeys = true
+		}
+	}
+	require.True(t, serverReceived0RTTKeys)
+
+	require.True(t, server.ConnectionState().DidResume)
+	require.True(t, client.ConnectionState().DidResume)
+	require.True(t, server.ConnectionState().Used0RTT)
+	require.True(t, client.ConnectionState().Used0RTT)
+}
+
+func Test0RTTRejectionOnTransportParametersChanged(t *testing.T) {
+	clientConf, serverConf := getTLSConfigs()
+	csc := newMockClientSessionCache()
+	clientConf.ClientSessionCache = csc
+	const clientRTT = 30 * time.Millisecond // RTT as measured by the client. Should be restored.
+	clientOrigRTTStats := newRTTStatsWithRTT(t, clientRTT)
+	const initialMaxData protocol.ByteCount = 1337
+	client, _, clientErr, server, _, serverErr := handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		clientOrigRTTStats, &utils.RTTStats{},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2, InitialMaxData: initialMaxData},
+		true,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+	select {
+	case <-csc.puts:
+	case <-time.After(time.Second):
+		t.Fatal("didn't receive a session ticket")
+	}
+	require.False(t, server.ConnectionState().DidResume)
+	require.False(t, client.ConnectionState().DidResume)
+
+	clientRTTStats := &utils.RTTStats{}
+	client, clientEvents, clientErr, server, _, serverErr := handshakeWithTLSConf(
+		t,
+		clientConf, serverConf,
+		clientRTTStats, &utils.RTTStats{},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2},
+		&wire.TransportParameters{ActiveConnectionIDLimit: 2, InitialMaxData: initialMaxData - 1},
+		true,
+	)
+	require.NoError(t, clientErr)
+	require.NoError(t, serverErr)
+	require.Equal(t, clientRTT, clientRTTStats.SmoothedRTT())
+
+	var tp *wire.TransportParameters
+	var clientReceived0RTTKeys bool
+	for _, ev := range clientEvents {
+		switch ev.Kind {
+		case EventRestoredTransportParameters:
+			tp = ev.TransportParameters
+		case EventReceivedReadKeys:
+			clientReceived0RTTKeys = true
+		}
+	}
+	require.True(t, clientReceived0RTTKeys)
+	require.NotNil(t, tp)
+	require.Equal(t, initialMaxData, tp.InitialMaxData)
+
+	require.True(t, server.ConnectionState().DidResume)
+	require.True(t, client.ConnectionState().DidResume)
+	require.False(t, server.ConnectionState().Used0RTT)
+	require.False(t, client.ConnectionState().Used0RTT)
+}
diff --git a/internal/handshake/handshake_helpers_test.go b/internal/handshake/handshake_helpers_test.go
new file mode 100644
index 00000000..a684c2cc
--- /dev/null
+++ b/internal/handshake/handshake_helpers_test.go
@@ -0,0 +1,35 @@
+package handshake
+
+import (
+	"crypto/tls"
+	"encoding/hex"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func splitHexString(t *testing.T, s string) (slice []byte) {
+	t.Helper()
+	for _, ss := range strings.Split(s, " ") {
+		if ss[0:2] == "0x" {
+			ss = ss[2:]
+		}
+		d, err := hex.DecodeString(ss)
+		require.NoError(t, err)
+		slice = append(slice, d...)
+	}
+	return
+}
+
+func TestSplitHexString(t *testing.T) {
+	require.Equal(t, []byte{0xde, 0xad, 0xbe, 0xef}, splitHexString(t, "0xdeadbeef"))
+	require.Equal(t, []byte{0xde, 0xad, 0xbe, 0xef}, splitHexString(t, "deadbeef"))
+	require.Equal(t, []byte{0xde, 0xad, 0xbe, 0xef}, splitHexString(t, "dead beef"))
+}
+
+var cipherSuites = []*cipherSuite{
+	getCipherSuite(tls.TLS_AES_128_GCM_SHA256),
+	getCipherSuite(tls.TLS_AES_256_GCM_SHA384),
+	getCipherSuite(tls.TLS_CHACHA20_POLY1305_SHA256),
+}
diff --git a/internal/handshake/handshake_suite_test.go b/internal/handshake/handshake_suite_test.go
deleted file mode 100644
index 80e48a30..00000000
--- a/internal/handshake/handshake_suite_test.go
+++ /dev/null
@@ -1,45 +0,0 @@
-package handshake
-
-import (
-	"crypto/tls"
-	"encoding/hex"
-	"strings"
-	"testing"
-
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
-	"go.uber.org/mock/gomock"
-)
-
-func TestHandshake(t *testing.T) {
-	RegisterFailHandler(Fail)
-	RunSpecs(t, "Handshake Suite")
-}
-
-var mockCtrl *gomock.Controller
-
-var _ = BeforeEach(func() {
-	mockCtrl = gomock.NewController(GinkgoT())
-})
-
-var _ = AfterEach(func() {
-	mockCtrl.Finish()
-})
-
-func splitHexString(s string) (slice []byte) {
-	for _, ss := range strings.Split(s, " ") {
-		if ss[0:2] == "0x" {
-			ss = ss[2:]
-		}
-		d, err := hex.DecodeString(ss)
-		ExpectWithOffset(1, err).ToNot(HaveOccurred())
-		slice = append(slice, d...)
-	}
-	return
-}
-
-var cipherSuites = []*cipherSuite{
-	getCipherSuite(tls.TLS_AES_128_GCM_SHA256),
-	getCipherSuite(tls.TLS_AES_256_GCM_SHA384),
-	getCipherSuite(tls.TLS_CHACHA20_POLY1305_SHA256),
-}
diff --git a/internal/handshake/hkdf_test.go b/internal/handshake/hkdf_test.go
index f32e016e..59b4567d 100644
--- a/internal/handshake/hkdf_test.go
+++ b/internal/handshake/hkdf_test.go
@@ -9,8 +9,7 @@ import (
 
 	"golang.org/x/exp/rand"
 
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/require"
 )
 
 type cipherSuiteTLS13 struct {
@@ -36,19 +35,29 @@ func cipherSuiteTLS13ByID(id uint16) *cipherSuiteTLS13 {
 //go:linkname expandLabel crypto/tls.(*cipherSuiteTLS13).expandLabel
 func expandLabel(cs *cipherSuiteTLS13, secret []byte, label string, context []byte, length int) []byte
 
-var _ = Describe("HKDF", func() {
-	DescribeTable("gets the same results as crypto/tls",
-		func(cipherSuite uint16, secret, context []byte, label string, length int) {
-			cs := cipherSuiteTLS13ByID(cipherSuite)
-			expected := expandLabel(cs, secret, label, context, length)
-			expanded := hkdfExpandLabel(cs.Hash, secret, context, label, length)
-			Expect(expanded).To(Equal(expected))
-		},
-		Entry("TLS_AES_128_GCM_SHA256", tls.TLS_AES_128_GCM_SHA256, []byte("secret"), []byte("context"), "label", 42),
-		Entry("TLS_AES_256_GCM_SHA384", tls.TLS_AES_256_GCM_SHA384, []byte("secret"), []byte("context"), "label", 100),
-		Entry("TLS_CHACHA20_POLY1305_SHA256", tls.TLS_CHACHA20_POLY1305_SHA256, []byte("secret"), []byte("context"), "label", 77),
-	)
-})
+func TestHKDF(t *testing.T) {
+	testCases := []struct {
+		name        string
+		cipherSuite uint16
+		secret      []byte
+		context     []byte
+		label       string
+		length      int
+	}{
+		{"TLS_AES_128_GCM_SHA256", tls.TLS_AES_128_GCM_SHA256, []byte("secret"), []byte("context"), "label", 42},
+		{"TLS_AES_256_GCM_SHA384", tls.TLS_AES_256_GCM_SHA384, []byte("secret"), []byte("context"), "label", 100},
+		{"TLS_CHACHA20_POLY1305_SHA256", tls.TLS_CHACHA20_POLY1305_SHA256, []byte("secret"), []byte("context"), "label", 77},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			cs := cipherSuiteTLS13ByID(tc.cipherSuite)
+			expected := expandLabel(cs, tc.secret, tc.label, tc.context, tc.length)
+			expanded := hkdfExpandLabel(cs.Hash, tc.secret, tc.context, tc.label, tc.length)
+			require.Equal(t, expected, expanded)
+		})
+	}
+}
 
 func BenchmarkHKDFExpandLabelStandardLibrary(b *testing.B) {
 	b.Run("TLS_AES_128_GCM_SHA256", func(b *testing.B) { benchmarkHKDFExpandLabel(b, tls.TLS_AES_128_GCM_SHA256, true) })
diff --git a/internal/handshake/initial_aead_test.go b/internal/handshake/initial_aead_test.go
index 57388a4a..b1c1fbee 100644
--- a/internal/handshake/initial_aead_test.go
+++ b/internal/handshake/initial_aead_test.go
@@ -2,194 +2,247 @@ package handshake
 
 import (
 	"bytes"
-	"fmt"
 	"testing"
 
 	"golang.org/x/exp/rand"
 
 	"github.com/quic-go/quic-go/internal/protocol"
 
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/require"
 )
 
-var _ = Describe("Initial AEAD using AES-GCM", func() {
-	It("converts the string representation used in the draft into byte slices", func() {
-		Expect(splitHexString("0xdeadbeef")).To(Equal([]byte{0xde, 0xad, 0xbe, 0xef}))
-		Expect(splitHexString("deadbeef")).To(Equal([]byte{0xde, 0xad, 0xbe, 0xef}))
-		Expect(splitHexString("dead beef")).To(Equal([]byte{0xde, 0xad, 0xbe, 0xef}))
-	})
+func TestComputeClientKeyAndIV(t *testing.T) {
+	connID := protocol.ParseConnectionID(splitHexString(t, "0x8394c8f03e515708"))
+
+	tests := []struct {
+		name                 string
+		version              protocol.Version
+		expectedClientSecret []byte
+		expectedKey          []byte
+		expectedIV           []byte
+	}{
+		{
+			name:                 "QUIC v1",
+			version:              protocol.Version1,
+			expectedClientSecret: splitHexString(t, "c00cf151ca5be075ed0ebfb5c80323c4 2d6b7db67881289af4008f1f6c357aea"),
+			expectedKey:          splitHexString(t, "1f369613dd76d5467730efcbe3b1a22d"),
+			expectedIV:           splitHexString(t, "fa044b2f42a3fd3b46fb255c"),
+		},
+		{
+			name:                 "QUIC v2",
+			version:              protocol.Version2,
+			expectedClientSecret: splitHexString(t, "14ec9d6eb9fd7af83bf5a668bc17a7e2 83766aade7ecd0891f70f9ff7f4bf47b"),
+			expectedKey:          splitHexString(t, "8b1a0bc121284290a29e0971b5cd045d"),
+			expectedIV:           splitHexString(t, "91f73e2351d8fa91660e909f"),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			clientSecret, _ := computeSecrets(connID, tt.version)
+			require.Equal(t, tt.expectedClientSecret, clientSecret)
+			key, iv := computeInitialKeyAndIV(clientSecret, tt.version)
+			require.Equal(t, tt.expectedKey, key)
+			require.Equal(t, tt.expectedIV, iv)
+		})
+	}
+}
+
+func TestComputeServerKeyAndIV(t *testing.T) {
+	connID := protocol.ParseConnectionID(splitHexString(t, "0x8394c8f03e515708"))
+
+	tests := []struct {
+		name                 string
+		version              protocol.Version
+		expectedServerSecret []byte
+		expectedKey          []byte
+		expectedIV           []byte
+	}{
+		{
+			name:                 "QUIC v1",
+			version:              protocol.Version1,
+			expectedServerSecret: splitHexString(t, "3c199828fd139efd216c155ad844cc81 fb82fa8d7446fa7d78be803acdda951b"),
+			expectedKey:          splitHexString(t, "cf3a5331653c364c88f0f379b6067e37"),
+			expectedIV:           splitHexString(t, "0ac1493ca1905853b0bba03e"),
+		},
+		{
+			name:                 "QUIC v2",
+			version:              protocol.Version2,
+			expectedServerSecret: splitHexString(t, "0263db1782731bf4588e7e4d93b74639 07cb8cd8200b5da55a8bd488eafc37c1"),
+			expectedKey:          splitHexString(t, "82db637861d55e1d011f19ea71d5d2a7"),
+			expectedIV:           splitHexString(t, "dd13c276499c0249d3310652"),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, serverSecret := computeSecrets(connID, tt.version)
+			require.Equal(t, tt.expectedServerSecret, serverSecret)
+			key, iv := computeInitialKeyAndIV(serverSecret, tt.version)
+			require.Equal(t, tt.expectedKey, key)
+			require.Equal(t, tt.expectedIV, iv)
+		})
+	}
+}
 
-	connID := protocol.ParseConnectionID(splitHexString("0x8394c8f03e515708"))
+func TestClientInitial(t *testing.T) {
+	connID := protocol.ParseConnectionID(splitHexString(t, "0x8394c8f03e515708"))
 
-	DescribeTable("computes the client key and IV",
-		func(v protocol.Version, expectedClientSecret, expectedKey, expectedIV []byte) {
-			clientSecret, _ := computeSecrets(connID, v)
-			Expect(clientSecret).To(Equal(expectedClientSecret))
-			key, iv := computeInitialKeyAndIV(clientSecret, v)
-			Expect(key).To(Equal(expectedKey))
-			Expect(iv).To(Equal(expectedIV))
+	tests := []struct {
+		version              protocol.Version
+		header               []byte
+		data                 []byte
+		expectedSample       []byte
+		expectedHdrFirstByte byte
+		expectedHdr          []byte
+		expectedPacket       []byte
+	}{
+		{
+			version:              protocol.Version1,
+			header:               splitHexString(t, "c300000001088394c8f03e5157080000449e00000002"),
+			data:                 splitHexString(t, "060040f1010000ed0303ebf8fa56f129 39b9584a3896472ec40bb863cfd3e868 04fe3a47f06a2b69484c000004130113 02010000c000000010000e00000b6578 616d706c652e636f6dff01000100000a 00080006001d00170018001000070005 04616c706e0005000501000000000033 00260024001d00209370b2c9caa47fba baf4559fedba753de171fa71f50f1ce1 5d43e994ec74d748002b000302030400 0d0010000e0403050306030203080408 050806002d00020101001c0002400100 3900320408ffffffffffffffff050480 00ffff07048000ffff08011001048000 75300901100f088394c8f03e51570806 048000ffff"),
+			expectedSample:       splitHexString(t, "d1b1c98dd7689fb8ec11d242b123dc9b"),
+			expectedHdrFirstByte: byte(0xc0),
+			expectedHdr:          splitHexString(t, "7b9aec34"),
+			expectedPacket:       splitHexString(t, "c000000001088394c8f03e5157080000 449e7b9aec34d1b1c98dd7689fb8ec11 d242b123dc9bd8bab936b47d92ec356c 0bab7df5976d27cd449f63300099f399 1c260ec4c60d17b31f8429157bb35a12 82a643a8d2262cad67500cadb8e7378c 8eb7539ec4d4905fed1bee1fc8aafba1 7c750e2c7ace01e6005f80fcb7df6212 30c83711b39343fa028cea7f7fb5ff89 eac2308249a02252155e2347b63d58c5 457afd84d05dfffdb20392844ae81215 4682e9cf012f9021a6f0be17ddd0c208 4dce25ff9b06cde535d0f920a2db1bf3 62c23e596d11a4f5a6cf3948838a3aec 4e15daf8500a6ef69ec4e3feb6b1d98e 610ac8b7ec3faf6ad760b7bad1db4ba3 485e8a94dc250ae3fdb41ed15fb6a8e5 eba0fc3dd60bc8e30c5c4287e53805db 059ae0648db2f64264ed5e39be2e20d8 2df566da8dd5998ccabdae053060ae6c 7b4378e846d29f37ed7b4ea9ec5d82e7 961b7f25a9323851f681d582363aa5f8 9937f5a67258bf63ad6f1a0b1d96dbd4 faddfcefc5266ba6611722395c906556 be52afe3f565636ad1b17d508b73d874 3eeb524be22b3dcbc2c7468d54119c74 68449a13d8e3b95811a198f3491de3e7 fe942b330407abf82a4ed7c1b311663a c69890f4157015853d91e923037c227a 33cdd5ec281ca3f79c44546b9d90ca00 f064c99e3dd97911d39fe9c5d0b23a22 9a234cb36186c4819e8b9c5927726632 291d6a418211cc2962e20fe47feb3edf 330f2c603a9d48c0fcb5699dbfe58964 25c5bac4aee82e57a85aaf4e2513e4f0 5796b07ba2ee47d80506f8d2c25e50fd 14de71e6c418559302f939b0e1abd576 f279c4b2e0feb85c1f28ff18f58891ff ef132eef2fa09346aee33c28eb130ff2 8f5b766953334113211996d20011a198 e3fc433f9f2541010ae17c1bf202580f 6047472fb36857fe843b19f5984009dd c324044e847a4f4a0ab34f719595de37 252d6235365e9b84392b061085349d73 203a4a13e96f5432ec0fd4a1ee65accd d5e3904df54c1da510b0ff20dcc0c77f cb2c0e0eb605cb0504db87632cf3d8b4 dae6e705769d1de354270123cb11450e fc60ac47683d7b8d0f811365565fd98c 4c8eb936bcab8d069fc33bd801b03ade a2e1fbc5aa463d08ca19896d2bf59a07 1b851e6c239052172f296bfb5e724047 90a2181014f3b94a4e97d117b4381303 68cc39dbb2d198065ae3986547926cd2 162f40a29f0c3c8745c0f50fba3852e5 66d44575c29d39a03f0cda721984b6f4 40591f355e12d439ff150aab7613499d bd49adabc8676eef023b15b65bfc5ca0 6948109f23f350db82123535eb8a7433 bdabcb909271a6ecbcb58b936a88cd4e 8f2e6ff5800175f113253d8fa9ca8885 c2f552e657dc603f252e1a8e308f76f0 be79e2fb8f5d5fbbe2e30ecadd220723 c8c0aea8078cdfcb3868263ff8f09400 54da48781893a7e49ad5aff4af300cd8 04a6b6279ab3ff3afb64491c85194aab 760d58a606654f9f4400e8b38591356f bf6425aca26dc85244259ff2b19c41b9 f96f3ca9ec1dde434da7d2d392b905dd f3d1f9af93d1af5950bd493f5aa731b4 056df31bd267b6b90a079831aaf579be 0a39013137aac6d404f518cfd4684064 7e78bfe706ca4cf5e9c5453e9f7cfd2b 8b4c8d169a44e55c88d4a9a7f9474241 e221af44860018ab0856972e194cd934"),
 		},
-		Entry("QUIC v1",
-			protocol.Version1,
-			splitHexString("c00cf151ca5be075ed0ebfb5c80323c4 2d6b7db67881289af4008f1f6c357aea"),
-			splitHexString("1f369613dd76d5467730efcbe3b1a22d"),
-			splitHexString("fa044b2f42a3fd3b46fb255c"),
-		),
-		Entry("QUIC v2",
-			protocol.Version2,
-			splitHexString("14ec9d6eb9fd7af83bf5a668bc17a7e2 83766aade7ecd0891f70f9ff7f4bf47b"),
-			splitHexString("8b1a0bc121284290a29e0971b5cd045d"),
-			splitHexString("91f73e2351d8fa91660e909f"),
-		),
-	)
-
-	DescribeTable("computes the server key and IV",
-		func(v protocol.Version, expectedServerSecret, expectedKey, expectedIV []byte) {
-			_, serverSecret := computeSecrets(connID, v)
-			Expect(serverSecret).To(Equal(expectedServerSecret))
-			key, iv := computeInitialKeyAndIV(serverSecret, v)
-			Expect(key).To(Equal(expectedKey))
-			Expect(iv).To(Equal(expectedIV))
+		{
+			version:              protocol.Version2,
+			header:               splitHexString(t, "d36b3343cf088394c8f03e5157080000449e00000002"),
+			data:                 splitHexString(t, "060040f1010000ed0303ebf8fa56f129 39b9584a3896472ec40bb863cfd3e868 04fe3a47f06a2b69484c000004130113 02010000c000000010000e00000b6578 616d706c652e636f6dff01000100000a 00080006001d00170018001000070005 04616c706e0005000501000000000033 00260024001d00209370b2c9caa47fba baf4559fedba753de171fa71f50f1ce1 5d43e994ec74d748002b000302030400 0d0010000e0403050306030203080408 050806002d00020101001c0002400100 3900320408ffffffffffffffff050480 00ffff07048000ffff08011001048000 75300901100f088394c8f03e51570806 048000ffff"),
+			expectedSample:       splitHexString(t, "ffe67b6abcdb4298b485dd04de806071"),
+			expectedHdrFirstByte: byte(0xd7),
+			expectedHdr:          splitHexString(t, "a0c95e82"),
+			expectedPacket:       splitHexString(t, "d76b3343cf088394c8f03e5157080000 449ea0c95e82ffe67b6abcdb4298b485 dd04de806071bf03dceebfa162e75d6c 96058bdbfb127cdfcbf903388e99ad04 9f9a3dd4425ae4d0992cfff18ecf0fdb 5a842d09747052f17ac2053d21f57c5d 250f2c4f0e0202b70785b7946e992e58 a59ac52dea6774d4f03b55545243cf1a 12834e3f249a78d395e0d18f4d766004 f1a2674802a747eaa901c3f10cda5500 cb9122faa9f1df66c392079a1b40f0de 1c6054196a11cbea40afb6ef5253cd68 18f6625efce3b6def6ba7e4b37a40f77 32e093daa7d52190935b8da58976ff33 12ae50b187c1433c0f028edcc4c2838b 6a9bfc226ca4b4530e7a4ccee1bfa2a3 d396ae5a3fb512384b2fdd851f784a65 e03f2c4fbe11a53c7777c023462239dd 6f7521a3f6c7d5dd3ec9b3f233773d4b 46d23cc375eb198c63301c21801f6520 bcfb7966fc49b393f0061d974a2706df 8c4a9449f11d7f3d2dcbb90c6b877045 636e7c0c0fe4eb0f697545460c806910 d2c355f1d253bc9d2452aaa549e27a1f ac7cf4ed77f322e8fa894b6a83810a34 b361901751a6f5eb65a0326e07de7c12 16ccce2d0193f958bb3850a833f7ae43 2b65bc5a53975c155aa4bcb4f7b2c4e5 4df16efaf6ddea94e2c50b4cd1dfe060 17e0e9d02900cffe1935e0491d77ffb4 fdf85290fdd893d577b1131a610ef6a5 c32b2ee0293617a37cbb08b847741c3b 8017c25ca9052ca1079d8b78aebd4787 6d330a30f6a8c6d61dd1ab5589329de7 14d19d61370f8149748c72f132f0fc99 f34d766c6938597040d8f9e2bb522ff9 9c63a344d6a2ae8aa8e51b7b90a4a806 105fcbca31506c446151adfeceb51b91 abfe43960977c87471cf9ad4074d30e1 0d6a7f03c63bd5d4317f68ff325ba3bd 80bf4dc8b52a0ba031758022eb025cdd 770b44d6d6cf0670f4e990b22347a7db 848265e3e5eb72dfe8299ad7481a4083 22cac55786e52f633b2fb6b614eaed18 d703dd84045a274ae8bfa73379661388 d6991fe39b0d93debb41700b41f90a15 c4d526250235ddcd6776fc77bc97e7a4 17ebcb31600d01e57f32162a8560cacc 7e27a096d37a1a86952ec71bd89a3e9a 30a2a26162984d7740f81193e8238e61 f6b5b984d4d3dfa033c1bb7e4f0037fe bf406d91c0dccf32acf423cfa1e70710 10d3f270121b493ce85054ef58bada42 310138fe081adb04e2bd901f2f13458b 3d6758158197107c14ebb193230cd115 7380aa79cae1374a7c1e5bbcb80ee23e 06ebfde206bfb0fcbc0edc4ebec30966 1bdd908d532eb0c6adc38b7ca7331dce 8dfce39ab71e7c32d318d136b6100671 a1ae6a6600e3899f31f0eed19e3417d1 34b90c9058f8632c798d4490da498730 7cba922d61c39805d072b589bd52fdf1 e86215c2d54e6670e07383a27bbffb5a ddf47d66aa85a0c6f9f32e59d85a44dd 5d3b22dc2be80919b490437ae4f36a0a e55edf1d0b5cb4e9a3ecabee93dfc6e3 8d209d0fa6536d27a5d6fbb17641cde2 7525d61093f1b28072d111b2b4ae5f89 d5974ee12e5cf7d5da4d6a31123041f3 3e61407e76cffcdcfd7e19ba58cf4b53 6f4c4938ae79324dc402894b44faf8af bab35282ab659d13c93f70412e85cb19 9a37ddec600545473cfb5a05e08d0b20 9973b2172b4d21fb69745a262ccde96b a18b2faa745b6fe189cf772a9f84cbfc"),
 		},
-		Entry("QUIC v1",
-			protocol.Version1,
-			splitHexString("3c199828fd139efd216c155ad844cc81 fb82fa8d7446fa7d78be803acdda951b"),
-			splitHexString("cf3a5331653c364c88f0f379b6067e37"),
-			splitHexString("0ac1493ca1905853b0bba03e"),
-		),
-		Entry("QUIC v2",
-			protocol.Version2,
-			splitHexString("0263db1782731bf4588e7e4d93b74639 07cb8cd8200b5da55a8bd488eafc37c1"),
-			splitHexString("82db637861d55e1d011f19ea71d5d2a7"),
-			splitHexString("dd13c276499c0249d3310652"),
-		),
-	)
-
-	DescribeTable("encrypts the client's Initial",
-		func(v protocol.Version, header, data, expectedSample []byte, expectedHdrFirstByte byte, expectedHdr, expectedPacket []byte) {
-			sealer, _ := NewInitialAEAD(connID, protocol.PerspectiveClient, v)
-			data = append(data, make([]byte, 1162-len(data))...) // add PADDING
-			sealed := sealer.Seal(nil, data, 2, header)
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.version.String(), func(t *testing.T) {
+			sealer, _ := NewInitialAEAD(connID, protocol.PerspectiveClient, tt.version)
+			tt.data = append(tt.data, make([]byte, 1162-len(tt.data))...) // add PADDING
+			sealed := sealer.Seal(nil, tt.data, 2, tt.header)
 			sample := sealed[0:16]
-			Expect(sample).To(Equal(expectedSample))
-			sealer.EncryptHeader(sample, &header[0], header[len(header)-4:])
-			Expect(header[0]).To(Equal(expectedHdrFirstByte))
-			Expect(header[len(header)-4:]).To(Equal(expectedHdr))
-			packet := append(header, sealed...)
-			Expect(packet).To(Equal(expectedPacket))
+			require.Equal(t, tt.expectedSample, sample)
+			sealer.EncryptHeader(sample, &tt.header[0], tt.header[len(tt.header)-4:])
+			require.Equal(t, tt.expectedHdrFirstByte, tt.header[0])
+			require.Equal(t, tt.expectedHdr, tt.header[len(tt.header)-4:])
+			packet := append(tt.header, sealed...)
+			require.Equal(t, tt.expectedPacket, packet)
+		})
+	}
+}
+
+func TestServersInitial(t *testing.T) {
+	connID := protocol.ParseConnectionID(splitHexString(t, "0x8394c8f03e515708"))
+
+	testCases := []struct {
+		name           string
+		version        protocol.Version
+		header         []byte
+		data           []byte
+		expectedSample []byte
+		expectedHdr    []byte
+		expectedPacket []byte
+	}{
+		{
+			name:           "QUIC v1",
+			version:        protocol.Version1,
+			header:         splitHexString(t, "c1000000010008f067a5502a4262b50040750001"),
+			data:           splitHexString(t, "02000000000600405a020000560303ee fce7f7b37ba1d1632e96677825ddf739 88cfc79825df566dc5430b9a045a1200 130100002e00330024001d00209d3c94 0d89690b84d08a60993c144eca684d10 81287c834d5311bcf32bb9da1a002b00 020304"),
+			expectedSample: splitHexString(t, "2cd0991cd25b0aac406a5816b6394100"),
+			expectedHdr:    splitHexString(t, "cf000000010008f067a5502a4262b5004075c0d9"),
+			expectedPacket: splitHexString(t, "cf000000010008f067a5502a4262b500 4075c0d95a482cd0991cd25b0aac406a 5816b6394100f37a1c69797554780bb3 8cc5a99f5ede4cf73c3ec2493a1839b3 dbcba3f6ea46c5b7684df3548e7ddeb9 c3bf9c73cc3f3bded74b562bfb19fb84 022f8ef4cdd93795d77d06edbb7aaf2f 58891850abbdca3d20398c276456cbc4 2158407dd074ee"),
 		},
-		Entry("QUIC v1",
-			protocol.Version1,
-			splitHexString("c300000001088394c8f03e5157080000449e00000002"),
-			splitHexString("060040f1010000ed0303ebf8fa56f129 39b9584a3896472ec40bb863cfd3e868 04fe3a47f06a2b69484c000004130113 02010000c000000010000e00000b6578 616d706c652e636f6dff01000100000a 00080006001d00170018001000070005 04616c706e0005000501000000000033 00260024001d00209370b2c9caa47fba baf4559fedba753de171fa71f50f1ce1 5d43e994ec74d748002b000302030400 0d0010000e0403050306030203080408 050806002d00020101001c0002400100 3900320408ffffffffffffffff050480 00ffff07048000ffff08011001048000 75300901100f088394c8f03e51570806 048000ffff"),
-			splitHexString("d1b1c98dd7689fb8ec11d242b123dc9b"),
-			byte(0xc0),
-			splitHexString("7b9aec34"),
-			splitHexString("c000000001088394c8f03e5157080000 449e7b9aec34d1b1c98dd7689fb8ec11 d242b123dc9bd8bab936b47d92ec356c 0bab7df5976d27cd449f63300099f399 1c260ec4c60d17b31f8429157bb35a12 82a643a8d2262cad67500cadb8e7378c 8eb7539ec4d4905fed1bee1fc8aafba1 7c750e2c7ace01e6005f80fcb7df6212 30c83711b39343fa028cea7f7fb5ff89 eac2308249a02252155e2347b63d58c5 457afd84d05dfffdb20392844ae81215 4682e9cf012f9021a6f0be17ddd0c208 4dce25ff9b06cde535d0f920a2db1bf3 62c23e596d11a4f5a6cf3948838a3aec 4e15daf8500a6ef69ec4e3feb6b1d98e 610ac8b7ec3faf6ad760b7bad1db4ba3 485e8a94dc250ae3fdb41ed15fb6a8e5 eba0fc3dd60bc8e30c5c4287e53805db 059ae0648db2f64264ed5e39be2e20d8 2df566da8dd5998ccabdae053060ae6c 7b4378e846d29f37ed7b4ea9ec5d82e7 961b7f25a9323851f681d582363aa5f8 9937f5a67258bf63ad6f1a0b1d96dbd4 faddfcefc5266ba6611722395c906556 be52afe3f565636ad1b17d508b73d874 3eeb524be22b3dcbc2c7468d54119c74 68449a13d8e3b95811a198f3491de3e7 fe942b330407abf82a4ed7c1b311663a c69890f4157015853d91e923037c227a 33cdd5ec281ca3f79c44546b9d90ca00 f064c99e3dd97911d39fe9c5d0b23a22 9a234cb36186c4819e8b9c5927726632 291d6a418211cc2962e20fe47feb3edf 330f2c603a9d48c0fcb5699dbfe58964 25c5bac4aee82e57a85aaf4e2513e4f0 5796b07ba2ee47d80506f8d2c25e50fd 14de71e6c418559302f939b0e1abd576 f279c4b2e0feb85c1f28ff18f58891ff ef132eef2fa09346aee33c28eb130ff2 8f5b766953334113211996d20011a198 e3fc433f9f2541010ae17c1bf202580f 6047472fb36857fe843b19f5984009dd c324044e847a4f4a0ab34f719595de37 252d6235365e9b84392b061085349d73 203a4a13e96f5432ec0fd4a1ee65accd d5e3904df54c1da510b0ff20dcc0c77f cb2c0e0eb605cb0504db87632cf3d8b4 dae6e705769d1de354270123cb11450e fc60ac47683d7b8d0f811365565fd98c 4c8eb936bcab8d069fc33bd801b03ade a2e1fbc5aa463d08ca19896d2bf59a07 1b851e6c239052172f296bfb5e724047 90a2181014f3b94a4e97d117b4381303 68cc39dbb2d198065ae3986547926cd2 162f40a29f0c3c8745c0f50fba3852e5 66d44575c29d39a03f0cda721984b6f4 40591f355e12d439ff150aab7613499d bd49adabc8676eef023b15b65bfc5ca0 6948109f23f350db82123535eb8a7433 bdabcb909271a6ecbcb58b936a88cd4e 8f2e6ff5800175f113253d8fa9ca8885 c2f552e657dc603f252e1a8e308f76f0 be79e2fb8f5d5fbbe2e30ecadd220723 c8c0aea8078cdfcb3868263ff8f09400 54da48781893a7e49ad5aff4af300cd8 04a6b6279ab3ff3afb64491c85194aab 760d58a606654f9f4400e8b38591356f bf6425aca26dc85244259ff2b19c41b9 f96f3ca9ec1dde434da7d2d392b905dd f3d1f9af93d1af5950bd493f5aa731b4 056df31bd267b6b90a079831aaf579be 0a39013137aac6d404f518cfd4684064 7e78bfe706ca4cf5e9c5453e9f7cfd2b 8b4c8d169a44e55c88d4a9a7f9474241 e221af44860018ab0856972e194cd934"),
-		),
-		Entry("QUIC v2",
-			protocol.Version2,
-			splitHexString("d36b3343cf088394c8f03e5157080000449e00000002"),
-			splitHexString("060040f1010000ed0303ebf8fa56f129 39b9584a3896472ec40bb863cfd3e868 04fe3a47f06a2b69484c000004130113 02010000c000000010000e00000b6578 616d706c652e636f6dff01000100000a 00080006001d00170018001000070005 04616c706e0005000501000000000033 00260024001d00209370b2c9caa47fba baf4559fedba753de171fa71f50f1ce1 5d43e994ec74d748002b000302030400 0d0010000e0403050306030203080408 050806002d00020101001c0002400100 3900320408ffffffffffffffff050480 00ffff07048000ffff08011001048000 75300901100f088394c8f03e51570806 048000ffff"),
-			splitHexString("ffe67b6abcdb4298b485dd04de806071"),
-			byte(0xd7),
-			splitHexString("a0c95e82"),
-			splitHexString("d76b3343cf088394c8f03e5157080000 449ea0c95e82ffe67b6abcdb4298b485 dd04de806071bf03dceebfa162e75d6c 96058bdbfb127cdfcbf903388e99ad04 9f9a3dd4425ae4d0992cfff18ecf0fdb 5a842d09747052f17ac2053d21f57c5d 250f2c4f0e0202b70785b7946e992e58 a59ac52dea6774d4f03b55545243cf1a 12834e3f249a78d395e0d18f4d766004 f1a2674802a747eaa901c3f10cda5500 cb9122faa9f1df66c392079a1b40f0de 1c6054196a11cbea40afb6ef5253cd68 18f6625efce3b6def6ba7e4b37a40f77 32e093daa7d52190935b8da58976ff33 12ae50b187c1433c0f028edcc4c2838b 6a9bfc226ca4b4530e7a4ccee1bfa2a3 d396ae5a3fb512384b2fdd851f784a65 e03f2c4fbe11a53c7777c023462239dd 6f7521a3f6c7d5dd3ec9b3f233773d4b 46d23cc375eb198c63301c21801f6520 bcfb7966fc49b393f0061d974a2706df 8c4a9449f11d7f3d2dcbb90c6b877045 636e7c0c0fe4eb0f697545460c806910 d2c355f1d253bc9d2452aaa549e27a1f ac7cf4ed77f322e8fa894b6a83810a34 b361901751a6f5eb65a0326e07de7c12 16ccce2d0193f958bb3850a833f7ae43 2b65bc5a53975c155aa4bcb4f7b2c4e5 4df16efaf6ddea94e2c50b4cd1dfe060 17e0e9d02900cffe1935e0491d77ffb4 fdf85290fdd893d577b1131a610ef6a5 c32b2ee0293617a37cbb08b847741c3b 8017c25ca9052ca1079d8b78aebd4787 6d330a30f6a8c6d61dd1ab5589329de7 14d19d61370f8149748c72f132f0fc99 f34d766c6938597040d8f9e2bb522ff9 9c63a344d6a2ae8aa8e51b7b90a4a806 105fcbca31506c446151adfeceb51b91 abfe43960977c87471cf9ad4074d30e1 0d6a7f03c63bd5d4317f68ff325ba3bd 80bf4dc8b52a0ba031758022eb025cdd 770b44d6d6cf0670f4e990b22347a7db 848265e3e5eb72dfe8299ad7481a4083 22cac55786e52f633b2fb6b614eaed18 d703dd84045a274ae8bfa73379661388 d6991fe39b0d93debb41700b41f90a15 c4d526250235ddcd6776fc77bc97e7a4 17ebcb31600d01e57f32162a8560cacc 7e27a096d37a1a86952ec71bd89a3e9a 30a2a26162984d7740f81193e8238e61 f6b5b984d4d3dfa033c1bb7e4f0037fe bf406d91c0dccf32acf423cfa1e70710 10d3f270121b493ce85054ef58bada42 310138fe081adb04e2bd901f2f13458b 3d6758158197107c14ebb193230cd115 7380aa79cae1374a7c1e5bbcb80ee23e 06ebfde206bfb0fcbc0edc4ebec30966 1bdd908d532eb0c6adc38b7ca7331dce 8dfce39ab71e7c32d318d136b6100671 a1ae6a6600e3899f31f0eed19e3417d1 34b90c9058f8632c798d4490da498730 7cba922d61c39805d072b589bd52fdf1 e86215c2d54e6670e07383a27bbffb5a ddf47d66aa85a0c6f9f32e59d85a44dd 5d3b22dc2be80919b490437ae4f36a0a e55edf1d0b5cb4e9a3ecabee93dfc6e3 8d209d0fa6536d27a5d6fbb17641cde2 7525d61093f1b28072d111b2b4ae5f89 d5974ee12e5cf7d5da4d6a31123041f3 3e61407e76cffcdcfd7e19ba58cf4b53 6f4c4938ae79324dc402894b44faf8af bab35282ab659d13c93f70412e85cb19 9a37ddec600545473cfb5a05e08d0b20 9973b2172b4d21fb69745a262ccde96b a18b2faa745b6fe189cf772a9f84cbfc"),
-		),
-	)
-
-	DescribeTable("encrypts the server's Initial",
-		func(v protocol.Version, header, data, expectedSample, expectedHdr, expectedPacket []byte) {
-			sealer, _ := NewInitialAEAD(connID, protocol.PerspectiveServer, v)
-			sealed := sealer.Seal(nil, data, 1, header)
-			sample := sealed[2 : 2+16]
-			Expect(sample).To(Equal(expectedSample))
-			sealer.EncryptHeader(sample, &header[0], header[len(header)-2:])
-			Expect(header).To(Equal(expectedHdr))
-			packet := append(header, sealed...)
-			Expect(packet).To(Equal(expectedPacket))
+		{
+			name:           "QUIC v2",
+			version:        protocol.Version2,
+			header:         splitHexString(t, "d16b3343cf0008f067a5502a4262b50040750001"),
+			data:           splitHexString(t, "02000000000600405a020000560303ee fce7f7b37ba1d1632e96677825ddf739 88cfc79825df566dc5430b9a045a1200 130100002e00330024001d00209d3c94 0d89690b84d08a60993c144eca684d10 81287c834d5311bcf32bb9da1a002b00 020304"),
+			expectedSample: splitHexString(t, "6f05d8a4398c47089698baeea26b91eb"),
+			expectedHdr:    splitHexString(t, "dc6b3343cf0008f067a5502a4262b5004075d92f"),
+			expectedPacket: splitHexString(t, "dc6b3343cf0008f067a5502a4262b500 4075d92faaf16f05d8a4398c47089698 baeea26b91eb761d9b89237bbf872630 17915358230035f7fd3945d88965cf17 f9af6e16886c61bfc703106fbaf3cb4c fa52382dd16a393e42757507698075b2 c984c707f0a0812d8cd5a6881eaf21ce da98f4bd23f6fe1a3e2c43edd9ce7ca8 4bed8521e2e140"),
 		},
-		Entry("QUIC v1",
-			protocol.Version1,
-			splitHexString("c1000000010008f067a5502a4262b50040750001"),
-			splitHexString("02000000000600405a020000560303ee fce7f7b37ba1d1632e96677825ddf739 88cfc79825df566dc5430b9a045a1200 130100002e00330024001d00209d3c94 0d89690b84d08a60993c144eca684d10 81287c834d5311bcf32bb9da1a002b00 020304"),
-			splitHexString("2cd0991cd25b0aac406a5816b6394100"),
-			splitHexString("cf000000010008f067a5502a4262b5004075c0d9"),
-			splitHexString("cf000000010008f067a5502a4262b500 4075c0d95a482cd0991cd25b0aac406a 5816b6394100f37a1c69797554780bb3 8cc5a99f5ede4cf73c3ec2493a1839b3 dbcba3f6ea46c5b7684df3548e7ddeb9 c3bf9c73cc3f3bded74b562bfb19fb84 022f8ef4cdd93795d77d06edbb7aaf2f 58891850abbdca3d20398c276456cbc4 2158407dd074ee"),
-		),
-		Entry("QUIC v2",
-			protocol.Version2,
-			splitHexString("d16b3343cf0008f067a5502a4262b50040750001"),
-			splitHexString("02000000000600405a020000560303ee fce7f7b37ba1d1632e96677825ddf739 88cfc79825df566dc5430b9a045a1200 130100002e00330024001d00209d3c94 0d89690b84d08a60993c144eca684d10 81287c834d5311bcf32bb9da1a002b00 020304"),
-			splitHexString("6f05d8a4398c47089698baeea26b91eb"),
-			splitHexString("dc6b3343cf0008f067a5502a4262b5004075d92f"),
-			splitHexString("dc6b3343cf0008f067a5502a4262b500 4075d92faaf16f05d8a4398c47089698 baeea26b91eb761d9b89237bbf872630 17915358230035f7fd3945d88965cf17 f9af6e16886c61bfc703106fbaf3cb4c fa52382dd16a393e42757507698075b2 c984c707f0a0812d8cd5a6881eaf21ce da98f4bd23f6fe1a3e2c43edd9ce7ca8 4bed8521e2e140"),
-		),
-	)
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			sealer, _ := NewInitialAEAD(connID, protocol.PerspectiveServer, tc.version)
+			sealed := sealer.Seal(nil, tc.data, 1, tc.header)
+			sample := sealed[2 : 2+16]
+			require.Equal(t, tc.expectedSample, sample)
+			sealer.EncryptHeader(sample, &tc.header[0], tc.header[len(tc.header)-2:])
+			require.Equal(t, tc.expectedHdr, tc.header)
+			packet := append(tc.header, sealed...)
+			require.Equal(t, tc.expectedPacket, packet)
+		})
+	}
+}
 
+func TestInitialAEADSealsAndOpens(t *testing.T) {
 	for _, ver := range []protocol.Version{protocol.Version1, protocol.Version2} {
-		v := ver
-
-		Context(fmt.Sprintf("using version %s", v), func() {
-			It("seals and opens", func() {
-				connectionID := protocol.ParseConnectionID([]byte{0x12, 0x34, 0x56, 0x78, 0x90, 0xab, 0xcd, 0xef})
-				clientSealer, clientOpener := NewInitialAEAD(connectionID, protocol.PerspectiveClient, v)
-				serverSealer, serverOpener := NewInitialAEAD(connectionID, protocol.PerspectiveServer, v)
-
-				clientMessage := clientSealer.Seal(nil, []byte("foobar"), 42, []byte("aad"))
-				m, err := serverOpener.Open(nil, clientMessage, 42, []byte("aad"))
-				Expect(err).ToNot(HaveOccurred())
-				Expect(m).To(Equal([]byte("foobar")))
-				serverMessage := serverSealer.Seal(nil, []byte("raboof"), 99, []byte("daa"))
-				m, err = clientOpener.Open(nil, serverMessage, 99, []byte("daa"))
-				Expect(err).ToNot(HaveOccurred())
-				Expect(m).To(Equal([]byte("raboof")))
-			})
-
-			It("doesn't work if initialized with different connection IDs", func() {
-				c1 := protocol.ParseConnectionID([]byte{0, 0, 0, 0, 0, 0, 0, 1})
-				c2 := protocol.ParseConnectionID([]byte{0, 0, 0, 0, 0, 0, 0, 2})
-				clientSealer, _ := NewInitialAEAD(c1, protocol.PerspectiveClient, v)
-				_, serverOpener := NewInitialAEAD(c2, protocol.PerspectiveServer, v)
-
-				clientMessage := clientSealer.Seal(nil, []byte("foobar"), 42, []byte("aad"))
-				_, err := serverOpener.Open(nil, clientMessage, 42, []byte("aad"))
-				Expect(err).To(MatchError(ErrDecryptionFailed))
-			})
-
-			It("encrypts und decrypts the header", func() {
-				connID := protocol.ParseConnectionID([]byte{0xde, 0xca, 0xfb, 0xad})
-				clientSealer, clientOpener := NewInitialAEAD(connID, protocol.PerspectiveClient, v)
-				serverSealer, serverOpener := NewInitialAEAD(connID, protocol.PerspectiveServer, v)
-
-				// the first byte and the last 4 bytes should be encrypted
-				header := []byte{0x5e, 0, 1, 2, 3, 4, 0xde, 0xad, 0xbe, 0xef}
-				sample := make([]byte, 16)
-				rand.Read(sample)
-				clientSealer.EncryptHeader(sample, &header[0], header[6:10])
-				// only the last 4 bits of the first byte are encrypted. Check that the first 4 bits are unmodified
-				Expect(header[0] & 0xf0).To(Equal(byte(0x5e & 0xf0)))
-				Expect(header[1:6]).To(Equal([]byte{0, 1, 2, 3, 4}))
-				Expect(header[6:10]).ToNot(Equal([]byte{0xde, 0xad, 0xbe, 0xef}))
-				serverOpener.DecryptHeader(sample, &header[0], header[6:10])
-				Expect(header[0]).To(Equal(byte(0x5e)))
-				Expect(header[1:6]).To(Equal([]byte{0, 1, 2, 3, 4}))
-				Expect(header[6:10]).To(Equal([]byte{0xde, 0xad, 0xbe, 0xef}))
-
-				serverSealer.EncryptHeader(sample, &header[0], header[6:10])
-				// only the last 4 bits of the first byte are encrypted. Check that the first 4 bits are unmodified
-				Expect(header[0] & 0xf0).To(Equal(byte(0x5e & 0xf0)))
-				Expect(header[1:6]).To(Equal([]byte{0, 1, 2, 3, 4}))
-				Expect(header[6:10]).ToNot(Equal([]byte{0xde, 0xad, 0xbe, 0xef}))
-				clientOpener.DecryptHeader(sample, &header[0], header[6:10])
-				Expect(header[0]).To(Equal(byte(0x5e)))
-				Expect(header[1:6]).To(Equal([]byte{0, 1, 2, 3, 4}))
-				Expect(header[6:10]).To(Equal([]byte{0xde, 0xad, 0xbe, 0xef}))
-			})
+		t.Run(ver.String(), func(t *testing.T) {
+			connectionID := protocol.ParseConnectionID([]byte{0x12, 0x34, 0x56, 0x78, 0x90, 0xab, 0xcd, 0xef})
+			clientSealer, clientOpener := NewInitialAEAD(connectionID, protocol.PerspectiveClient, ver)
+			serverSealer, serverOpener := NewInitialAEAD(connectionID, protocol.PerspectiveServer, ver)
+
+			clientMessage := clientSealer.Seal(nil, []byte("foobar"), 42, []byte("aad"))
+			m, err := serverOpener.Open(nil, clientMessage, 42, []byte("aad"))
+			require.NoError(t, err)
+			require.Equal(t, []byte("foobar"), m)
+			serverMessage := serverSealer.Seal(nil, []byte("raboof"), 99, []byte("daa"))
+			m, err = clientOpener.Open(nil, serverMessage, 99, []byte("daa"))
+			require.NoError(t, err)
+			require.Equal(t, []byte("raboof"), m)
 		})
 	}
-})
+}
+
+func TestInitialAEADFailsWithDifferentConnectionIDs(t *testing.T) {
+	for _, ver := range []protocol.Version{protocol.Version1, protocol.Version2} {
+		t.Run(ver.String(), func(t *testing.T) {
+			c1 := protocol.ParseConnectionID([]byte{0, 0, 0, 0, 0, 0, 0, 1})
+			c2 := protocol.ParseConnectionID([]byte{0, 0, 0, 0, 0, 0, 0, 2})
+			clientSealer, _ := NewInitialAEAD(c1, protocol.PerspectiveClient, ver)
+			_, serverOpener := NewInitialAEAD(c2, protocol.PerspectiveServer, ver)
+
+			clientMessage := clientSealer.Seal(nil, []byte("foobar"), 42, []byte("aad"))
+			_, err := serverOpener.Open(nil, clientMessage, 42, []byte("aad"))
+			require.Equal(t, ErrDecryptionFailed, err)
+		})
+	}
+}
+
+func TestInitialAEADEncryptsAndDecryptsHeader(t *testing.T) {
+	for _, ver := range []protocol.Version{protocol.Version1, protocol.Version2} {
+		t.Run(ver.String(), func(t *testing.T) {
+			connID := protocol.ParseConnectionID([]byte{0xde, 0xca, 0xfb, 0xad})
+			clientSealer, clientOpener := NewInitialAEAD(connID, protocol.PerspectiveClient, ver)
+			serverSealer, serverOpener := NewInitialAEAD(connID, protocol.PerspectiveServer, ver)
+
+			header := []byte{0x5e, 0, 1, 2, 3, 4, 0xde, 0xad, 0xbe, 0xef}
+			sample := make([]byte, 16)
+			rand.Read(sample)
+			clientSealer.EncryptHeader(sample, &header[0], header[6:10])
+			require.Equal(t, byte(0x5e&0xf0), header[0]&0xf0)
+			require.Equal(t, []byte{0, 1, 2, 3, 4}, header[1:6])
+			require.NotEqual(t, []byte{0xde, 0xad, 0xbe, 0xef}, header[6:10])
+			serverOpener.DecryptHeader(sample, &header[0], header[6:10])
+			require.Equal(t, byte(0x5e), header[0])
+			require.Equal(t, []byte{0, 1, 2, 3, 4}, header[1:6])
+			require.Equal(t, []byte{0xde, 0xad, 0xbe, 0xef}, header[6:10])
+
+			serverSealer.EncryptHeader(sample, &header[0], header[6:10])
+			require.Equal(t, byte(0x5e&0xf0), header[0]&0xf0)
+			require.Equal(t, []byte{0, 1, 2, 3, 4}, header[1:6])
+			require.NotEqual(t, []byte{0xde, 0xad, 0xbe, 0xef}, header[6:10])
+			clientOpener.DecryptHeader(sample, &header[0], header[6:10])
+			require.Equal(t, byte(0x5e), header[0])
+			require.Equal(t, []byte{0, 1, 2, 3, 4}, header[1:6])
+			require.Equal(t, []byte{0xde, 0xad, 0xbe, 0xef}, header[6:10])
+		})
+	}
+}
 
 func BenchmarkInitialAEADCreate(b *testing.B) {
 	b.ReportAllocs()
diff --git a/internal/handshake/retry_test.go b/internal/handshake/retry_test.go
index c30f94e7..6ccbb473 100644
--- a/internal/handshake/retry_test.go
+++ b/internal/handshake/retry_test.go
@@ -2,45 +2,55 @@ package handshake
 
 import (
 	"encoding/binary"
+	"testing"
 
 	"github.com/quic-go/quic-go/internal/protocol"
 
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/require"
 )
 
-var _ = Describe("Retry Integrity Check", func() {
-	It("calculates retry integrity tags", func() {
-		connID := protocol.ParseConnectionID([]byte{1, 2, 3, 4})
-		fooTag := GetRetryIntegrityTag([]byte("foo"), connID, protocol.Version1)
-		barTag := GetRetryIntegrityTag([]byte("bar"), connID, protocol.Version1)
-		Expect(fooTag).ToNot(BeNil())
-		Expect(barTag).ToNot(BeNil())
-		Expect(*fooTag).ToNot(Equal(*barTag))
-	})
+func TestRetryIntegrityTagCalculation(t *testing.T) {
+	connID := protocol.ParseConnectionID([]byte{1, 2, 3, 4})
+	fooTag := GetRetryIntegrityTag([]byte("foo"), connID, protocol.Version1)
+	barTag := GetRetryIntegrityTag([]byte("bar"), connID, protocol.Version1)
+	require.NotNil(t, fooTag)
+	require.NotNil(t, barTag)
+	require.NotEqual(t, *fooTag, *barTag)
+}
 
-	It("includes the original connection ID in the tag calculation", func() {
-		connID1 := protocol.ParseConnectionID([]byte{1, 2, 3, 4})
-		connID2 := protocol.ParseConnectionID([]byte{4, 3, 2, 1})
-		t1 := GetRetryIntegrityTag([]byte("foobar"), connID1, protocol.Version1)
-		t2 := GetRetryIntegrityTag([]byte("foobar"), connID2, protocol.Version1)
-		Expect(*t1).ToNot(Equal(*t2))
-	})
+func TestRetryIntegrityTagWithDifferentConnectionIDs(t *testing.T) {
+	connID1 := protocol.ParseConnectionID([]byte{1, 2, 3, 4})
+	connID2 := protocol.ParseConnectionID([]byte{4, 3, 2, 1})
+	t1 := GetRetryIntegrityTag([]byte("foobar"), connID1, protocol.Version1)
+	t2 := GetRetryIntegrityTag([]byte("foobar"), connID2, protocol.Version1)
+	require.NotEqual(t, *t1, *t2)
+}
 
-	DescribeTable("using the test vectors",
-		func(version protocol.Version, data []byte) {
-			v := binary.BigEndian.Uint32(data[1:5])
-			Expect(protocol.Version(v)).To(Equal(version))
-			connID := protocol.ParseConnectionID(splitHexString("0x8394c8f03e515708"))
-			Expect(GetRetryIntegrityTag(data[:len(data)-16], connID, version)[:]).To(Equal(data[len(data)-16:]))
+func TestRetryIntegrityTagWithTestVectors(t *testing.T) {
+	tests := []struct {
+		name    string
+		version protocol.Version
+		data    []byte
+	}{
+		{
+			name:    "v1",
+			version: protocol.Version1,
+			data:    splitHexString(t, "ff000000010008f067a5502a4262b574 6f6b656e04a265ba2eff4d829058fb3f 0f2496ba"),
 		},
-		Entry("v1",
-			protocol.Version1,
-			splitHexString("ff000000010008f067a5502a4262b574 6f6b656e04a265ba2eff4d829058fb3f 0f2496ba"),
-		),
-		Entry("v2",
-			protocol.Version2,
-			splitHexString("cf6b3343cf0008f067a5502a4262b574 6f6b656ec8646ce8bfe33952d9555436 65dcc7b6"),
-		),
-	)
-})
+		{
+			name:    "v2",
+			version: protocol.Version2,
+			data:    splitHexString(t, "cf6b3343cf0008f067a5502a4262b574 6f6b656ec8646ce8bfe33952d9555436 65dcc7b6"),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			v := binary.BigEndian.Uint32(tt.data[1:5])
+			require.Equal(t, tt.version, protocol.Version(v))
+			connID := protocol.ParseConnectionID(splitHexString(t, "0x8394c8f03e515708"))
+			tag := GetRetryIntegrityTag(tt.data[:len(tt.data)-16], connID, tt.version)
+			require.Equal(t, tt.data[len(tt.data)-16:], tag[:])
+		})
+	}
+}
diff --git a/internal/handshake/session_ticket_test.go b/internal/handshake/session_ticket_test.go
index 8e209d3a..9d9741e2 100644
--- a/internal/handshake/session_ticket_test.go
+++ b/internal/handshake/session_ticket_test.go
@@ -1,86 +1,93 @@
 package handshake
 
 import (
+	"testing"
 	"time"
 
 	"github.com/quic-go/quic-go/internal/wire"
 	"github.com/quic-go/quic-go/quicvarint"
 
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/require"
 )
 
-var _ = Describe("Session Ticket", func() {
-	It("marshals and unmarshals a 0-RTT session ticket", func() {
-		ticket := &sessionTicket{
-			Parameters: &wire.TransportParameters{
-				InitialMaxStreamDataBidiLocal:  1,
-				InitialMaxStreamDataBidiRemote: 2,
-				ActiveConnectionIDLimit:        10,
-				MaxDatagramFrameSize:           20,
-			},
-			RTT: 1337 * time.Microsecond,
-		}
-		var t sessionTicket
-		Expect(t.Unmarshal(ticket.Marshal(), true)).To(Succeed())
-		Expect(t.Parameters.InitialMaxStreamDataBidiLocal).To(BeEquivalentTo(1))
-		Expect(t.Parameters.InitialMaxStreamDataBidiRemote).To(BeEquivalentTo(2))
-		Expect(t.Parameters.ActiveConnectionIDLimit).To(BeEquivalentTo(10))
-		Expect(t.Parameters.MaxDatagramFrameSize).To(BeEquivalentTo(20))
-		Expect(t.RTT).To(Equal(1337 * time.Microsecond))
-		// fails to unmarshal the ticket as a non-0-RTT ticket
-		Expect(t.Unmarshal(ticket.Marshal(), false)).To(MatchError("the session ticket has more bytes than expected"))
-	})
+func TestMarshalUnmarshal0RTTSessionTicket(t *testing.T) {
+	ticket := &sessionTicket{
+		Parameters: &wire.TransportParameters{
+			InitialMaxStreamDataBidiLocal:  1,
+			InitialMaxStreamDataBidiRemote: 2,
+			ActiveConnectionIDLimit:        10,
+			MaxDatagramFrameSize:           20,
+		},
+		RTT: 1337 * time.Microsecond,
+	}
+	var t2 sessionTicket
+	require.NoError(t, t2.Unmarshal(ticket.Marshal(), true))
+	require.EqualValues(t, 1, t2.Parameters.InitialMaxStreamDataBidiLocal)
+	require.EqualValues(t, 2, t2.Parameters.InitialMaxStreamDataBidiRemote)
+	require.EqualValues(t, 10, t2.Parameters.ActiveConnectionIDLimit)
+	require.EqualValues(t, 20, t2.Parameters.MaxDatagramFrameSize)
+	require.Equal(t, 1337*time.Microsecond, t2.RTT)
+	// fails to unmarshal the ticket as a non-0-RTT ticket
+	require.Error(t, t2.Unmarshal(ticket.Marshal(), false))
+	require.EqualError(t, t2.Unmarshal(ticket.Marshal(), false), "the session ticket has more bytes than expected")
+}
 
-	It("marshals and unmarshals a non-0-RTT session ticket", func() {
-		ticket := &sessionTicket{
-			RTT: 1337 * time.Microsecond,
-		}
-		var t sessionTicket
-		Expect(t.Unmarshal(ticket.Marshal(), false)).To(Succeed())
-		Expect(t.Parameters).To(BeNil())
-		Expect(t.RTT).To(Equal(1337 * time.Microsecond))
-		// fails to unmarshal the ticket as a 0-RTT ticket
-		Expect(t.Unmarshal(ticket.Marshal(), true)).To(MatchError(ContainSubstring("unmarshaling transport parameters from session ticket failed")))
-	})
+func TestMarshalUnmarshalNon0RTTSessionTicket(t *testing.T) {
+	ticket := &sessionTicket{
+		RTT: 1337 * time.Microsecond,
+	}
+	var t2 sessionTicket
+	require.NoError(t, t2.Unmarshal(ticket.Marshal(), false))
+	require.Nil(t, t2.Parameters)
+	require.Equal(t, 1337*time.Microsecond, t2.RTT)
+	// fails to unmarshal the ticket as a 0-RTT ticket
+	err := t2.Unmarshal(ticket.Marshal(), true)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "unmarshaling transport parameters from session ticket failed")
+}
 
-	It("refuses to unmarshal if the ticket is too short for the revision", func() {
-		Expect((&sessionTicket{}).Unmarshal([]byte{}, true)).To(MatchError("failed to read session ticket revision"))
-		Expect((&sessionTicket{}).Unmarshal([]byte{}, false)).To(MatchError("failed to read session ticket revision"))
-	})
+func TestUnmarshalRefusesTooShortTicket(t *testing.T) {
+	err := (&sessionTicket{}).Unmarshal([]byte{}, true)
+	require.EqualError(t, err, "failed to read session ticket revision")
+	err = (&sessionTicket{}).Unmarshal([]byte{}, false)
+	require.EqualError(t, err, "failed to read session ticket revision")
+}
 
-	It("refuses to unmarshal if the revision doesn't match", func() {
-		b := quicvarint.Append(nil, 1337)
-		Expect((&sessionTicket{}).Unmarshal(b, true)).To(MatchError("unknown session ticket revision: 1337"))
-		Expect((&sessionTicket{}).Unmarshal(b, false)).To(MatchError("unknown session ticket revision: 1337"))
-	})
+func TestUnmarshalRefusesUnknownRevision(t *testing.T) {
+	b := quicvarint.Append(nil, 1337)
+	err := (&sessionTicket{}).Unmarshal(b, true)
+	require.EqualError(t, err, "unknown session ticket revision: 1337")
+	err = (&sessionTicket{}).Unmarshal(b, false)
+	require.EqualError(t, err, "unknown session ticket revision: 1337")
+}
 
-	It("refuses to unmarshal if the RTT cannot be read", func() {
-		b := quicvarint.Append(nil, sessionTicketRevision)
-		Expect((&sessionTicket{}).Unmarshal(b, true)).To(MatchError("failed to read RTT"))
-		Expect((&sessionTicket{}).Unmarshal(b, false)).To(MatchError("failed to read RTT"))
-	})
+func TestUnmarshalRefusesInvalidRTT(t *testing.T) {
+	b := quicvarint.Append(nil, sessionTicketRevision)
+	err := (&sessionTicket{}).Unmarshal(b, true)
+	require.EqualError(t, err, "failed to read RTT")
+	err = (&sessionTicket{}).Unmarshal(b, false)
+	require.EqualError(t, err, "failed to read RTT")
+}
 
-	It("refuses to unmarshal a 0-RTT session ticket if unmarshaling the transport parameters fails", func() {
-		b := quicvarint.Append(nil, sessionTicketRevision)
-		b = append(b, []byte("foobar")...)
-		err := (&sessionTicket{}).Unmarshal(b, true)
-		Expect(err).To(HaveOccurred())
-		Expect(err.Error()).To(ContainSubstring("unmarshaling transport parameters from session ticket failed"))
-	})
+func TestUnmarshal0RTTRefusesInvalidTransportParameters(t *testing.T) {
+	b := quicvarint.Append(nil, sessionTicketRevision)
+	b = append(b, []byte("foobar")...)
+	err := (&sessionTicket{}).Unmarshal(b, true)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "unmarshaling transport parameters from session ticket failed")
+}
 
-	It("refuses to unmarshal if the non-0-RTT session ticket has more bytes than expected", func() {
-		ticket := &sessionTicket{
-			Parameters: &wire.TransportParameters{
-				InitialMaxStreamDataBidiLocal:  1,
-				InitialMaxStreamDataBidiRemote: 2,
-				ActiveConnectionIDLimit:        10,
-				MaxDatagramFrameSize:           20,
-			},
-			RTT: 1234 * time.Microsecond,
-		}
-		err := (&sessionTicket{}).Unmarshal(ticket.Marshal(), false)
-		Expect(err).To(HaveOccurred())
-		Expect(err.Error()).To(ContainSubstring("the session ticket has more bytes than expected"))
-	})
-})
+func TestUnmarshalNon0RTTRefusesExtraBytes(t *testing.T) {
+	ticket := &sessionTicket{
+		Parameters: &wire.TransportParameters{
+			InitialMaxStreamDataBidiLocal:  1,
+			InitialMaxStreamDataBidiRemote: 2,
+			ActiveConnectionIDLimit:        10,
+			MaxDatagramFrameSize:           20,
+		},
+		RTT: 1234 * time.Microsecond,
+	}
+	err := (&sessionTicket{}).Unmarshal(ticket.Marshal(), false)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "the session ticket has more bytes than expected")
+}
diff --git a/internal/handshake/token_generator_test.go b/internal/handshake/token_generator_test.go
index 7abc53b4..c09917bd 100644
--- a/internal/handshake/token_generator_test.go
+++ b/internal/handshake/token_generator_test.go
@@ -4,120 +4,117 @@ import (
 	"crypto/rand"
 	"encoding/asn1"
 	"net"
+	"testing"
 	"time"
 
 	"github.com/quic-go/quic-go/internal/protocol"
 
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/require"
 )
 
-var _ = Describe("Token Generator", func() {
-	var tokenGen *TokenGenerator
-
-	BeforeEach(func() {
-		var key TokenProtectorKey
-		rand.Read(key[:])
-		tokenGen = NewTokenGenerator(key)
-	})
-
-	It("generates a token", func() {
-		ip := net.IPv4(127, 0, 0, 1)
-		token, err := tokenGen.NewRetryToken(&net.UDPAddr{IP: ip, Port: 1337}, protocol.ConnectionID{}, protocol.ConnectionID{})
-		Expect(err).ToNot(HaveOccurred())
-		Expect(token).ToNot(BeEmpty())
-	})
-
-	It("works with nil tokens", func() {
-		token, err := tokenGen.DecodeToken(nil)
-		Expect(err).ToNot(HaveOccurred())
-		Expect(token).To(BeNil())
-	})
-
-	It("accepts a valid token", func() {
-		addr := &net.UDPAddr{IP: net.IPv4(192, 168, 0, 1), Port: 1337}
-		tokenEnc, err := tokenGen.NewRetryToken(addr, protocol.ConnectionID{}, protocol.ConnectionID{})
-		Expect(err).ToNot(HaveOccurred())
-		token, err := tokenGen.DecodeToken(tokenEnc)
-		Expect(err).ToNot(HaveOccurred())
-		Expect(token.ValidateRemoteAddr(addr)).To(BeTrue())
-		Expect(token.ValidateRemoteAddr(&net.UDPAddr{IP: net.IPv4(192, 168, 0, 2), Port: 1337})).To(BeFalse())
-		Expect(token.SentTime).To(BeTemporally("~", time.Now(), 100*time.Millisecond))
-		Expect(token.OriginalDestConnectionID.Len()).To(BeZero())
-		Expect(token.RetrySrcConnectionID.Len()).To(BeZero())
-	})
-
-	It("saves the connection ID", func() {
-		connID1 := protocol.ParseConnectionID([]byte{0xde, 0xad, 0xbe, 0xef})
-		connID2 := protocol.ParseConnectionID([]byte{0xde, 0xad, 0xc0, 0xde})
-		tokenEnc, err := tokenGen.NewRetryToken(&net.UDPAddr{}, connID1, connID2)
-		Expect(err).ToNot(HaveOccurred())
-		token, err := tokenGen.DecodeToken(tokenEnc)
-		Expect(err).ToNot(HaveOccurred())
-		Expect(token.OriginalDestConnectionID).To(Equal(connID1))
-		Expect(token.RetrySrcConnectionID).To(Equal(connID2))
-	})
-
-	It("rejects invalid tokens", func() {
-		_, err := tokenGen.DecodeToken([]byte("invalid token"))
-		Expect(err).To(HaveOccurred())
-	})
-
-	It("rejects tokens that cannot be decoded", func() {
-		token, err := tokenGen.tokenProtector.NewToken([]byte("foobar"))
-		Expect(err).ToNot(HaveOccurred())
-		_, err = tokenGen.DecodeToken(token)
-		Expect(err).To(HaveOccurred())
-	})
-
-	It("rejects tokens that can be decoded, but have additional payload", func() {
-		t, err := asn1.Marshal(token{RemoteAddr: []byte("foobar")})
-		Expect(err).ToNot(HaveOccurred())
-		t = append(t, []byte("rest")...)
-		enc, err := tokenGen.tokenProtector.NewToken(t)
-		Expect(err).ToNot(HaveOccurred())
-		_, err = tokenGen.DecodeToken(enc)
-		Expect(err).To(MatchError("rest when unpacking token: 4"))
-	})
-
-	// we don't generate tokens that have no data, but we should be able to handle them if we receive one for whatever reason
-	It("doesn't panic if a tokens has no data", func() {
-		t, err := asn1.Marshal(token{RemoteAddr: []byte("")})
-		Expect(err).ToNot(HaveOccurred())
-		enc, err := tokenGen.tokenProtector.NewToken(t)
-		Expect(err).ToNot(HaveOccurred())
-		_, err = tokenGen.DecodeToken(enc)
-		Expect(err).ToNot(HaveOccurred())
-	})
-
-	It("works with an IPv6 addresses ", func() {
-		addresses := []string{
-			"2001:db8::68",
-			"2001:0000:4136:e378:8000:63bf:3fff:fdd2",
-			"2001::1",
-			"ff01:0:0:0:0:0:0:2",
-		}
-		for _, addr := range addresses {
-			ip := net.ParseIP(addr)
-			Expect(ip).ToNot(BeNil())
-			raddr := &net.UDPAddr{IP: ip, Port: 1337}
-			tokenEnc, err := tokenGen.NewRetryToken(raddr, protocol.ConnectionID{}, protocol.ConnectionID{})
-			Expect(err).ToNot(HaveOccurred())
-			token, err := tokenGen.DecodeToken(tokenEnc)
-			Expect(err).ToNot(HaveOccurred())
-			Expect(token.ValidateRemoteAddr(raddr)).To(BeTrue())
-			Expect(token.SentTime).To(BeTemporally("~", time.Now(), 100*time.Millisecond))
-		}
-	})
-
-	It("uses the string representation an address that is not a UDP address", func() {
-		raddr := &net.TCPAddr{IP: net.IPv4(192, 168, 13, 37), Port: 1337}
+func newTokenGenerator(t *testing.T) *TokenGenerator {
+	var key TokenProtectorKey
+	_, err := rand.Read(key[:])
+	require.NoError(t, err)
+	return NewTokenGenerator(key)
+}
+
+func TestTokenGeneratorNilTokens(t *testing.T) {
+	tokenGen := newTokenGenerator(t)
+	nilToken, err := tokenGen.DecodeToken(nil)
+	require.NoError(t, err)
+	require.Nil(t, nilToken)
+}
+
+func TestTokenGeneratorValidToken(t *testing.T) {
+	tokenGen := newTokenGenerator(t)
+
+	addr := &net.UDPAddr{IP: net.IPv4(192, 168, 0, 1), Port: 1337}
+	connID1 := protocol.ParseConnectionID([]byte{0xde, 0xad, 0xbe, 0xef})
+	connID2 := protocol.ParseConnectionID([]byte{0xde, 0xad, 0xc0, 0xde})
+	tokenEnc, err := tokenGen.NewRetryToken(addr, connID1, connID2)
+	require.NoError(t, err)
+	decodedToken, err := tokenGen.DecodeToken(tokenEnc)
+	require.NoError(t, err)
+	require.True(t, decodedToken.ValidateRemoteAddr(addr))
+	require.False(t, decodedToken.ValidateRemoteAddr(&net.UDPAddr{IP: net.IPv4(192, 168, 0, 2), Port: 1337}))
+	require.WithinDuration(t, time.Now(), decodedToken.SentTime, 100*time.Millisecond)
+	require.Equal(t, connID1, decodedToken.OriginalDestConnectionID)
+	require.Equal(t, connID2, decodedToken.RetrySrcConnectionID)
+}
+
+func TestTokenGeneratorRejectsInvalidTokens(t *testing.T) {
+	tokenGen := newTokenGenerator(t)
+
+	_, err := tokenGen.DecodeToken([]byte("invalid token"))
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "too short")
+}
+
+func TestTokenGeneratorDecodingFailed(t *testing.T) {
+	tokenGen := newTokenGenerator(t)
+
+	invalidToken, err := tokenGen.tokenProtector.NewToken([]byte("foobar"))
+	require.NoError(t, err)
+	_, err = tokenGen.DecodeToken(invalidToken)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "asn1")
+}
+
+func TestTokenGeneratorAdditionalPayload(t *testing.T) {
+	tokenGen := newTokenGenerator(t)
+
+	tok, err := asn1.Marshal(token{RemoteAddr: []byte("foobar")})
+	require.NoError(t, err)
+	tok = append(tok, []byte("rest")...)
+	enc, err := tokenGen.tokenProtector.NewToken(tok)
+	require.NoError(t, err)
+	_, err = tokenGen.DecodeToken(enc)
+	require.EqualError(t, err, "rest when unpacking token: 4")
+}
+
+func TestTokenGeneratorEmptyTokens(t *testing.T) {
+	tokenGen := newTokenGenerator(t)
+
+	emptyTok, err := asn1.Marshal(token{RemoteAddr: []byte("")})
+	require.NoError(t, err)
+	emptyEnc, err := tokenGen.tokenProtector.NewToken(emptyTok)
+	require.NoError(t, err)
+	_, err = tokenGen.DecodeToken(emptyEnc)
+	require.NoError(t, err)
+}
+
+func TestTokenGeneratorIPv6(t *testing.T) {
+	tokenGen := newTokenGenerator(t)
+
+	addresses := []string{
+		"2001:db8::68",
+		"2001:0000:4136:e378:8000:63bf:3fff:fdd2",
+		"2001::1",
+		"ff01:0:0:0:0:0:0:2",
+	}
+	for _, addr := range addresses {
+		ip := net.ParseIP(addr)
+		require.NotNil(t, ip)
+		raddr := &net.UDPAddr{IP: ip, Port: 1337}
 		tokenEnc, err := tokenGen.NewRetryToken(raddr, protocol.ConnectionID{}, protocol.ConnectionID{})
-		Expect(err).ToNot(HaveOccurred())
+		require.NoError(t, err)
 		token, err := tokenGen.DecodeToken(tokenEnc)
-		Expect(err).ToNot(HaveOccurred())
-		Expect(token.ValidateRemoteAddr(raddr)).To(BeTrue())
-		Expect(token.ValidateRemoteAddr(&net.TCPAddr{IP: net.IPv4(192, 168, 13, 37), Port: 1338})).To(BeFalse())
-		Expect(token.SentTime).To(BeTemporally("~", time.Now(), 100*time.Millisecond))
-	})
-})
+		require.NoError(t, err)
+		require.True(t, token.ValidateRemoteAddr(raddr))
+		require.WithinDuration(t, time.Now(), token.SentTime, 100*time.Millisecond)
+	}
+}
+
+func TestTokenGeneratorNonUDPAddr(t *testing.T) {
+	tokenGen := newTokenGenerator(t)
+
+	raddr := &net.TCPAddr{IP: net.IPv4(192, 168, 13, 37), Port: 1337}
+	tokenEnc, err := tokenGen.NewRetryToken(raddr, protocol.ConnectionID{}, protocol.ConnectionID{})
+	require.NoError(t, err)
+	token, err := tokenGen.DecodeToken(tokenEnc)
+	require.NoError(t, err)
+	require.True(t, token.ValidateRemoteAddr(raddr))
+	require.False(t, token.ValidateRemoteAddr(&net.TCPAddr{IP: net.IPv4(192, 168, 13, 37), Port: 1338}))
+	require.WithinDuration(t, time.Now(), token.SentTime, 100*time.Millisecond)
+}
diff --git a/internal/handshake/token_protector_test.go b/internal/handshake/token_protector_test.go
index fb342f3a..497b4547 100644
--- a/internal/handshake/token_protector_test.go
+++ b/internal/handshake/token_protector_test.go
@@ -2,65 +2,66 @@ package handshake
 
 import (
 	"crypto/rand"
+	"testing"
 
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/require"
 )
 
-var _ = Describe("Token Protector", func() {
-	var tp *tokenProtector
+func TestTokenProtectorEncodeAndDecode(t *testing.T) {
+	var key TokenProtectorKey
+	rand.Read(key[:])
+	tp := newTokenProtector(key)
 
-	BeforeEach(func() {
-		var key TokenProtectorKey
-		rand.Read(key[:])
-		var err error
-		tp = newTokenProtector(key)
-		Expect(err).ToNot(HaveOccurred())
-	})
+	token, err := tp.NewToken([]byte("foobar"))
+	require.NoError(t, err)
+	require.NotContains(t, string(token), "foobar")
 
-	It("encodes and decodes tokens", func() {
-		token, err := tp.NewToken([]byte("foobar"))
-		Expect(err).ToNot(HaveOccurred())
-		Expect(token).ToNot(ContainSubstring("foobar"))
-		decoded, err := tp.DecodeToken(token)
-		Expect(err).ToNot(HaveOccurred())
-		Expect(decoded).To(Equal([]byte("foobar")))
-	})
+	decoded, err := tp.DecodeToken(token)
+	require.NoError(t, err)
+	require.Equal(t, []byte("foobar"), decoded)
+}
 
-	It("uses the different keys", func() {
-		var key1, key2 TokenProtectorKey
-		rand.Read(key1[:])
-		rand.Read(key2[:])
-		tp1 := newTokenProtector(key1)
-		tp2 := newTokenProtector(key2)
-		t1, err := tp1.NewToken([]byte("foo"))
-		Expect(err).ToNot(HaveOccurred())
-		t2, err := tp2.NewToken([]byte("foo"))
-		Expect(err).ToNot(HaveOccurred())
+func TestTokenProtectorDifferentKeys(t *testing.T) {
+	var key1, key2 TokenProtectorKey
+	rand.Read(key1[:])
+	rand.Read(key2[:])
+	tp1 := newTokenProtector(key1)
+	tp2 := newTokenProtector(key2)
 
-		_, err = tp1.DecodeToken(t1)
-		Expect(err).ToNot(HaveOccurred())
-		_, err = tp1.DecodeToken(t2)
-		Expect(err).To(HaveOccurred())
+	t1, err := tp1.NewToken([]byte("foo"))
+	require.NoError(t, err)
+	t2, err := tp2.NewToken([]byte("foo"))
+	require.NoError(t, err)
 
-		// now create another token protector, reusing key1
-		tp3 := newTokenProtector(key1)
-		_, err = tp3.DecodeToken(t1)
-		Expect(err).ToNot(HaveOccurred())
-		_, err = tp3.DecodeToken(t2)
-		Expect(err).To(HaveOccurred())
-	})
+	_, err = tp1.DecodeToken(t1)
+	require.NoError(t, err)
+	_, err = tp1.DecodeToken(t2)
+	require.Error(t, err)
 
-	It("doesn't decode invalid tokens", func() {
-		token, err := tp.NewToken([]byte("foobar"))
-		Expect(err).ToNot(HaveOccurred())
-		_, err = tp.DecodeToken(token[1:]) // the token is invalid without the first byte
-		Expect(err).To(HaveOccurred())
-		Expect(err.Error()).To(ContainSubstring("message authentication failed"))
-	})
+	tp3 := newTokenProtector(key1)
+	_, err = tp3.DecodeToken(t1)
+	require.NoError(t, err)
+	_, err = tp3.DecodeToken(t2)
+	require.Error(t, err)
+}
 
-	It("errors when decoding too short tokens", func() {
-		_, err := tp.DecodeToken([]byte("foobar"))
-		Expect(err).To(MatchError("token too short: 6"))
-	})
-})
+func TestTokenProtectorInvalidTokens(t *testing.T) {
+	var key TokenProtectorKey
+	rand.Read(key[:])
+	tp := newTokenProtector(key)
+
+	token, err := tp.NewToken([]byte("foobar"))
+	require.NoError(t, err)
+	_, err = tp.DecodeToken(token[1:])
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "message authentication failed")
+}
+
+func TestTokenProtectorTooShortTokens(t *testing.T) {
+	var key TokenProtectorKey
+	rand.Read(key[:])
+	tp := newTokenProtector(key)
+
+	_, err := tp.DecodeToken([]byte("foobar"))
+	require.EqualError(t, err, "token too short: 6")
+}
diff --git a/internal/handshake/updatable_aead_test.go b/internal/handshake/updatable_aead_test.go
index 17cbe018..b5053f02 100644
--- a/internal/handshake/updatable_aead_test.go
+++ b/internal/handshake/updatable_aead_test.go
@@ -4,6 +4,7 @@ import (
 	"crypto/rand"
 	"crypto/tls"
 	"fmt"
+	mrand "math/rand/v2"
 	"testing"
 	"time"
 
@@ -11,531 +12,574 @@ import (
 	"github.com/quic-go/quic-go/internal/protocol"
 	"github.com/quic-go/quic-go/internal/qerr"
 	"github.com/quic-go/quic-go/internal/utils"
-	"github.com/quic-go/quic-go/logging"
-
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
 	"go.uber.org/mock/gomock"
+
+	"github.com/stretchr/testify/require"
+)
+
+const (
+	msg = "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua."
+	ad  = "Donec in velit neque."
 )
 
-var _ = Describe("Updatable AEAD", func() {
-	DescribeTable("ChaCha test vector",
-		func(v protocol.Version, expectedPayload, expectedPacket []byte) {
-			secret := splitHexString("9ac312a7f877468ebe69422748ad00a1 5443f18203a07d6060f688f30f21632b")
-			aead := newUpdatableAEAD(&utils.RTTStats{}, nil, nil, v)
+func randomCipherSuite() *cipherSuite { return cipherSuites[mrand.IntN(len(cipherSuites))] }
+
+func setupEndpoints(t *testing.T, serverRTTStats *utils.RTTStats) (client, server *updatableAEAD, serverTracer *mocklogging.MockConnectionTracer) {
+	cs := randomCipherSuite()
+	mockCtrl := gomock.NewController(t)
+	tr, serverTracer := mocklogging.NewMockConnectionTracer(mockCtrl)
+
+	trafficSecret1 := make([]byte, 16)
+	trafficSecret2 := make([]byte, 16)
+	rand.Read(trafficSecret1)
+	rand.Read(trafficSecret2)
+
+	client = newUpdatableAEAD(&utils.RTTStats{}, nil, utils.DefaultLogger, protocol.Version1)
+	server = newUpdatableAEAD(serverRTTStats, tr, utils.DefaultLogger, protocol.Version1)
+	client.SetReadKey(cs, trafficSecret2)
+	client.SetWriteKey(cs, trafficSecret1)
+	server.SetReadKey(cs, trafficSecret1)
+	server.SetWriteKey(cs, trafficSecret2)
+	return client, server, serverTracer
+}
+
+func TestChaChaTestVector(t *testing.T) {
+	testCases := []struct {
+		name            string
+		version         protocol.Version
+		expectedPayload []byte
+		expectedPacket  []byte
+	}{
+		{
+			version:         protocol.Version1,
+			expectedPayload: splitHexString(t, "655e5cd55c41f69080575d7999c25a5bfb"),
+			expectedPacket:  splitHexString(t, "4cfe4189655e5cd55c41f69080575d7999c25a5bfb"),
+		},
+		{
+			version:         protocol.Version2,
+			expectedPayload: splitHexString(t, "0ae7b6b932bc27d786f4bc2bb20f2162ba"),
+			expectedPacket:  splitHexString(t, "5558b1c60ae7b6b932bc27d786f4bc2bb20f2162ba"),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("QUIC %s", tc.version), func(t *testing.T) {
+			secret := splitHexString(t, "9ac312a7f877468ebe69422748ad00a1 5443f18203a07d6060f688f30f21632b")
+			aead := newUpdatableAEAD(&utils.RTTStats{}, nil, nil, tc.version)
 			chacha := cipherSuites[2]
-			Expect(chacha.ID).To(Equal(tls.TLS_CHACHA20_POLY1305_SHA256))
+			require.Equal(t, tls.TLS_CHACHA20_POLY1305_SHA256, chacha.ID)
 			aead.SetWriteKey(chacha, secret)
 			const pnOffset = 1
-			header := splitHexString("4200bff4")
+			header := splitHexString(t, "4200bff4")
 			payloadOffset := len(header)
-			plaintext := splitHexString("01")
+			plaintext := splitHexString(t, "01")
 			payload := aead.Seal(nil, plaintext, 654360564, header)
-			Expect(payload).To(Equal(expectedPayload))
+			require.Equal(t, tc.expectedPayload, payload)
 			packet := append(header, payload...)
 			aead.EncryptHeader(packet[pnOffset+4:pnOffset+4+16], &packet[0], packet[pnOffset:payloadOffset])
-			Expect(packet).To(Equal(expectedPacket))
-		},
-		Entry("QUIC v1",
-			protocol.Version1,
-			splitHexString("655e5cd55c41f69080575d7999c25a5bfb"),
-			splitHexString("4cfe4189655e5cd55c41f69080575d7999c25a5bfb"),
-		),
-		Entry("QUIC v2",
-			protocol.Version2,
-			splitHexString("0ae7b6b932bc27d786f4bc2bb20f2162ba"),
-			splitHexString("5558b1c60ae7b6b932bc27d786f4bc2bb20f2162ba"),
-		),
+			require.Equal(t, tc.expectedPacket, packet)
+		})
+	}
+}
+
+func TestUpdatableAEADHeaderProtection(t *testing.T) {
+	for _, v := range []protocol.Version{protocol.Version1, protocol.Version2} {
+		for _, cs := range cipherSuites {
+			t.Run(fmt.Sprintf("QUIC %s/%s", v, tls.CipherSuiteName(cs.ID)), func(t *testing.T) {
+				trafficSecret1 := make([]byte, 16)
+				trafficSecret2 := make([]byte, 16)
+				rand.Read(trafficSecret1)
+				rand.Read(trafficSecret2)
+
+				client := newUpdatableAEAD(&utils.RTTStats{}, nil, utils.DefaultLogger, v)
+				server := newUpdatableAEAD(&utils.RTTStats{}, nil, utils.DefaultLogger, v)
+				client.SetReadKey(cs, trafficSecret2)
+				client.SetWriteKey(cs, trafficSecret1)
+				server.SetReadKey(cs, trafficSecret1)
+				server.SetWriteKey(cs, trafficSecret2)
+
+				var lastFiveBitsDifferent int
+				for i := 0; i < 100; i++ {
+					sample := make([]byte, 16)
+					rand.Read(sample)
+					header := []byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}
+					client.EncryptHeader(sample, &header[0], header[9:13])
+					if header[0]&0x1f != 0xb5&0x1f {
+						lastFiveBitsDifferent++
+					}
+					require.Equal(t, byte(0xb5&0xe0), header[0]&0xe0)
+					require.Equal(t, []byte{1, 2, 3, 4, 5, 6, 7, 8}, header[1:9])
+					require.NotEqual(t, []byte{0xde, 0xad, 0xbe, 0xef}, header[9:13])
+					server.DecryptHeader(sample, &header[0], header[9:13])
+					require.Equal(t, []byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}, header)
+				}
+				require.Greater(t, lastFiveBitsDifferent, 75)
+			})
+		}
+	}
+}
+
+func TestUpdatableAEADEncryptDecryptMessage(t *testing.T) {
+	for _, v := range []protocol.Version{protocol.Version1, protocol.Version2} {
+		for _, cs := range cipherSuites {
+			t.Run(fmt.Sprintf("QUIC %s/%s", v, tls.CipherSuiteName(cs.ID)), func(t *testing.T) {
+				rttStats := utils.RTTStats{}
+				trafficSecret1 := make([]byte, 16)
+				trafficSecret2 := make([]byte, 16)
+				rand.Read(trafficSecret1)
+				rand.Read(trafficSecret2)
+
+				client := newUpdatableAEAD(&rttStats, nil, utils.DefaultLogger, v)
+				server := newUpdatableAEAD(&rttStats, nil, utils.DefaultLogger, v)
+				client.SetReadKey(cs, trafficSecret2)
+				client.SetWriteKey(cs, trafficSecret1)
+				server.SetReadKey(cs, trafficSecret1)
+				server.SetWriteKey(cs, trafficSecret2)
+
+				msg := []byte("Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.")
+				ad := []byte("Donec in velit neque.")
+
+				encrypted := server.Seal(nil, msg, 0x1337, ad)
+
+				opened, err := client.Open(nil, encrypted, time.Now(), 0x1337, protocol.KeyPhaseZero, ad)
+				require.NoError(t, err)
+				require.Equal(t, msg, opened)
+
+				_, err = client.Open(nil, encrypted, time.Now(), 0x1337, protocol.KeyPhaseZero, []byte("wrong ad"))
+				require.Equal(t, ErrDecryptionFailed, err)
+
+				_, err = client.Open(nil, encrypted, time.Now(), 0x42, protocol.KeyPhaseZero, ad)
+				require.Equal(t, ErrDecryptionFailed, err)
+			})
+		}
+	}
+}
+
+func TestUpdatableAEADPacketNumbers(t *testing.T) {
+	client, server, _ := setupEndpoints(t, &utils.RTTStats{})
+	msg := []byte("Lorem ipsum")
+	ad := []byte("Donec in velit neque.")
+
+	encrypted := server.Seal(nil, msg, 0x1337, ad)
+	require.Equal(t, protocol.PacketNumber(0x1337), server.FirstPacketNumber()) // make sure we save the first packet number
+	_ = server.Seal(nil, msg, 0x1338, ad)
+	require.Equal(t, protocol.PacketNumber(0x1337), server.FirstPacketNumber()) // make sure we save the first packet number
+
+	// check that decoding the packet number works as expected
+	_, err := client.Open(nil, encrypted[:len(encrypted)-1], time.Now(), 0x1337, protocol.KeyPhaseZero, ad)
+	require.Error(t, err)
+	require.Equal(t, protocol.PacketNumber(0x38), client.DecodePacketNumber(0x38, protocol.PacketNumberLen1))
+
+	_, err = client.Open(nil, encrypted, time.Now(), 0x1337, protocol.KeyPhaseZero, ad)
+	require.NoError(t, err)
+	require.Equal(t, protocol.PacketNumber(0x1338), client.DecodePacketNumber(0x38, protocol.PacketNumberLen1))
+}
+
+func TestAEADLimitReached(t *testing.T) {
+	client, _, _ := setupEndpoints(t, &utils.RTTStats{})
+	client.invalidPacketLimit = 10
+	for i := 0; i < 9; i++ {
+		_, err := client.Open(nil, []byte("foobar"), time.Now(), protocol.PacketNumber(i), protocol.KeyPhaseZero, []byte("ad"))
+		require.Equal(t, ErrDecryptionFailed, err)
+	}
+	_, err := client.Open(nil, []byte("foobar"), time.Now(), 10, protocol.KeyPhaseZero, []byte("ad"))
+	require.Error(t, err)
+	var transportErr *qerr.TransportError
+	require.ErrorAs(t, err, &transportErr)
+	require.Equal(t, qerr.AEADLimitReached, transportErr.ErrorCode)
+}
+
+func TestKeyUpdates(t *testing.T) {
+	client, server, _ := setupEndpoints(t, &utils.RTTStats{})
+
+	now := time.Now()
+	require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+	encrypted0 := server.Seal(nil, []byte(msg), 0x1337, []byte(ad))
+	server.rollKeys()
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+	encrypted1 := server.Seal(nil, []byte(msg), 0x1337, []byte(ad))
+	require.NotEqual(t, encrypted0, encrypted1)
+
+	_, err := client.Open(nil, encrypted1, now, 0x1337, protocol.KeyPhaseZero, []byte(ad))
+	require.Equal(t, ErrDecryptionFailed, err)
+
+	client.rollKeys()
+	decrypted, err := client.Open(nil, encrypted1, now, 0x1337, protocol.KeyPhaseOne, []byte(ad))
+	require.NoError(t, err)
+	require.Equal(t, msg, string(decrypted))
+}
+
+// func TestUpdatesKeysWhenReceivingPacketWithNextKeyPhase(t *testing.T) {
+// 	rttStats := utils.RTTStats{}
+// 	mockCtrl := gomock.NewController(t)
+// 	serverTracer := mocklogging.NewMockConnectionTracer(mockCtrl)
+
+// 	trafficSecret1 := make([]byte, 16)
+// 	trafficSecret2 := make([]byte, 16)
+// 	rand.Read(trafficSecret1)
+// 	rand.Read(trafficSecret2)
+
+// 	client := newUpdatableAEAD(&rttStats, nil, utils.DefaultLogger, protocol.Version1)
+// 	server := newUpdatableAEAD(&rttStats, serverTracer, utils.DefaultLogger, protocol.Version1)
+// 	client.SetReadKey(cs, trafficSecret2)
+// 	client.SetWriteKey(cs, trafficSecret1)
+// 	server.SetReadKey(cs, trafficSecret1)
+// 	server.SetWriteKey(cs, trafficSecret2)
+
+// 	now := time.Now()
+// 	encrypted0 := client.Seal(nil, []byte(msg), 0x42, ad)
+// 	decrypted, err := server.Open(nil, encrypted0, now, 0x42, protocol.KeyPhaseZero, ad)
+// 	require.NoError(t, err)
+// 	require.Equal(t, msg, decrypted)
+
+// 	require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+// 	_ = server.Seal(nil, msg, 0x1, ad)
+
+// 	client.rollKeys()
+// 	encrypted1 := client.Seal(nil, msg, 0x43, ad)
+// 	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), true)
+// 	decrypted, err = server.Open(nil, encrypted1, now, 0x43, protocol.KeyPhaseOne, ad)
+// 	require.NoError(t, err)
+// 	require.Equal(t, msg, decrypted)
+// 	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+// }
+
+func TestReorderedPacketAfterKeyUpdate(t *testing.T) {
+	client, server, serverTracer := setupEndpoints(t, &utils.RTTStats{})
+
+	now := time.Now()
+	encrypted01 := client.Seal(nil, []byte(msg), 0x42, []byte(ad))
+	encrypted02 := client.Seal(nil, []byte(msg), 0x43, []byte(ad))
+	_, err := server.Open(nil, encrypted01, now, 0x42, protocol.KeyPhaseZero, []byte(ad))
+	require.NoError(t, err)
+	_ = server.Seal(nil, []byte(msg), 0x1, []byte(ad))
+
+	client.rollKeys()
+	encrypted1 := client.Seal(nil, []byte(msg), 0x44, []byte(ad))
+	require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), true)
+	_, err = server.Open(nil, encrypted1, now, 0x44, protocol.KeyPhaseOne, []byte(ad))
+	require.NoError(t, err)
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+
+	// now receive a reordered packet
+	decrypted, err := server.Open(nil, encrypted02, now, 0x43, protocol.KeyPhaseZero, []byte(ad))
+	require.NoError(t, err)
+	require.Equal(t, msg, string(decrypted))
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+}
+
+func TestDropsKeys3PTOsAfterKeyUpdate(t *testing.T) {
+	var rttStats utils.RTTStats
+	client, server, serverTracer := setupEndpoints(t, &rttStats)
+
+	now := time.Now()
+	rttStats.UpdateRTT(10*time.Millisecond, 0, now)
+	pto := rttStats.PTO(true)
+	encrypted01 := client.Seal(nil, []byte(msg), 0x42, []byte(ad))
+	encrypted02 := client.Seal(nil, []byte(msg), 0x43, []byte(ad))
+	_, err := server.Open(nil, encrypted01, now, 0x42, protocol.KeyPhaseZero, []byte(ad))
+	require.NoError(t, err)
+	_ = server.Seal(nil, []byte(msg), 0x1, []byte(ad))
+
+	client.rollKeys()
+	encrypted1 := client.Seal(nil, []byte(msg), 0x44, []byte(ad))
+	require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), true)
+	serverTracer.EXPECT().DroppedKey(protocol.KeyPhase(0))
+	_, err = server.Open(nil, encrypted1, now, 0x44, protocol.KeyPhaseOne, []byte(ad))
+	require.NoError(t, err)
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+
+	// packet arrived too late, the key was already dropped
+	_, err = server.Open(nil, encrypted02, now.Add(3*pto).Add(time.Nanosecond), 0x43, protocol.KeyPhaseZero, []byte(ad))
+	require.Equal(t, ErrKeysDropped, err)
+}
+
+func TestAllowsFirstKeyUpdateImmediately(t *testing.T) {
+	client, server, serverTracer := setupEndpoints(t, &utils.RTTStats{})
+	client.rollKeys()
+	encrypted := client.Seal(nil, []byte(msg), 0x1337, []byte(ad))
+
+	// if decryption failed, we don't expect a key phase update
+	_, err := server.Open(nil, encrypted[:len(encrypted)-1], time.Now(), 0x1337, protocol.KeyPhaseOne, []byte(ad))
+	require.Equal(t, ErrDecryptionFailed, err)
+
+	// the key phase is updated on first successful decryption
+	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), true)
+	_, err = server.Open(nil, encrypted, time.Now(), 0x1337, protocol.KeyPhaseOne, []byte(ad))
+	require.NoError(t, err)
+}
+
+func TestRejectFrequentKeyUpdates(t *testing.T) {
+	client, server, _ := setupEndpoints(t, &utils.RTTStats{})
+
+	server.rollKeys()
+	client.rollKeys()
+	encrypted0 := client.Seal(nil, []byte(msg), 0x42, []byte(ad))
+	_, err := server.Open(nil, encrypted0, time.Now(), 0x42, protocol.KeyPhaseOne, []byte(ad))
+	require.NoError(t, err)
+
+	client.rollKeys()
+	encrypted1 := client.Seal(nil, []byte(msg), 0x42, []byte(ad))
+	_, err = server.Open(nil, encrypted1, time.Now(), 0x42, protocol.KeyPhaseZero, []byte(ad))
+	require.Equal(t, &qerr.TransportError{
+		ErrorCode:    qerr.KeyUpdateError,
+		ErrorMessage: "keys updated too quickly",
+	}, err)
+}
+
+//nolint:unparam
+func setKeyUpdateIntervals(t *testing.T, firstKeyUpdateInterval, keyUpdateInterval uint64) {
+	origKeyUpdateInterval := KeyUpdateInterval
+	origFirstKeyUpdateInterval := FirstKeyUpdateInterval
+	KeyUpdateInterval = keyUpdateInterval
+	FirstKeyUpdateInterval = firstKeyUpdateInterval
+
+	t.Cleanup(func() {
+		KeyUpdateInterval = origKeyUpdateInterval
+		FirstKeyUpdateInterval = origFirstKeyUpdateInterval
+	})
+}
+
+func TestInitiateKeyUpdateAfterSendingMaxPackets(t *testing.T) {
+	const firstKeyUpdateInterval = 5
+	const keyUpdateInterval = 20
+	setKeyUpdateIntervals(t, firstKeyUpdateInterval, keyUpdateInterval)
+
+	client, server, serverTracer := setupEndpoints(t, &utils.RTTStats{})
+	server.SetHandshakeConfirmed()
+
+	var pn protocol.PacketNumber
+	// first key update
+	for i := 0; i < firstKeyUpdateInterval; i++ {
+		require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+		server.Seal(nil, []byte(msg), pn, []byte(ad))
+		pn++
+	}
+	// the first update is allowed without receiving an acknowledgement
+	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+
+	// subsequent key update
+	for i := 0; i < 2*keyUpdateInterval; i++ {
+		require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+		server.Seal(nil, []byte(msg), pn, []byte(ad))
+		pn++
+	}
+	// no update allowed before receiving an acknowledgement for the current key phase
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+	// receive an ACK for a packet sent in key phase 1
+	client.rollKeys()
+	b := client.Seal(nil, []byte("foobar"), 1, []byte("ad"))
+	_, err := server.Open(nil, b, time.Now(), 1, protocol.KeyPhaseOne, []byte("ad"))
+	require.NoError(t, err)
+	require.NoError(t, server.SetLargestAcked(firstKeyUpdateInterval))
+
+	serverTracer.EXPECT().DroppedKey(protocol.KeyPhase(0))
+	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(2), false)
+	require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+}
+
+func TestKeyUpdateEnforceACKKeyPhase(t *testing.T) {
+	const firstKeyUpdateInterval = 5
+	setKeyUpdateIntervals(t, firstKeyUpdateInterval, KeyUpdateInterval)
+
+	_, server, serverTracer := setupEndpoints(t, &utils.RTTStats{})
+	server.SetHandshakeConfirmed()
+
+	// First make sure that we update our keys.
+	for i := 0; i < firstKeyUpdateInterval; i++ {
+		pn := protocol.PacketNumber(i)
+		require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+		server.Seal(nil, []byte(msg), pn, []byte(ad))
+	}
+	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+
+	// Now that our keys are updated, send a packet using the new keys.
+	const nextPN = firstKeyUpdateInterval + 1
+	server.Seal(nil, []byte(msg), nextPN, []byte(ad))
+
+	for i := 0; i < firstKeyUpdateInterval; i++ {
+		// We haven't decrypted any packet in the new key phase yet.
+		// This means that the ACK must have been sent in the old key phase.
+		require.NoError(t, server.SetLargestAcked(protocol.PacketNumber(i)))
+	}
+
+	// We haven't decrypted any packet in the new key phase yet.
+	// This means that the ACK must have been sent in the old key phase.
+	err := server.SetLargestAcked(nextPN)
+	require.Error(t, err)
+	var transportErr *qerr.TransportError
+	require.ErrorAs(t, err, &transportErr)
+	require.Equal(t, qerr.KeyUpdateError, transportErr.ErrorCode)
+	require.Equal(t, "received ACK for key phase 1, but peer didn't update keys", transportErr.ErrorMessage)
+}
+
+func TestKeyUpdateAfterOpeningMaxPackets(t *testing.T) {
+	const firstKeyUpdateInterval = 5
+	const keyUpdateInterval = 20
+	setKeyUpdateIntervals(t, firstKeyUpdateInterval, keyUpdateInterval)
+
+	client, server, serverTracer := setupEndpoints(t, &utils.RTTStats{})
+	server.SetHandshakeConfirmed()
+
+	msg := []byte("message")
+	ad := []byte("additional data")
+
+	// first key update
+	var pn protocol.PacketNumber
+	for i := 0; i < firstKeyUpdateInterval; i++ {
+		require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+		encrypted := client.Seal(nil, msg, pn, ad)
+		_, err := server.Open(nil, encrypted, time.Now(), pn, protocol.KeyPhaseZero, ad)
+		require.NoError(t, err)
+		pn++
+	}
+
+	// the first update is allowed without receiving an acknowledgement
+	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+
+	// subsequent key update
+	client.rollKeys()
+	for i := 0; i < keyUpdateInterval; i++ {
+		require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+		encrypted := client.Seal(nil, msg, pn, ad)
+		_, err := server.Open(nil, encrypted, time.Now(), pn, protocol.KeyPhaseOne, ad)
+		require.NoError(t, err)
+		pn++
+	}
+
+	// No update allowed before receiving an acknowledgement for the current key phase
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+	server.Seal(nil, msg, 1, ad)
+	require.NoError(t, server.SetLargestAcked(firstKeyUpdateInterval+1))
+	serverTracer.EXPECT().DroppedKey(protocol.KeyPhase(0))
+	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(2), false)
+	require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+}
+
+func TestKeyUpdateKeyPhaseSkipping(t *testing.T) {
+	const firstKeyUpdateInterval = 5
+	const keyUpdateInterval = 20
+	setKeyUpdateIntervals(t, firstKeyUpdateInterval, keyUpdateInterval)
+
+	var rttStats utils.RTTStats
+	rttStats.UpdateRTT(10*time.Millisecond, 0, time.Now())
+	client, server, serverTracer := setupEndpoints(t, &rttStats)
+	server.SetHandshakeConfirmed()
+
+	now := time.Now()
+	data1 := client.Seal(nil, []byte(msg), 1, []byte(ad))
+	_, err := server.Open(nil, data1, now, 1, protocol.KeyPhaseZero, []byte(ad))
+	require.NoError(t, err)
+	for i := 0; i < firstKeyUpdateInterval; i++ {
+		pn := protocol.PacketNumber(i)
+		require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+		server.Seal(nil, []byte(msg), pn, []byte(ad))
+		require.NoError(t, server.SetLargestAcked(pn))
+	}
+	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+	// The server never received a packet at key phase 1.
+	// Make sure the key phase 0 is still there at a much later point.
+	data2 := client.Seal(nil, []byte(msg), 1, []byte(ad))
+	_, err = server.Open(nil, data2, now.Add(10*rttStats.PTO(true)), 1, protocol.KeyPhaseZero, []byte(ad))
+	require.NoError(t, err)
+}
+
+func TestFastKeyUpdatesByPeer(t *testing.T) {
+	const firstKeyUpdateInterval = 5
+	const keyUpdateInterval = 20
+	setKeyUpdateIntervals(t, firstKeyUpdateInterval, keyUpdateInterval)
+
+	client, server, serverTracer := setupEndpoints(t, &utils.RTTStats{})
+	server.SetHandshakeConfirmed()
+
+	var pn protocol.PacketNumber
+	for i := 0; i < firstKeyUpdateInterval; i++ {
+		require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+		server.Seal(nil, []byte(msg), pn, []byte(ad))
+		pn++
+	}
+	b := client.Seal(nil, []byte("foobar"), 1, []byte("ad"))
+	_, err := server.Open(nil, b, time.Now(), 1, protocol.KeyPhaseZero, []byte("ad"))
+	require.NoError(t, err)
+	require.NoError(t, server.SetLargestAcked(0))
+	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+
+	// Send and receive an acknowledgement for a packet in key phase 1.
+	// We are now running a timer to drop the keys with 3 PTO.
+	server.Seal(nil, []byte(msg), pn, []byte(ad))
+	client.rollKeys()
+	dataKeyPhaseOne := client.Seal(nil, []byte(msg), 2, []byte(ad))
+	now := time.Now()
+	_, err = server.Open(nil, dataKeyPhaseOne, now, 2, protocol.KeyPhaseOne, []byte(ad))
+	require.NoError(t, err)
+	require.NoError(t, server.SetLargestAcked(pn))
+	// Now the client sends us a packet in key phase 2, forcing us to update keys before the 3 PTO period is over.
+	// This mean that we need to drop the keys for key phase 0 immediately.
+	client.rollKeys()
+	dataKeyPhaseTwo := client.Seal(nil, []byte(msg), 3, []byte(ad))
+	gomock.InOrder(
+		serverTracer.EXPECT().DroppedKey(protocol.KeyPhase(0)),
+		serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(2), true),
 	)
+	_, err = server.Open(nil, dataKeyPhaseTwo, now, 3, protocol.KeyPhaseZero, []byte(ad))
+	require.NoError(t, err)
+	require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+}
 
-	for _, ver := range []protocol.Version{protocol.Version1, protocol.Version2} {
-		v := ver
-
-		Context(fmt.Sprintf("using version %s", v), func() {
-			for i := range cipherSuites {
-				cs := cipherSuites[i]
-
-				Context(fmt.Sprintf("using %s", tls.CipherSuiteName(cs.ID)), func() {
-					var (
-						client, server *updatableAEAD
-						serverTracer   *mocklogging.MockConnectionTracer
-						rttStats       utils.RTTStats
-					)
-
-					BeforeEach(func() {
-						var tr *logging.ConnectionTracer
-						rttStats = utils.RTTStats{}
-						tr, serverTracer = mocklogging.NewMockConnectionTracer(mockCtrl)
-						trafficSecret1 := make([]byte, 16)
-						trafficSecret2 := make([]byte, 16)
-						rand.Read(trafficSecret1)
-						rand.Read(trafficSecret2)
-
-						client = newUpdatableAEAD(&rttStats, nil, utils.DefaultLogger, v)
-						server = newUpdatableAEAD(&rttStats, tr, utils.DefaultLogger, v)
-						client.SetReadKey(cs, trafficSecret2)
-						client.SetWriteKey(cs, trafficSecret1)
-						server.SetReadKey(cs, trafficSecret1)
-						server.SetWriteKey(cs, trafficSecret2)
-					})
-
-					Context("header protection", func() {
-						It("encrypts and decrypts the header", func() {
-							var lastFiveBitsDifferent int
-							for i := 0; i < 100; i++ {
-								sample := make([]byte, 16)
-								rand.Read(sample)
-								header := []byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}
-								client.EncryptHeader(sample, &header[0], header[9:13])
-								if header[0]&0x1f != 0xb5&0x1f {
-									lastFiveBitsDifferent++
-								}
-								Expect(header[0] & 0xe0).To(Equal(byte(0xb5 & 0xe0)))
-								Expect(header[1:9]).To(Equal([]byte{1, 2, 3, 4, 5, 6, 7, 8}))
-								Expect(header[9:13]).ToNot(Equal([]byte{0xde, 0xad, 0xbe, 0xef}))
-								server.DecryptHeader(sample, &header[0], header[9:13])
-								Expect(header).To(Equal([]byte{0xb5, 1, 2, 3, 4, 5, 6, 7, 8, 0xde, 0xad, 0xbe, 0xef}))
-							}
-							Expect(lastFiveBitsDifferent).To(BeNumerically(">", 75))
-						})
-					})
-
-					Context("message encryption", func() {
-						var msg, ad []byte
-
-						BeforeEach(func() {
-							msg = []byte("Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.")
-							ad = []byte("Donec in velit neque.")
-						})
-
-						It("encrypts and decrypts a message", func() {
-							encrypted := server.Seal(nil, msg, 0x1337, ad)
-							opened, err := client.Open(nil, encrypted, time.Now(), 0x1337, protocol.KeyPhaseZero, ad)
-							Expect(err).ToNot(HaveOccurred())
-							Expect(opened).To(Equal(msg))
-						})
-
-						It("saves the first packet number", func() {
-							client.Seal(nil, msg, 0x1337, ad)
-							Expect(client.FirstPacketNumber()).To(Equal(protocol.PacketNumber(0x1337)))
-							client.Seal(nil, msg, 0x1338, ad)
-							Expect(client.FirstPacketNumber()).To(Equal(protocol.PacketNumber(0x1337)))
-						})
-
-						It("fails to open a message if the associated data is not the same", func() {
-							encrypted := client.Seal(nil, msg, 0x1337, ad)
-							_, err := server.Open(nil, encrypted, time.Now(), 0x1337, protocol.KeyPhaseZero, []byte("wrong ad"))
-							Expect(err).To(MatchError(ErrDecryptionFailed))
-						})
-
-						It("fails to open a message if the packet number is not the same", func() {
-							encrypted := server.Seal(nil, msg, 0x1337, ad)
-							_, err := client.Open(nil, encrypted, time.Now(), 0x42, protocol.KeyPhaseZero, ad)
-							Expect(err).To(MatchError(ErrDecryptionFailed))
-						})
-
-						It("decodes the packet number", func() {
-							encrypted := server.Seal(nil, msg, 0x1337, ad)
-							_, err := client.Open(nil, encrypted, time.Now(), 0x1337, protocol.KeyPhaseZero, ad)
-							Expect(err).ToNot(HaveOccurred())
-							Expect(client.DecodePacketNumber(0x38, protocol.PacketNumberLen1)).To(BeEquivalentTo(0x1338))
-						})
-
-						It("ignores packets it can't decrypt for packet number derivation", func() {
-							encrypted := server.Seal(nil, msg, 0x1337, ad)
-							_, err := client.Open(nil, encrypted[:len(encrypted)-1], time.Now(), 0x1337, protocol.KeyPhaseZero, ad)
-							Expect(err).To(HaveOccurred())
-							Expect(client.DecodePacketNumber(0x38, protocol.PacketNumberLen1)).To(BeEquivalentTo(0x38))
-						})
-
-						It("returns an AEAD_LIMIT_REACHED error when reaching the AEAD limit", func() {
-							client.invalidPacketLimit = 10
-							for i := 0; i < 9; i++ {
-								_, err := client.Open(nil, []byte("foobar"), time.Now(), protocol.PacketNumber(i), protocol.KeyPhaseZero, []byte("ad"))
-								Expect(err).To(MatchError(ErrDecryptionFailed))
-							}
-							_, err := client.Open(nil, []byte("foobar"), time.Now(), 10, protocol.KeyPhaseZero, []byte("ad"))
-							Expect(err).To(HaveOccurred())
-							Expect(err).To(BeAssignableToTypeOf(&qerr.TransportError{}))
-							Expect(err.(*qerr.TransportError).ErrorCode).To(Equal(qerr.AEADLimitReached))
-						})
-
-						Context("key updates", func() {
-							Context("receiving key updates", func() {
-								It("updates keys", func() {
-									now := time.Now()
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-									encrypted0 := server.Seal(nil, msg, 0x1337, ad)
-									server.rollKeys()
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-									encrypted1 := server.Seal(nil, msg, 0x1337, ad)
-									Expect(encrypted0).ToNot(Equal(encrypted1))
-									// expect opening to fail. The client didn't roll keys yet
-									_, err := client.Open(nil, encrypted1, now, 0x1337, protocol.KeyPhaseZero, ad)
-									Expect(err).To(MatchError(ErrDecryptionFailed))
-									client.rollKeys()
-									decrypted, err := client.Open(nil, encrypted1, now, 0x1337, protocol.KeyPhaseOne, ad)
-									Expect(err).ToNot(HaveOccurred())
-									Expect(decrypted).To(Equal(msg))
-								})
-
-								It("updates the keys when receiving a packet with the next key phase", func() {
-									now := time.Now()
-									// receive the first packet at key phase zero
-									encrypted0 := client.Seal(nil, msg, 0x42, ad)
-									decrypted, err := server.Open(nil, encrypted0, now, 0x42, protocol.KeyPhaseZero, ad)
-									Expect(err).ToNot(HaveOccurred())
-									Expect(decrypted).To(Equal(msg))
-									// send one packet at key phase zero
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-									_ = server.Seal(nil, msg, 0x1, ad)
-									// now received a message at key phase one
-									client.rollKeys()
-									encrypted1 := client.Seal(nil, msg, 0x43, ad)
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), true)
-									decrypted, err = server.Open(nil, encrypted1, now, 0x43, protocol.KeyPhaseOne, ad)
-									Expect(err).ToNot(HaveOccurred())
-									Expect(decrypted).To(Equal(msg))
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-								})
-
-								It("opens a reordered packet with the old keys after an update", func() {
-									now := time.Now()
-									encrypted01 := client.Seal(nil, msg, 0x42, ad)
-									encrypted02 := client.Seal(nil, msg, 0x43, ad)
-									// receive the first packet with key phase 0
-									_, err := server.Open(nil, encrypted01, now, 0x42, protocol.KeyPhaseZero, ad)
-									Expect(err).ToNot(HaveOccurred())
-									// send one packet at key phase zero
-									_ = server.Seal(nil, msg, 0x1, ad)
-									// now receive a packet with key phase 1
-									client.rollKeys()
-									encrypted1 := client.Seal(nil, msg, 0x44, ad)
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), true)
-									_, err = server.Open(nil, encrypted1, now, 0x44, protocol.KeyPhaseOne, ad)
-									Expect(err).ToNot(HaveOccurred())
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-									// now receive a reordered packet with key phase 0
-									decrypted, err := server.Open(nil, encrypted02, now, 0x43, protocol.KeyPhaseZero, ad)
-									Expect(err).ToNot(HaveOccurred())
-									Expect(decrypted).To(Equal(msg))
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-								})
-
-								It("drops keys 3 PTOs after a key update", func() {
-									now := time.Now()
-									rttStats.UpdateRTT(10*time.Millisecond, 0, now)
-									pto := rttStats.PTO(true)
-									encrypted01 := client.Seal(nil, msg, 0x42, ad)
-									encrypted02 := client.Seal(nil, msg, 0x43, ad)
-									// receive the first packet with key phase 0
-									_, err := server.Open(nil, encrypted01, now, 0x42, protocol.KeyPhaseZero, ad)
-									Expect(err).ToNot(HaveOccurred())
-									// send one packet at key phase zero
-									_ = server.Seal(nil, msg, 0x1, ad)
-									// now receive a packet with key phase 1
-									client.rollKeys()
-									encrypted1 := client.Seal(nil, msg, 0x44, ad)
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), true)
-									serverTracer.EXPECT().DroppedKey(protocol.KeyPhase(0))
-									_, err = server.Open(nil, encrypted1, now, 0x44, protocol.KeyPhaseOne, ad)
-									Expect(err).ToNot(HaveOccurred())
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-									// now receive a reordered packet with key phase 0
-									_, err = server.Open(nil, encrypted02, now.Add(3*pto).Add(time.Nanosecond), 0x43, protocol.KeyPhaseZero, ad)
-									Expect(err).To(MatchError(ErrKeysDropped))
-								})
-
-								It("allows the first key update immediately", func() {
-									// receive a packet at key phase one, before having sent or received any packets at key phase 0
-									client.rollKeys()
-									encrypted1 := client.Seal(nil, msg, 0x1337, ad)
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), true)
-									_, err := server.Open(nil, encrypted1, time.Now(), 0x1337, protocol.KeyPhaseOne, ad)
-									Expect(err).ToNot(HaveOccurred())
-								})
-
-								It("only errors when the peer starts with key phase 1 if decrypting the packet succeeds", func() {
-									client.rollKeys()
-									encrypted := client.Seal(nil, msg, 0x1337, ad)
-									encrypted = encrypted[:len(encrypted)-1]
-									_, err := server.Open(nil, encrypted, time.Now(), 0x1337, protocol.KeyPhaseOne, ad)
-									Expect(err).To(MatchError(ErrDecryptionFailed))
-								})
-
-								It("errors when the peer updates keys too frequently", func() {
-									server.rollKeys()
-									client.rollKeys()
-									// receive the first packet at key phase one
-									encrypted0 := client.Seal(nil, msg, 0x42, ad)
-									_, err := server.Open(nil, encrypted0, time.Now(), 0x42, protocol.KeyPhaseOne, ad)
-									Expect(err).ToNot(HaveOccurred())
-									// now receive a packet at key phase two, before having sent any packets
-									client.rollKeys()
-									encrypted1 := client.Seal(nil, msg, 0x42, ad)
-									_, err = server.Open(nil, encrypted1, time.Now(), 0x42, protocol.KeyPhaseZero, ad)
-									Expect(err).To(MatchError(&qerr.TransportError{
-										ErrorCode:    qerr.KeyUpdateError,
-										ErrorMessage: "keys updated too quickly",
-									}))
-								})
-							})
-
-							Context("initiating key updates", func() {
-								const firstKeyUpdateInterval = 5
-								const keyUpdateInterval = 20
-								var origKeyUpdateInterval, origFirstKeyUpdateInterval uint64
-
-								BeforeEach(func() {
-									origKeyUpdateInterval = KeyUpdateInterval
-									origFirstKeyUpdateInterval = FirstKeyUpdateInterval
-									KeyUpdateInterval = keyUpdateInterval
-									FirstKeyUpdateInterval = firstKeyUpdateInterval
-									server.SetHandshakeConfirmed()
-								})
-
-								AfterEach(func() {
-									KeyUpdateInterval = origKeyUpdateInterval
-									FirstKeyUpdateInterval = origFirstKeyUpdateInterval
-								})
-
-								It("initiates a key update after sealing the maximum number of packets, for the first update", func() {
-									for i := 0; i < firstKeyUpdateInterval; i++ {
-										pn := protocol.PacketNumber(i)
-										Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-										server.Seal(nil, msg, pn, ad)
-									}
-									// the first update is allowed without receiving an acknowledgement
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-								})
-
-								It("initiates a key update after sealing the maximum number of packets, for subsequent updates", func() {
-									server.rollKeys()
-									client.rollKeys()
-									for i := 0; i < keyUpdateInterval; i++ {
-										pn := protocol.PacketNumber(i)
-										Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-										server.Seal(nil, msg, pn, ad)
-									}
-									// no update allowed before receiving an acknowledgement for the current key phase
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-									// receive an ACK for a packet sent in key phase 0
-									b := client.Seal(nil, []byte("foobar"), 1, []byte("ad"))
-									_, err := server.Open(nil, b, time.Now(), 1, protocol.KeyPhaseOne, []byte("ad"))
-									Expect(err).ToNot(HaveOccurred())
-									ExpectWithOffset(1, server.SetLargestAcked(0)).To(Succeed())
-									serverTracer.EXPECT().DroppedKey(protocol.KeyPhase(0))
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(2), false)
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-								})
-
-								It("errors if the peer acknowledges a packet sent in the next key phase using the old key phase", func() {
-									// First make sure that we update our keys.
-									for i := 0; i < firstKeyUpdateInterval; i++ {
-										pn := protocol.PacketNumber(i)
-										Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-										server.Seal(nil, msg, pn, ad)
-									}
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-									// Now that our keys are updated, send a packet using the new keys.
-									const nextPN = firstKeyUpdateInterval + 1
-									server.Seal(nil, msg, nextPN, ad)
-									// We haven't decrypted any packet in the new key phase yet.
-									// This means that the ACK must have been sent in the old key phase.
-									Expect(server.SetLargestAcked(nextPN)).To(MatchError(&qerr.TransportError{
-										ErrorCode:    qerr.KeyUpdateError,
-										ErrorMessage: "received ACK for key phase 1, but peer didn't update keys",
-									}))
-								})
-
-								It("doesn't error before actually sending a packet in the new key phase", func() {
-									// First make sure that we update our keys.
-									for i := 0; i < firstKeyUpdateInterval; i++ {
-										pn := protocol.PacketNumber(i)
-										Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-										server.Seal(nil, msg, pn, ad)
-									}
-									b := client.Seal(nil, []byte("foobar"), 1, []byte("ad"))
-									_, err := server.Open(nil, b, time.Now(), 1, protocol.KeyPhaseZero, []byte("ad"))
-									Expect(err).ToNot(HaveOccurred())
-									ExpectWithOffset(1, server.SetLargestAcked(0)).To(Succeed())
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
-									// Now that our keys are updated, send a packet using the new keys.
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-									// We haven't decrypted any packet in the new key phase yet.
-									// This means that the ACK must have been sent in the old key phase.
-									Expect(server.SetLargestAcked(1)).ToNot(HaveOccurred())
-								})
-
-								It("initiates a key update after opening the maximum number of packets, for the first update", func() {
-									for i := 0; i < firstKeyUpdateInterval; i++ {
-										pn := protocol.PacketNumber(i)
-										Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-										encrypted := client.Seal(nil, msg, pn, ad)
-										_, err := server.Open(nil, encrypted, time.Now(), pn, protocol.KeyPhaseZero, ad)
-										Expect(err).ToNot(HaveOccurred())
-									}
-									// the first update is allowed without receiving an acknowledgement
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-								})
-
-								It("initiates a key update after opening the maximum number of packets, for subsequent updates", func() {
-									server.rollKeys()
-									client.rollKeys()
-									for i := 0; i < keyUpdateInterval; i++ {
-										pn := protocol.PacketNumber(i)
-										Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-										encrypted := client.Seal(nil, msg, pn, ad)
-										_, err := server.Open(nil, encrypted, time.Now(), pn, protocol.KeyPhaseOne, ad)
-										Expect(err).ToNot(HaveOccurred())
-									}
-									// no update allowed before receiving an acknowledgement for the current key phase
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-									server.Seal(nil, msg, 1, ad)
-									Expect(server.SetLargestAcked(1)).To(Succeed())
-									serverTracer.EXPECT().DroppedKey(protocol.KeyPhase(0))
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(2), false)
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-								})
-
-								It("drops keys 3 PTOs after a key update", func() {
-									now := time.Now()
-									for i := 0; i < firstKeyUpdateInterval; i++ {
-										pn := protocol.PacketNumber(i)
-										Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-										server.Seal(nil, msg, pn, ad)
-									}
-									b := client.Seal(nil, []byte("foobar"), 1, []byte("ad"))
-									_, err := server.Open(nil, b, now, 1, protocol.KeyPhaseZero, []byte("ad"))
-									Expect(err).ToNot(HaveOccurred())
-									Expect(server.SetLargestAcked(0)).To(Succeed())
-									// Now we've initiated the first key update.
-									// Decrypt a message sent from the client more than 3 PTO later to make sure the key is still there
-									threePTO := 3 * rttStats.PTO(false)
-									dataKeyPhaseZero := client.Seal(nil, msg, 1, ad)
-									_, err = server.Open(nil, dataKeyPhaseZero, now.Add(threePTO).Add(time.Second), 1, protocol.KeyPhaseZero, ad)
-									Expect(err).ToNot(HaveOccurred())
-									// Now receive a packet with key phase 1.
-									// This should start the timer to drop the keys after 3 PTOs.
-									client.rollKeys()
-									dataKeyPhaseOne := client.Seal(nil, msg, 10, ad)
-									t := now.Add(threePTO).Add(time.Second)
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), true)
-									_, err = server.Open(nil, dataKeyPhaseOne, t, 10, protocol.KeyPhaseOne, ad)
-									Expect(err).ToNot(HaveOccurred())
-									// Make sure the keys are still here.
-									_, err = server.Open(nil, dataKeyPhaseZero, t.Add(threePTO*9/10), 1, protocol.KeyPhaseZero, ad)
-									Expect(err).ToNot(HaveOccurred())
-									serverTracer.EXPECT().DroppedKey(protocol.KeyPhase(0))
-									_, err = server.Open(nil, dataKeyPhaseZero, t.Add(threePTO).Add(time.Nanosecond), 1, protocol.KeyPhaseZero, ad)
-									Expect(err).To(MatchError(ErrKeysDropped))
-								})
-
-								It("doesn't drop the first key generation too early", func() {
-									now := time.Now()
-									data1 := client.Seal(nil, msg, 1, ad)
-									_, err := server.Open(nil, data1, now, 1, protocol.KeyPhaseZero, ad)
-									Expect(err).ToNot(HaveOccurred())
-									for i := 0; i < firstKeyUpdateInterval; i++ {
-										pn := protocol.PacketNumber(i)
-										Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-										server.Seal(nil, msg, pn, ad)
-										Expect(server.SetLargestAcked(pn)).To(Succeed())
-									}
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-									// The server never received a packet at key phase 1.
-									// Make sure the key phase 0 is still there at a much later point.
-									data2 := client.Seal(nil, msg, 1, ad)
-									_, err = server.Open(nil, data2, now.Add(10*rttStats.PTO(true)), 1, protocol.KeyPhaseZero, ad)
-									Expect(err).ToNot(HaveOccurred())
-								})
-
-								It("drops keys early when the peer forces initiates a key update within the 3 PTO period", func() {
-									for i := 0; i < firstKeyUpdateInterval; i++ {
-										pn := protocol.PacketNumber(i)
-										Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-										server.Seal(nil, msg, pn, ad)
-									}
-									b := client.Seal(nil, []byte("foobar"), 1, []byte("ad"))
-									_, err := server.Open(nil, b, time.Now(), 1, protocol.KeyPhaseZero, []byte("ad"))
-									Expect(err).ToNot(HaveOccurred())
-									ExpectWithOffset(1, server.SetLargestAcked(0)).To(Succeed())
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-									const nextPN = keyUpdateInterval + 1
-									// Send and receive an acknowledgement for a packet in key phase 1.
-									// We are now running a timer to drop the keys with 3 PTO.
-									server.Seal(nil, msg, nextPN, ad)
-									client.rollKeys()
-									dataKeyPhaseOne := client.Seal(nil, msg, 2, ad)
-									now := time.Now()
-									_, err = server.Open(nil, dataKeyPhaseOne, now, 2, protocol.KeyPhaseOne, ad)
-									Expect(err).ToNot(HaveOccurred())
-									Expect(server.SetLargestAcked(nextPN))
-									// Now the client sends us a packet in key phase 2, forcing us to update keys before the 3 PTO period is over.
-									// This mean that we need to drop the keys for key phase 0 immediately.
-									client.rollKeys()
-									dataKeyPhaseTwo := client.Seal(nil, msg, 3, ad)
-									gomock.InOrder(
-										serverTracer.EXPECT().DroppedKey(protocol.KeyPhase(0)),
-										serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(2), true),
-									)
-									_, err = server.Open(nil, dataKeyPhaseTwo, now, 3, protocol.KeyPhaseZero, ad)
-									Expect(err).ToNot(HaveOccurred())
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-								})
-
-								It("drops keys early when we initiate another key update within the 3 PTO period", func() {
-									server.SetHandshakeConfirmed()
-									// send so many packets that we initiate the first key update
-									for i := 0; i < firstKeyUpdateInterval; i++ {
-										pn := protocol.PacketNumber(i)
-										Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-										server.Seal(nil, msg, pn, ad)
-									}
-									b := client.Seal(nil, []byte("foobar"), 1, []byte("ad"))
-									_, err := server.Open(nil, b, time.Now(), 1, protocol.KeyPhaseZero, []byte("ad"))
-									Expect(err).ToNot(HaveOccurred())
-									ExpectWithOffset(1, server.SetLargestAcked(0)).To(Succeed())
-									serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-									// send so many packets that we initiate the next key update
-									for i := keyUpdateInterval; i < 2*keyUpdateInterval; i++ {
-										pn := protocol.PacketNumber(i)
-										Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
-										server.Seal(nil, msg, pn, ad)
-									}
-									client.rollKeys()
-									b = client.Seal(nil, []byte("foobar"), 2, []byte("ad"))
-									now := time.Now()
-									_, err = server.Open(nil, b, now, 2, protocol.KeyPhaseOne, []byte("ad"))
-									Expect(err).ToNot(HaveOccurred())
-									ExpectWithOffset(1, server.SetLargestAcked(keyUpdateInterval)).To(Succeed())
-									gomock.InOrder(
-										serverTracer.EXPECT().DroppedKey(protocol.KeyPhase(0)),
-										serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(2), false),
-									)
-									Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
-									// We haven't received an ACK for a packet sent in key phase 2 yet.
-									// Make sure we canceled the timer to drop the previous key phase.
-									b = client.Seal(nil, []byte("foobar"), 3, []byte("ad"))
-									_, err = server.Open(nil, b, now.Add(10*rttStats.PTO(true)), 3, protocol.KeyPhaseOne, []byte("ad"))
-									Expect(err).ToNot(HaveOccurred())
-								})
-							})
-						})
-					})
-				})
-			}
-		})
+func TestFastKeyUpdateByUs(t *testing.T) {
+	const firstKeyUpdateInterval = 5
+	const keyUpdateInterval = 20
+	setKeyUpdateIntervals(t, firstKeyUpdateInterval, keyUpdateInterval)
+
+	var rttStats utils.RTTStats
+	rttStats.UpdateRTT(10*time.Millisecond, 0, time.Now())
+	client, server, serverTracer := setupEndpoints(t, &rttStats)
+	server.SetHandshakeConfirmed()
+
+	// send so many packets that we initiate the first key update
+	for i := 0; i < firstKeyUpdateInterval; i++ {
+		pn := protocol.PacketNumber(i)
+		require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+		server.Seal(nil, []byte(msg), pn, []byte(ad))
+	}
+	b := client.Seal(nil, []byte("foobar"), 1, []byte("ad"))
+	_, err := server.Open(nil, b, time.Now(), 1, protocol.KeyPhaseZero, []byte("ad"))
+	require.NoError(t, err)
+	require.NoError(t, server.SetLargestAcked(0))
+	serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(1), false)
+	require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+
+	// send so many packets that we initiate the next key update
+	for i := keyUpdateInterval; i < 2*keyUpdateInterval; i++ {
+		pn := protocol.PacketNumber(i)
+		require.Equal(t, protocol.KeyPhaseOne, server.KeyPhase())
+		server.Seal(nil, []byte(msg), pn, []byte(ad))
 	}
-})
+	client.rollKeys()
+	b = client.Seal(nil, []byte("foobar"), 2, []byte("ad"))
+	now := time.Now()
+	_, err = server.Open(nil, b, now, 2, protocol.KeyPhaseOne, []byte("ad"))
+	require.NoError(t, err)
+	require.NoError(t, server.SetLargestAcked(keyUpdateInterval))
+	gomock.InOrder(
+		serverTracer.EXPECT().DroppedKey(protocol.KeyPhase(0)),
+		serverTracer.EXPECT().UpdatedKey(protocol.KeyPhase(2), false),
+	)
+	require.Equal(t, protocol.KeyPhaseZero, server.KeyPhase())
+
+	// We haven't received an ACK for a packet sent in key phase 2 yet.
+	// Make sure we canceled the timer to drop the previous key phase.
+	b = client.Seal(nil, []byte("foobar"), 3, []byte("ad"))
+	_, err = server.Open(nil, b, now.Add(10*rttStats.PTO(true)), 3, protocol.KeyPhaseOne, []byte("ad"))
+	require.NoError(t, err)
+}
 
 func getClientAndServer() (client, server *updatableAEAD) {
 	trafficSecret1 := make([]byte, 16)