Create setter functions for the msg_callback and msg_callback_arg

We create setter functions for the msg_callback and msg_callback_arg so
that these values can be properly propagated to the QRX/QTX/TXP even
after the channel has been created.

Reviewed-by: Tomas Mraz <tomas@openssl.org>
Reviewed-by: Hugo Landau <hlandau@openssl.org>
(Merged from https://github.com/openssl/openssl/pull/20914)

6 files changed, 89 insertions, 52 deletions

diff --git a/ssl/quic/quic_channel.c b/ssl/quic/quic_channel.c
index 76546e2bd3..328c8bee8b 100644
--- a/ssl/quic/quic_channel.c
+++ b/ssl/quic/quic_channel.c
@@ -130,10 +130,6 @@ static int ch_init(QUIC_CHANNEL *ch)
     /* We plug in a network write BIO to the QTX later when we get one. */
     qtx_args.libctx = ch->libctx;
     qtx_args.mdpl = QUIC_MIN_INITIAL_DGRAM_LEN;
-    /* Callback related arguments */
-    qtx_args.msg_callback       = ch->msg_callback;
-    qtx_args.msg_callback_arg   = ch->msg_callback_arg;
-    qtx_args.msg_callback_s     = ch->msg_callback_s;
     ch->rx_max_udp_payload_size = qtx_args.mdpl;
 
     ch->qtx = ossl_qtx_new(&qtx_args);
@@ -212,10 +208,6 @@ static int ch_init(QUIC_CHANNEL *ch)
     txp_args.cc_data                = ch->cc_data;
     txp_args.now                    = get_time;
     txp_args.now_arg                = ch;
-    /* Callback related arguments */
-    txp_args.msg_callback           = ch->msg_callback;
-    txp_args.msg_callback_arg       = ch->msg_callback_arg;
-    txp_args.msg_callback_s         = ch->msg_callback_s;
 
     for (pn_space = QUIC_PN_SPACE_INITIAL; pn_space < QUIC_PN_SPACE_NUM; ++pn_space) {
         ch->crypto_send[pn_space] = ossl_quic_sstream_new(INIT_CRYPTO_BUF_LEN);
@@ -248,10 +240,6 @@ static int ch_init(QUIC_CHANNEL *ch)
     qrx_args.demux              = ch->demux;
     qrx_args.short_conn_id_len  = rx_short_cid_len;
     qrx_args.max_deferred       = 32;
-    /* Callback related arguments */
-    qrx_args.msg_callback       = ch->msg_callback;
-    qrx_args.msg_callback_arg   = ch->msg_callback_arg;
-    qrx_args.msg_callback_s     = ch->msg_callback_s;
 
     if ((ch->qrx = ossl_qrx_new(&qrx_args)) == NULL)
         goto err;
@@ -360,16 +348,13 @@ QUIC_CHANNEL *ossl_quic_channel_new(const QUIC_CHANNEL_ARGS *args)
     if ((ch = OPENSSL_zalloc(sizeof(*ch))) == NULL)
         return NULL;
 
-    ch->libctx           = args->libctx;
-    ch->propq            = args->propq;
-    ch->is_server        = args->is_server;
-    ch->tls              = args->tls;
-    ch->mutex            = args->mutex;
-    ch->now_cb           = args->now_cb;
-    ch->now_cb_arg       = args->now_cb_arg;
-    ch->msg_callback     = args->msg_callback;
-    ch->msg_callback_arg = args->msg_callback_arg;
-    ch->msg_callback_s   = args->msg_callback_s;
+    ch->libctx      = args->libctx;
+    ch->propq       = args->propq;
+    ch->is_server   = args->is_server;
+    ch->tls         = args->tls;
+    ch->mutex       = args->mutex;
+    ch->now_cb      = args->now_cb;
+    ch->now_cb_arg  = args->now_cb_arg;
 
     if (!ch_init(ch)) {
         OPENSSL_free(ch);
@@ -2524,3 +2509,24 @@ int ossl_quic_channel_replace_local_cid(QUIC_CHANNEL *ch,
         return 0;
     return 1;
 }
+
+void ossl_quic_channel_set_msg_callback(QUIC_CHANNEL *ch,
+                                        ossl_msg_cb msg_callback,
+                                        SSL *msg_callback_s)
+{
+    ch->msg_callback = msg_callback;
+    ch->msg_callback_s = msg_callback_s;
+    ossl_qtx_set_msg_callback(ch->qtx, msg_callback, msg_callback_s);
+    ossl_quic_tx_packetiser_set_msg_callback(ch->txp, msg_callback,
+                                             msg_callback_s);
+    ossl_qrx_set_msg_callback(ch->qrx, msg_callback, msg_callback_s);
+}
+
+void ossl_quic_channel_set_msg_callback_arg(QUIC_CHANNEL *ch,
+                                            void *msg_callback_arg)
+{
+    ch->msg_callback_arg = msg_callback_arg;
+    ossl_qtx_set_msg_callback_arg(ch->qtx, msg_callback_arg);
+    ossl_quic_tx_packetiser_set_msg_callback_arg(ch->txp, msg_callback_arg);
+    ossl_qrx_set_msg_callback_arg(ch->qrx, msg_callback_arg);
+}
diff --git a/ssl/quic/quic_impl.c b/ssl/quic/quic_impl.c
index c623a3c0b9..4ba87a32af 100644
--- a/ssl/quic/quic_impl.c
+++ b/ssl/quic/quic_impl.c
@@ -326,13 +326,13 @@ SSL *ossl_quic_new(SSL_CTX *ctx)
     qc->default_blocking        = 1;
     qc->incoming_stream_policy  = SSL_INCOMING_STREAM_POLICY_AUTO;
     qc->last_error              = SSL_ERROR_NONE;
-    qc->msg_callback            = ctx->msg_callback;
-    qc->msg_callback_arg        = ctx->msg_callback_arg;
-    qc->msg_callback_s          = ssl_base;
 
     if (!create_channel(qc))
         goto err;
 
+    ossl_quic_channel_set_msg_callback(qc->ch, ctx->msg_callback, ssl_base);
+    ossl_quic_channel_set_msg_callback_arg(qc->ch, ctx->msg_callback_arg);
+
     qc_update_reject_policy(qc);
 
     /*
@@ -1045,7 +1045,7 @@ long ossl_quic_ctrl(SSL *s, int cmd, long larg, void *parg)
         return ctx.qc->default_ssl_mode;
 
     case SSL_CTRL_SET_MSG_CALLBACK_ARG:
-        ctx.qc->msg_callback_arg = parg;
+        ossl_quic_channel_set_msg_callback_arg(ctx.qc->ch, parg);
         /* This ctrl also needs to be passed to the internal SSL object */
         return SSL_ctrl(ctx.qc->tls, cmd, larg, parg);
 
@@ -1120,16 +1120,13 @@ static int create_channel(QUIC_CONNECTION *qc)
 {
     QUIC_CHANNEL_ARGS args = {0};
 
-    args.libctx           = qc->ssl.ctx->libctx;
-    args.propq            = qc->ssl.ctx->propq;
-    args.is_server        = qc->as_server;
-    args.tls              = qc->tls;
-    args.mutex            = qc->mutex;
-    args.now_cb           = qc->override_now_cb;
-    args.now_cb_arg       = qc->override_now_cb_arg;
-    args.msg_callback     = qc->msg_callback;
-    args.msg_callback_arg = qc->msg_callback_arg;
-    args.msg_callback_s   = qc->msg_callback_s;
+    args.libctx     = qc->ssl.ctx->libctx;
+    args.propq      = qc->ssl.ctx->propq;
+    args.is_server  = qc->as_server;
+    args.tls        = qc->tls;
+    args.mutex      = qc->mutex;
+    args.now_cb     = qc->override_now_cb;
+    args.now_cb_arg = qc->override_now_cb_arg;
 
     qc->ch = ossl_quic_channel_new(&args);
     if (qc->ch == NULL)
@@ -2672,7 +2669,8 @@ long ossl_quic_callback_ctrl(SSL *s, int cmd, void (*fp) (void))
 
     switch (cmd) {
     case SSL_CTRL_SET_MSG_CALLBACK:
-        ctx.qc->msg_callback = (ossl_msg_cb)fp;
+        ossl_quic_channel_set_msg_callback(ctx.qc->ch, (ossl_msg_cb)fp,
+                                           &ctx.qc->ssl);
         /* This callback also needs to be set on the internal SSL object */
         return ssl3_callback_ctrl(ctx.qc->tls, cmd, fp);;
 
diff --git a/ssl/quic/quic_local.h b/ssl/quic/quic_local.h
index d4088d4b69..46b0e72014 100644
--- a/ssl/quic/quic_local.h
+++ b/ssl/quic/quic_local.h
@@ -195,11 +195,6 @@ struct quic_conn_st {
      * and SSL_ERROR_WANT_WRITE.
      */
     int                             last_error;
-
-    /* Message callback related arguments */
-    ossl_msg_cb                     msg_callback;
-    void                            *msg_callback_arg;
-    SSL                             *msg_callback_s;
 };
 
 /* Internal calls to the QUIC CSM which come from various places. */
diff --git a/ssl/quic/quic_record_rx.c b/ssl/quic/quic_record_rx.c
index 266dee9d31..db70c119fa 100644
--- a/ssl/quic/quic_record_rx.c
+++ b/ssl/quic/quic_record_rx.c
@@ -176,9 +176,6 @@ OSSL_QRX *ossl_qrx_new(const OSSL_QRX_ARGS *args)
     qrx->short_conn_id_len      = args->short_conn_id_len;
     qrx->init_key_phase_bit     = args->init_key_phase_bit;
     qrx->max_deferred           = args->max_deferred;
-    qrx->msg_callback           = args->msg_callback;
-    qrx->msg_callback_arg       = args->msg_callback_arg;
-    qrx->msg_callback_s         = args->msg_callback_s;
     return qrx;
 }
 
@@ -1207,3 +1204,15 @@ uint64_t ossl_qrx_get_max_forged_pkt_count(OSSL_QRX *qrx,
     return el == NULL ? UINT64_MAX
         : ossl_qrl_get_suite_max_forged_pkt(el->suite_id);
 }
+
+void ossl_qrx_set_msg_callback(OSSL_QRX *qrx, ossl_msg_cb msg_callback,
+                               SSL *msg_callback_s)
+{
+    qrx->msg_callback = msg_callback;
+    qrx->msg_callback_s = msg_callback_s;
+}
+
+void ossl_qrx_set_msg_callback_arg(OSSL_QRX *qrx, void *msg_callback_arg)
+{
+    qrx->msg_callback_arg = msg_callback_arg;
+}
diff --git a/ssl/quic/quic_record_tx.c b/ssl/quic/quic_record_tx.c
index 748b4c358c..09862c7751 100644
--- a/ssl/quic/quic_record_tx.c
+++ b/ssl/quic/quic_record_tx.c
@@ -117,9 +117,6 @@ OSSL_QTX *ossl_qtx_new(const OSSL_QTX_ARGS *args)
     qtx->propq              = args->propq;
     qtx->bio                = args->bio;
     qtx->mdpl               = args->mdpl;
-    qtx->msg_callback       = args->msg_callback;
-    qtx->msg_callback_arg   = args->msg_callback_arg;
-    qtx->msg_callback_s     = args->msg_callback_s;
     return qtx;
 }
 
@@ -1005,3 +1002,15 @@ uint64_t ossl_qtx_get_max_epoch_pkt_count(OSSL_QTX *qtx, uint32_t enc_level)
 
     return ossl_qrl_get_suite_max_pkt(el->suite_id);
 }
+
+void ossl_qtx_set_msg_callback(OSSL_QTX *qtx, ossl_msg_cb msg_callback,
+                               SSL *msg_callback_s)
+{
+    qtx->msg_callback = msg_callback;
+    qtx->msg_callback_s = msg_callback_s;
+}
+
+void ossl_qtx_set_msg_callback_arg(OSSL_QTX *qtx, void *msg_callback_arg)
+{
+    qtx->msg_callback_arg = msg_callback_arg;
+}
diff --git a/ssl/quic/quic_txp.c b/ssl/quic/quic_txp.c
index b8bfb3592e..b04fe01115 100644
--- a/ssl/quic/quic_txp.c
+++ b/ssl/quic/quic_txp.c
@@ -69,6 +69,12 @@ struct ossl_quic_tx_packetiser_st {
     size_t          scratch_len;    /* number of bytes allocated for scratch */
     OSSL_QTX_IOVEC  *iovec;         /* scratch iovec array for use with QTX */
     size_t          alloc_iovec;    /* size of iovec array */
+
+    /* Message callback related arguments */
+    ossl_msg_cb msg_callback;
+    void *msg_callback_arg;
+    SSL *msg_callback_s;
+
 };
 
 /*
@@ -301,7 +307,7 @@ static int tx_helper_commit(struct tx_helper *h)
         return 0;
     }
 
-    if (h->txp->args.msg_callback != NULL && l > 0) {
+    if (h->txp->msg_callback != NULL && l > 0) {
         uint64_t ftype;
         int ctype = SSL3_RT_QUIC_FRAME_FULL;
         PACKET pkt;
@@ -318,9 +324,9 @@ static int tx_helper_commit(struct tx_helper *h)
                 || ftype == OSSL_QUIC_FRAME_TYPE_CRYPTO)
             ctype = SSL3_RT_QUIC_FRAME_HEADER;
 
-        h->txp->args.msg_callback(1, OSSL_QUIC1_VERSION, ctype, h->txn.data, l,
-                                  h->txp->args.msg_callback_s,
-                                  h->txp->args.msg_callback_arg);
+        h->txp->msg_callback(1, OSSL_QUIC1_VERSION, ctype, h->txn.data, l,
+                             h->txp->msg_callback_s,
+                             h->txp->msg_callback_arg);
     }
 
     h->scratch_bytes += l;
@@ -2376,3 +2382,17 @@ int ossl_quic_tx_packetiser_schedule_conn_close(OSSL_QUIC_TX_PACKETISER *txp,
     txp->want_conn_close                = 1;
     return 1;
 }
+
+void ossl_quic_tx_packetiser_set_msg_callback(OSSL_QUIC_TX_PACKETISER *txp,
+                                              ossl_msg_cb msg_callback,
+                                              SSL *msg_callback_s)
+{
+    txp->msg_callback = msg_callback;
+    txp->msg_callback_s = msg_callback_s;
+}
+
+void ossl_quic_tx_packetiser_set_msg_callback_arg(OSSL_QUIC_TX_PACKETISER *txp,
+                                                  void *msg_callback_arg)
+{
+    txp->msg_callback_arg = msg_callback_arg;
+}