chacha: Clarify the dispatching to assembly implementations.

Clarify that there are 4 different functions, one for each architecture,
which have different capabilities w.r.t. overlapping input/output.

we plan to extend this to support all the details of the dispatching
that is currently within each assembly implementation of
`ChaCha20_ctr32` when that dispatching logic is moved to Rust.

src/aead/chacha.rs

@@ -15,17 +15,23 @@
 
 use super::{overlapping, quic::Sample, Nonce};
 use crate::cpu;
+use cfg_if::cfg_if;
 
-#[cfg(any(
-    test,
-    not(any(
-        all(target_arch = "aarch64", target_endian = "little"),
-        all(target_arch = "arm", target_endian = "little"),
-        target_arch = "x86",
-        target_arch = "x86_64"
-    ))
-))]
-mod fallback;
+cfg_if! {
+    if #[cfg(any(
+                all(target_arch = "aarch64", target_endian = "little"),
+                all(target_arch = "arm", target_endian = "little"),
+                target_arch = "x86",
+                target_arch = "x86_64"
+            ))] {
+        #[macro_use]
+        mod ffi;
+        #[cfg(test)]
+        mod fallback;
+    } else {
+        mod fallback;
+    }
+}
 
 use crate::polyfill::ArraySplitMap;
 
@@ -74,52 +80,31 @@ impl Key {
     }
 
     #[inline(always)]
-    pub(super) fn encrypt(&self, counter: Counter, in_out: Overlapping<'_>, _cpu: cpu::Features) {
-        #[cfg(any(
-            all(target_arch = "aarch64", target_endian = "little"),
-            all(target_arch = "arm", target_endian = "little"),
-            target_arch = "x86",
-            target_arch = "x86_64"
-        ))]
-        #[inline(always)]
-        pub(super) fn ChaCha20_ctr32(key: &Key, counter: Counter, in_out: Overlapping<'_>) {
-            // XXX: The x86 and at least one branch of the ARM assembly language
-            // code doesn't allow overlapping input and output unless they are
-            // exactly overlapping. TODO: Figure out which branch of the ARM code
-            // has this limitation and come up with a better solution.
-            //
-            // https://rt.openssl.org/Ticket/Display.html?id=4362
-            #[cfg(not(any(
-                all(target_arch = "aarch64", target_endian = "little"),
-                target_arch = "x86_64"
-            )))]
-            let in_out = Overlapping::from(in_out.copy_within());
-
-            let (input, output, len) = in_out.into_input_output_len();
-
-            // There's no need to worry if `counter` is incremented because it is
-            // owned here and we drop immediately after the call.
-            prefixed_extern! {
-                fn ChaCha20_ctr32(
-                    out: *mut u8,
-                    in_: *const u8,
-                    in_len: crate::c::size_t,
-                    key: &[u32; KEY_LEN / 4],
-                    counter: &Counter,
-                );
+    pub(super) fn encrypt(&self, counter: Counter, in_out: Overlapping<'_>, cpu: cpu::Features) {
+        // XXX: The x86 and at least one branch of the ARM assembly language
+        // code doesn't allow overlapping input and output unless they are
+        // "in place". See https://rt.openssl.org/Ticket/Display.html?id=4362.
+        cfg_if! {
+            if #[cfg(all(target_arch = "aarch64", target_endian = "little"))] {
+                chacha20_ctr32_ffi!(
+                    unsafe { (cpu::Features, Overlapping<'_>) => ChaCha20_ctr32 },
+                    self, counter, in_out, cpu)
+            } else if #[cfg(all(target_arch = "arm", target_endian = "little"))] {
+                chacha20_ctr32_ffi!(
+                    unsafe { (cpu::Features, &mut [u8]) => ChaCha20_ctr32 },
+                    self, counter, in_out.copy_within(), cpu)
+            } else if #[cfg(target_arch = "x86")] {
+                chacha20_ctr32_ffi!(
+                    unsafe { (cpu::Features, &mut [u8]) => ChaCha20_ctr32 },
+                    self, counter, in_out.copy_within(), cpu)
+            } else if #[cfg(target_arch = "x86_64")] {
+                chacha20_ctr32_ffi!(
+                    unsafe { (cpu::Features, Overlapping<'_>) => ChaCha20_ctr32 },
+                    self, counter, in_out, cpu)
+            } else {
+                fallback::ChaCha20_ctr32(key, counter, in_out)
             }
-            unsafe { ChaCha20_ctr32(output, input, len, key.words_less_safe(), &counter) }
         }
-
-        #[cfg(not(any(
-            all(target_arch = "aarch64", target_endian = "little"),
-            all(target_arch = "arm", target_endian = "little"),
-            target_arch = "x86",
-            target_arch = "x86_64"
-        )))]
-        use fallback::ChaCha20_ctr32;
-
-        ChaCha20_ctr32(self, counter, in_out)
     }
 
     #[inline]

src/aead/chacha/ffi.rs

@@ -0,0 +1,56 @@
+// Copyright 2016-2025 Brian Smith.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHORS DISCLAIM ALL WARRANTIES
+// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY
+// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
+// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
+// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+use super::{super::overlapping::Overlapping, Counter, Key};
+
+// `unsafe { (C, InOut) => f }` means that the function `f` is safe to call
+// iff CPU features `C` are available and the input type is `InOut`. If `f`
+// supports overlapping input/output then `InOut` should be
+// `Overlapping<'_, u8>`; otherwise it should be `&mut [u8]`.
+macro_rules! chacha20_ctr32_ffi {
+    ( unsafe { ($Cpu:ty, $InOut:ty) => $f:ident },
+      $key:expr, $counter:expr, $in_out:expr, $cpu:expr ) => {{
+        prefixed_extern! {
+            fn $f(
+                out: *mut u8,
+                in_: *const u8,
+                in_len: crate::c::size_t,
+                key: &[u32; 8],
+                counter: &crate::aead::chacha::Counter,
+            );
+        }
+        // SAFETY: The user asserts that $f has the signature above and is safe
+        // to call if additionally we have a value of type `$Cpu` and an in/out
+        // value of the indicated type, which we do.
+        unsafe {
+            crate::aead::chacha::ffi::chacha20_ctr32_ffi::<$InOut, $Cpu>(
+                $key, $counter, $in_out, $cpu, $f,
+            )
+        }
+    }};
+}
+
+pub(super) unsafe fn chacha20_ctr32_ffi<'o, InOut: 'o + Into<Overlapping<'o, u8>>, Cpu>(
+    key: &Key,
+    counter: Counter,
+    in_out: InOut,
+    cpu: Cpu,
+    f: unsafe extern "C" fn(*mut u8, *const u8, crate::c::size_t, &[u32; 8], &Counter),
+) {
+    let in_out: Overlapping<'_, u8> = in_out.into();
+    let (input, output, len) = in_out.into_input_output_len();
+    let key = key.words_less_safe();
+    let _: Cpu = cpu;
+    unsafe { f(output, input, len, key, &counter) }
+}