Store custom extension callbacks according to their extension types instead of function signature

openssl/src/ssl/callbacks.rs

@@ -562,7 +562,7 @@ pub struct CustomExtAddState<T>(Option<T>);
 #[cfg(ossl111)]
 pub extern "C" fn raw_custom_ext_add<F, T>(
     ssl: *mut ffi::SSL,
-    _: c_uint,
+    ext_type: c_uint,
     context: c_uint,
     out: *mut *const c_uchar,
     outlen: *mut size_t,
@@ -580,9 +580,11 @@ where
 {
     unsafe {
         let ssl = SslRef::from_ptr_mut(ssl);
+        let cb_key =
+            SslContext::get_custom_ext_cb_key(ext_type as u16, super::CustomExtCbType::Add);
         let callback = ssl
             .ssl_context()
-            .ex_data(SslContext::cached_ex_index::<F>())
+            .ex_data(SslContext::cached_custom_ext_ex_index(cb_key))
             .expect("BUG: custom ext add callback missing") as *const F;
         let ectx = ExtensionContext::from_bits_truncate(context);
         let cert = if ectx.contains(ExtensionContext::TLS1_3_CERTIFICATE) {
@@ -640,7 +642,7 @@ pub extern "C" fn raw_custom_ext_free<T>(
 #[cfg(ossl111)]
 pub extern "C" fn raw_custom_ext_parse<F>(
     ssl: *mut ffi::SSL,
-    _: c_uint,
+    ext_type: c_uint,
     context: c_uint,
     input: *const c_uchar,
     inlen: size_t,
@@ -657,9 +659,11 @@ where
 {
     unsafe {
         let ssl = SslRef::from_ptr_mut(ssl);
+        let cb_key =
+            SslContext::get_custom_ext_cb_key(ext_type as u16, super::CustomExtCbType::Parse);
         let callback = ssl
             .ssl_context()
-            .ex_data(SslContext::cached_ex_index::<F>())
+            .ex_data(SslContext::cached_custom_ext_ex_index(cb_key))
             .expect("BUG: custom ext parse callback missing") as *const F;
         let ectx = ExtensionContext::from_bits_truncate(context);
         #[allow(clippy::unnecessary_cast)]

openssl/src/ssl/mod.rs

@@ -554,9 +554,18 @@ impl NameType {
     }
 }
 
+#[cfg(ossl111)]
+enum CustomExtCbType {
+    Add,
+    Parse,
+}
+
 static INDEXES: Lazy<Mutex<HashMap<TypeId, c_int>>> = Lazy::new(|| Mutex::new(HashMap::new()));
 static SSL_INDEXES: Lazy<Mutex<HashMap<TypeId, c_int>>> = Lazy::new(|| Mutex::new(HashMap::new()));
 static SESSION_CTX_INDEX: OnceCell<Index<Ssl, SslContext>> = OnceCell::new();
+#[cfg(ossl111)]
+static CUSTOM_EXT_INDEXES: Lazy<Mutex<HashMap<u64, c_int>>> =
+    Lazy::new(|| Mutex::new(HashMap::new()));
 
 fn try_get_session_ctx_index() -> Result<&'static Index<Ssl, SslContext>, ErrorStack> {
     SESSION_CTX_INDEX.get_or_try_init(Ssl::new_ex_index)
@@ -1642,8 +1651,16 @@ impl SslContextBuilder {
             + Send,
     {
         let ret = unsafe {
-            self.set_ex_data(SslContext::cached_ex_index::<AddFn>(), add_cb);
-            self.set_ex_data(SslContext::cached_ex_index::<ParseFn>(), parse_cb);
+            let add_cb_key = SslContext::get_custom_ext_cb_key(ext_type, CustomExtCbType::Add);
+            let parse_cb_key = SslContext::get_custom_ext_cb_key(ext_type, CustomExtCbType::Parse);
+            self.set_ex_data(
+                SslContext::cached_custom_ext_ex_index::<AddFn>(add_cb_key),
+                add_cb,
+            );
+            self.set_ex_data(
+                SslContext::cached_custom_ext_ex_index::<ParseFn>(parse_cb_key),
+                parse_cb,
+            );
 
             ffi::SSL_CTX_add_custom_ext(
                 self.as_ptr(),
@@ -1827,6 +1844,29 @@ impl SslContext {
         }
     }
 
+    #[cfg(ossl111)]
+    fn get_custom_ext_cb_key(ext_type: u16, cb_type: CustomExtCbType) -> u64 {
+        match cb_type {
+            CustomExtCbType::Add => ext_type as u64,
+            CustomExtCbType::Parse => ext_type as u64 | 0x8000_0000_0000_0000,
+        }
+    }
+
+    #[cfg(ossl111)]
+    fn cached_custom_ext_ex_index<T>(key: u64) -> Index<SslContext, T>
+    where
+        T: 'static + Sync + Send,
+    {
+        unsafe {
+            let idx = *CUSTOM_EXT_INDEXES
+                .lock()
+                .unwrap_or_else(|e| e.into_inner())
+                .entry(key)
+                .or_insert_with(|| SslContext::new_ex_index::<T>().unwrap().as_raw());
+            Index::from_raw(idx)
+        }
+    }
+
     // FIXME should return a result?
     fn cached_ex_index<T>() -> Index<SslContext, T>
     where

openssl/src/ssl/test/mod.rs

@@ -1302,6 +1302,59 @@ fn custom_extensions() {
     assert!(FOUND_EXTENSION.load(Ordering::SeqCst));
 }
 
+#[test]
+#[cfg(ossl111)]
+fn custom_extensions_with_same_callback_signature() {
+    static FOUND_EXTENSION_1: AtomicBool = AtomicBool::new(false);
+    static FOUND_EXTENSION_2: AtomicBool = AtomicBool::new(false);
+
+    fn insert_custom_extension(builder: &mut SslContextBuilder, ext_type: u16, data: Vec<u8>) {
+        builder
+            .add_custom_ext(
+                ext_type,
+                ExtensionContext::CLIENT_HELLO,
+                move |_, _, _| Ok(Some(data.clone())),
+                |_, _, _, _| Ok(()),
+            )
+            .expect("Failed to add custom extension");
+    }
+
+    let mut server = Server::builder();
+    server
+        .ctx()
+        .add_custom_ext(
+            12345,
+            ExtensionContext::CLIENT_HELLO,
+            |_, _, _| -> Result<Option<&'static [u8]>, _> { unreachable!() },
+            move |_, _, data, _| {
+                FOUND_EXTENSION_1.store(data == b"some data".to_vec(), Ordering::SeqCst);
+                Ok(())
+            },
+        )
+        .unwrap();
+    server
+        .ctx()
+        .add_custom_ext(
+            23456,
+            ExtensionContext::CLIENT_HELLO,
+            |_, _, _| -> Result<Option<&'static [u8]>, _> { unreachable!() },
+            move |_, _, data, _| {
+                FOUND_EXTENSION_2.store(data == b"some other data".to_vec(), Ordering::SeqCst);
+                Ok(())
+            },
+        )
+        .unwrap();
+
+    let server = server.build();
+
+    let mut client = server.client();
+    insert_custom_extension(client.ctx(), 12345, b"some data".to_vec());
+    insert_custom_extension(client.ctx(), 23456, b"some other data".to_vec());
+    client.connect();
+
+    assert!(FOUND_EXTENSION_1.load(Ordering::SeqCst));
+    assert!(FOUND_EXTENSION_2.load(Ordering::SeqCst));
+}
 fn _check_kinds() {
     fn is_send<T: Send>() {}
     fn is_sync<T: Sync>() {}