Check selector argument count when verifying method

madsmtm · Jan 12, 2023 · 70874fb · 70874fb
1 parent 8c8482f
commit 70874fb
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 7 deletions.
diff --git a/crates/objc2/src/declare.rs b/crates/objc2/src/declare.rs
@@ -229,10 +229,6 @@ method_decl_impl!(A, B, C, D, E, F, G, H, I, J);
 method_decl_impl!(A, B, C, D, E, F, G, H, I, J, K);
 method_decl_impl!(A, B, C, D, E, F, G, H, I, J, K, L);
 
-fn count_args(sel: Sel) -> usize {
-    sel.name().chars().filter(|&c| c == ':').count()
-}
-
 fn method_type_encoding(ret: &Encoding, args: &[Encoding]) -> CString {
     // First two arguments are always self and the selector
     let mut types = format!("{ret}{}{}", <*mut Object>::ENCODING, Sel::ENCODING);
@@ -360,7 +356,7 @@ impl ClassBuilder {
         let enc_args = F::Args::ENCODINGS;
         let enc_ret = F::Ret::ENCODING;
 
-        let sel_args = count_args(sel);
+        let sel_args = sel.number_of_arguments();
         assert_eq!(
             sel_args,
             enc_args.len(),
@@ -417,7 +413,7 @@ impl ClassBuilder {
         let enc_args = F::Args::ENCODINGS;
         let enc_ret = F::Ret::ENCODING;
 
-        let sel_args = count_args(sel);
+        let sel_args = sel.number_of_arguments();
         assert_eq!(
             sel_args,
             enc_args.len(),
@@ -575,7 +571,7 @@ impl ProtocolBuilder {
         Ret: Encode,
     {
         let encs = Args::ENCODINGS;
-        let sel_args = count_args(sel);
+        let sel_args = sel.number_of_arguments();
         assert_eq!(
             sel_args,
             encs.len(),

diff --git a/crates/objc2/src/runtime.rs b/crates/objc2/src/runtime.rs
@@ -197,6 +197,14 @@ impl Sel {
         let name = unsafe { CStr::from_ptr(ptr) };
         str::from_utf8(name.to_bytes()).unwrap()
     }
+
+    pub(crate) fn number_of_arguments(self) -> usize {
+        self.name()
+            .as_bytes()
+            .iter()
+            .filter(|&&b| b == b':')
+            .count()
+    }
 }
 
 // `ffi::sel_isEqual` is just pointer comparison on Apple (the documentation

diff --git a/crates/objc2/src/verify.rs b/crates/objc2/src/verify.rs
@@ -106,6 +106,11 @@ pub(crate) fn verify_method_signature(
         return Err(Inner::MismatchedArgumentsCount(actual_count + remaining, actual_count).into());
     }
 
+    let expected_count = method.name().number_of_arguments();
+    if expected_count != actual_count {
+        return Err(Inner::MismatchedArgumentsCount(expected_count, actual_count).into());
+    }
+
     Ok(())
 }