quic: fix client key discard

Author: akhinvasara-jumptrading

src/waltz/quic/fd_quic.c

@@ -1886,7 +1886,7 @@ fd_quic_handle_v1_handshake(
   /* RFC 9000 Section 17.2.2.1. Abandoning Initial Packets
      > A server stops sending and processing Initial packets when it
      > receives its first Handshake packet. */
-  fd_quic_abandon_enc_level( conn, fd_quic_enc_level_initial_id );
+  if( FD_LIKELY( quic->config.role==FD_QUIC_ROLE_SERVER ) ) fd_quic_abandon_enc_level( conn, fd_quic_enc_level_initial_id );
   conn->peer_enc_level = (uchar)fd_uchar_max( conn->peer_enc_level, fd_quic_enc_level_handshake_id );
 
   /* handle frames */
@@ -3562,12 +3562,6 @@ fd_quic_conn_tx( fd_quic_t      * quic,
    * This ensures that ack-only packets only occur when nothing else needs
    * to be sent */
   uint enc_level = fd_quic_tx_enc_level( conn, 1 /* acks */ );
-  /* RFC 9000 Section 17.2.2.1. Abandoning Initial Packets
-     > A client stops both sending and processing Initial packets when
-     > it sends its first Handshake packet. */
-  if( quic->config.role==FD_QUIC_ROLE_CLIENT && enc_level==fd_quic_enc_level_handshake_id ) {
-    fd_quic_abandon_enc_level( conn, fd_quic_enc_level_initial_id );
-  }
 
   /* nothing to send / bad state? */
   if( enc_level == ~0u ) return;
@@ -3588,6 +3582,13 @@ fd_quic_conn_tx( fd_quic_t      * quic,
   //}
 
   while( enc_level != ~0u ) {
+    /* RFC 9000 Section 17.2.2.1. Abandoning Initial Packets
+       > A client stops both sending and processing Initial packets when
+       > it sends its first Handshake packet. */
+    if( FD_UNLIKELY( (quic->config.role==FD_QUIC_ROLE_CLIENT) & (enc_level==fd_quic_enc_level_handshake_id) ) ) {
+      fd_quic_abandon_enc_level( conn, fd_quic_enc_level_initial_id );
+    }
+
     uint initial_pkt = 0;    /* is this the first initial packet? */
 
 

src/waltz/quic/tests/test_quic_retry_integration.c

@@ -91,9 +91,10 @@ test_initial_token_odd_sz( fd_quic_t * server_quic,
   /* Verify server sent a retry packet (difference-based check) */
   FD_TEST( server_quic->metrics.conn_retry_cnt == initial_retry_count + 1 );
 
-  /* Verify the 46-byte token was counted in the "other sizes" category */
+  /* Verify the 46-byte token was counted in the "other sizes" category
+     See test_retry_integration for details on the expected increases */
   FD_TEST( server_quic->metrics.initial_token_len_cnt[2] == initial_token_len_other + 1 );
-  FD_TEST( server_quic->metrics.initial_token_len_cnt[1] == initial_token_len_our + 1 );
+  FD_TEST( server_quic->metrics.initial_token_len_cnt[1] == initial_token_len_our + 2 );
 
   /* Verify the connection was created */
   FD_TEST( server_quic->metrics.conn_created_cnt == conn_created_count + 1 );
@@ -126,10 +127,13 @@ test_retry_integration( fd_quic_t * server_quic,
   FD_TEST( server_quic->metrics.conn_created_cnt== 1 );
   FD_TEST( server_quic->metrics.conn_retry_cnt  == 1 );
 
-  /* Check initial token length metrics - should have seen 1 packet with no token (idx 0)
-     and 1 packet with retry token (idx 1) */
+  /* Check initial token length metrics - should have seen:
+     * 1 packet with no token (idx 0), the very first initial
+     * 2 packets with retry token with correct fd_quic len (idx 1):
+       - first the Initial with crypto frames triggered by Retry
+       - second an Initial with ack frame, ACKing server Initial */
   FD_TEST( server_quic->metrics.initial_token_len_cnt[0] == 1 ); /* no token */
-  FD_TEST( server_quic->metrics.initial_token_len_cnt[1] == 1 ); /* retry token */
+  FD_TEST( server_quic->metrics.initial_token_len_cnt[1] == 2 ); /* retry token */
   FD_TEST( server_quic->metrics.initial_token_len_cnt[2] == 0 ); /* other sizes */
   /* Server: Retry, Initial, Handshake
      Client: Initial, Initial, Handshake */