Make key export callback and context connection-specific

Fixes #2188

Signed-off-by: Hanno Becker <hanno.becker@arm.com>
diff --git a/ChangeLog.d/key-export.txt b/ChangeLog.d/key-export.txt
index 5882d23..10d8c89 100644
--- a/ChangeLog.d/key-export.txt
+++ b/ChangeLog.d/key-export.txt
@@ -1,8 +1,13 @@
 API changes
-   * mbedtls_ssl_conf_export_keys_ext_cb() has been removed.
+   * mbedtls_ssl_conf_export_keys_ext_cb() and
+     mbedtls_ssl_conf_export_keys_cb() have been removed
+     and replaced by a new API
+     mbedtls_ssl_set_export_keys_cb().
    * The signature of key export callbacks configured via
-     mbedtls_ssl_conf_export_keys_cb() has changed, and raw
-     keys and IVs are no longer exported. Further, callbacks
+     mbedtls_ssl_set_export_keys_cb() is different from that
+     of the previous mbedtls_ssl_conf_export_keys_cb(): First,
+     raw keys and IVs are no longer exported. Further, callbacks
      now receive an additional parameter indicating the type
      of secret that's being exported, paving the way for the
-     larger number of secrets in TLS 1.3.
+     larger number of secrets in TLS 1.3. Finally, the key export
+     callback and context are now connection-specific.
diff --git a/docs/3.0-migration-guide.d/key-export.md b/docs/3.0-migration-guide.d/key-export.md
index 43781a2..967ecf8 100644
--- a/docs/3.0-migration-guide.d/key-export.md
+++ b/docs/3.0-migration-guide.d/key-export.md
@@ -7,9 +7,13 @@
     mbedtls_ssl_conf_export_keys_ext_cb()
 ```
 
-The API `mbedtls_ssl_conf_export_keys_ext_cb()` has been removed,
-and the function type of key export callback passed to
-`mbedtls_ssl_conf_export_keys_cb()` has changed, as follows:
+Those APIs have been removed and replaced by the new API
+`mbedtls_ssl_set_export_keys_cb()`. This API differs from
+the previous key export API in the following ways:
+
+- It is no longer bound to an SSL configuration, but to an
+  SSL context. This allows users to more easily identify the
+  connection an exported key belongs to.
 - It no longer exports raw keys and IV.
 - A secret type parameter has been added to identify which key
   is being exported. For TLS 1.2, only the master secret is
diff --git a/include/mbedtls/ssl.h b/include/mbedtls/ssl.h
index 540f552..2580020 100644
--- a/include/mbedtls/ssl.h
+++ b/include/mbedtls/ssl.h
@@ -1074,12 +1074,6 @@
     void *MBEDTLS_PRIVATE(p_ticket);                 /*!< context for the ticket callbacks   */
 #endif /* MBEDTLS_SSL_SESSION_TICKETS && MBEDTLS_SSL_SRV_C */
 
-#if defined(MBEDTLS_SSL_EXPORT_KEYS)
-    /** Callback to export key block and master secret                      */
-    mbedtls_ssl_export_keys_t *MBEDTLS_PRIVATE(f_export_keys);
-    void *MBEDTLS_PRIVATE(p_export_keys);            /*!< context for key export callback    */
-#endif
-
 #if defined(MBEDTLS_SSL_DTLS_CONNECTION_ID)
     size_t MBEDTLS_PRIVATE(cid_len); /*!< The length of CIDs for incoming DTLS records.      */
 #endif /* MBEDTLS_SSL_DTLS_CONNECTION_ID */
@@ -1260,6 +1254,12 @@
     int MBEDTLS_PRIVATE(minor_ver);              /*!< one of MBEDTLS_SSL_MINOR_VERSION_x macros */
     unsigned MBEDTLS_PRIVATE(badmac_seen);       /*!< records with a bad MAC received    */
 
+#if defined(MBEDTLS_SSL_EXPORT_KEYS)
+    /** Callback to export key block and master secret                      */
+    mbedtls_ssl_export_keys_t *MBEDTLS_PRIVATE(f_export_keys);
+    void *MBEDTLS_PRIVATE(p_export_keys);            /*!< context for key export callback    */
+#endif
+
 #if defined(MBEDTLS_X509_CRT_PARSE_C)
     /** Callback to customize X.509 certificate chain verification          */
     int (*MBEDTLS_PRIVATE(f_vrfy))(void *, mbedtls_x509_crt *, int, uint32_t *);
@@ -2015,17 +2015,15 @@
  *            key exporters, e.g. for EAP-TLS or DTLS-SRTP.
  *
  *
- * \param conf           The SSL configuration to which the export
- *                       callback should be attached. All connections
- *                       subsequently bound to this configuration will
- *                       have their keys exported.
+ * \param ssl            The SSL context to which the export
+ *                       callback should be attached.
  * \param f_export_keys  The callback for the key export.
  * \param p_export_keys  The opaque context pointer to be passed to the
  *                       callback \p f_export_keys.
  */
-void mbedtls_ssl_conf_export_keys_cb( mbedtls_ssl_config *conf,
-                                      mbedtls_ssl_export_keys_t *f_export_keys,
-                                      void *p_export_keys );
+void mbedtls_ssl_set_export_keys_cb( mbedtls_ssl_context *ssl,
+                                     mbedtls_ssl_export_keys_t *f_export_keys,
+                                     void *p_export_keys );
 #endif /* MBEDTLS_SSL_EXPORT_KEYS */
 
 #if defined(MBEDTLS_SSL_ASYNC_PRIVATE)
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 4c69a08..9268ede 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -986,14 +986,14 @@
     ((void) mac_enc);
 
 #if defined(MBEDTLS_SSL_EXPORT_KEYS)
-    if( ssl->conf->f_export_keys != NULL )
+    if( ssl->f_export_keys != NULL )
     {
-        ssl->conf->f_export_keys( ssl->conf->p_export_keys,
-                                  MBEDTLS_SSL_KEY_EXPORT_TLS12_MASTER_SECRET,
-                                  master, 48,
-                                  randbytes + 32,
-                                  randbytes,
-                                  tls_prf_get_type( tls_prf ) );
+        ssl->f_export_keys( ssl->p_export_keys,
+                            MBEDTLS_SSL_KEY_EXPORT_TLS12_MASTER_SECRET,
+                            master, 48,
+                            randbytes + 32,
+                            randbytes,
+                            tls_prf_get_type( tls_prf ) );
     }
 #endif
 
@@ -4184,12 +4184,12 @@
 #endif /* MBEDTLS_SSL_SESSION_TICKETS */
 
 #if defined(MBEDTLS_SSL_EXPORT_KEYS)
-void mbedtls_ssl_conf_export_keys_cb( mbedtls_ssl_config *conf,
-                                      mbedtls_ssl_export_keys_t *f_export_keys,
-                                      void *p_export_keys )
+void mbedtls_ssl_set_export_keys_cb( mbedtls_ssl_context *ssl,
+                                     mbedtls_ssl_export_keys_t *f_export_keys,
+                                     void *p_export_keys )
 {
-    conf->f_export_keys = f_export_keys;
-    conf->p_export_keys = p_export_keys;
+    ssl->f_export_keys = f_export_keys;
+    ssl->p_export_keys = p_export_keys;
 }
 #endif
 
diff --git a/programs/ssl/ssl_client2.c b/programs/ssl/ssl_client2.c
index c25b9ee..a7b7ece 100644
--- a/programs/ssl/ssl_client2.c
+++ b/programs/ssl/ssl_client2.c
@@ -1736,27 +1736,6 @@
         mbedtls_ssl_conf_encrypt_then_mac( &conf, opt.etm );
 #endif
 
-#if defined(MBEDTLS_SSL_EXPORT_KEYS)
-    if( opt.eap_tls != 0 )
-    {
-        mbedtls_ssl_conf_export_keys_cb( &conf, eap_tls_key_derivation,
-                                         &eap_tls_keying );
-    }
-    else if( opt.nss_keylog != 0 )
-    {
-        mbedtls_ssl_conf_export_keys_cb( &conf,
-                                         nss_keylog_export,
-                                         NULL );
-    }
-#if defined( MBEDTLS_SSL_DTLS_SRTP )
-    else if( opt.use_srtp != 0 )
-    {
-        mbedtls_ssl_conf_export_keys_cb( &conf, dtls_srtp_key_derivation,
-                                         &dtls_srtp_keying );
-    }
-#endif /* MBEDTLS_SSL_DTLS_SRTP */
-#endif /* MBEDTLS_SSL_EXPORT_KEYS */
-
 #if defined(MBEDTLS_DHM_C)
     if( opt.dhmlen != DFL_DHMLEN )
         mbedtls_ssl_conf_dhm_min_bitlen( &conf, opt.dhmlen );
@@ -1886,6 +1865,27 @@
         goto exit;
     }
 
+#if defined(MBEDTLS_SSL_EXPORT_KEYS)
+    if( opt.eap_tls != 0 )
+    {
+        mbedtls_ssl_set_export_keys_cb( &ssl, eap_tls_key_derivation,
+                                        &eap_tls_keying );
+    }
+    else if( opt.nss_keylog != 0 )
+    {
+        mbedtls_ssl_set_export_keys_cb( &ssl,
+                                        nss_keylog_export,
+                                        NULL );
+    }
+#if defined( MBEDTLS_SSL_DTLS_SRTP )
+    else if( opt.use_srtp != 0 )
+    {
+        mbedtls_ssl_set_export_keys_cb( &ssl, dtls_srtp_key_derivation,
+                                        &dtls_srtp_keying );
+    }
+#endif /* MBEDTLS_SSL_DTLS_SRTP */
+#endif /* MBEDTLS_SSL_EXPORT_KEYS */
+
 #if defined(MBEDTLS_X509_CRT_PARSE_C)
     if( ( ret = mbedtls_ssl_set_hostname( &ssl, opt.server_name ) ) != 0 )
     {
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index 9cecf7f..cb15866 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -2525,27 +2525,6 @@
         mbedtls_ssl_conf_encrypt_then_mac( &conf, opt.etm );
 #endif
 
-#if defined(MBEDTLS_SSL_EXPORT_KEYS)
-    if( opt.eap_tls != 0 )
-    {
-        mbedtls_ssl_conf_export_keys_cb( &conf, eap_tls_key_derivation,
-                                         &eap_tls_keying );
-    }
-    else if( opt.nss_keylog != 0 )
-    {
-        mbedtls_ssl_conf_export_keys_cb( &conf,
-                                         nss_keylog_export,
-                                         NULL );
-    }
-#if defined( MBEDTLS_SSL_DTLS_SRTP )
-    else if( opt.use_srtp != 0 )
-    {
-        mbedtls_ssl_conf_export_keys_cb( &conf, dtls_srtp_key_derivation,
-                                         &dtls_srtp_keying );
-    }
-#endif /* MBEDTLS_SSL_DTLS_SRTP */
-#endif /* MBEDTLS_SSL_EXPORT_KEYS */
-
 #if defined(MBEDTLS_SSL_ALPN)
     if( opt.alpn_string != NULL )
         if( ( ret = mbedtls_ssl_conf_alpn_protocols( &conf, alpn_list ) ) != 0 )
@@ -2872,6 +2851,27 @@
         goto exit;
     }
 
+#if defined(MBEDTLS_SSL_EXPORT_KEYS)
+    if( opt.eap_tls != 0 )
+    {
+        mbedtls_ssl_set_export_keys_cb( &ssl, eap_tls_key_derivation,
+                                        &eap_tls_keying );
+    }
+    else if( opt.nss_keylog != 0 )
+    {
+        mbedtls_ssl_set_export_keys_cb( &ssl,
+                                        nss_keylog_export,
+                                        NULL );
+    }
+#if defined( MBEDTLS_SSL_DTLS_SRTP )
+    else if( opt.use_srtp != 0 )
+    {
+        mbedtls_ssl_set_export_keys_cb( &ssl, dtls_srtp_key_derivation,
+                                        &dtls_srtp_keying );
+    }
+#endif /* MBEDTLS_SSL_DTLS_SRTP */
+#endif /* MBEDTLS_SSL_EXPORT_KEYS */
+
     io_ctx.ssl = &ssl;
     io_ctx.net = &client_fd;
     mbedtls_ssl_set_bio( &ssl, &io_ctx, send_cb, recv_cb,