rust: bssl-tls: Refactor the stack iterator As we need to add more stack iterators with different element types, we need to generalise the current iterator and reuse some of the code. Signed-off-by: Xiangfei Ding <xfding@google.com> Change-Id: I44de3afa40fac49bf188c9f71cf0b5f96a6a6964 Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/98792 Reviewed-by: Adam Langley <agl@google.com>

diff --git a/rust/bssl-tls/src/credentials.rs b/rust/bssl-tls/src/credentials.rs
index 2fa6181..b2b4db1 100644
--- a/rust/bssl-tls/src/credentials.rs
+++ b/rust/bssl-tls/src/credentials.rs

@@ -65,7 +65,9 @@
     ffi::{
         Alloc,
         Bio,
+        CryptoBufferWrapper,
         Stack,
+        StackIterator,
         sanitize_slice,
         slice_into_ffi_raw_parts, //
     },
@@ -919,67 +921,55 @@
 /// Certificate chain iterator.
 ///
 /// This iterator will supply the peer leaf certificate as the first element in the chain, if any.
+pub type CertificateChainIterator<'a> = CryptoBufferIterator<'a, Certificate>;
+
 #[derive(Clone, Copy)]
-pub struct CertificateChainIterator<'a> {
-    certs: *const bssl_sys::stack_st_CRYPTO_BUFFER,
-    len: usize,
-    curr: usize,
-    _p: PhantomData<&'a ()>,
+#[doc(hidden)]
+pub struct CryptoBufferIterator<'a, T> {
+    inner: StackIterator<'a, bssl_sys::CRYPTO_BUFFER>,
+    _p: PhantomData<fn() -> T>,
 }
 
-impl<'a> CertificateChainIterator<'a> {
-    /// Safety: caller must ensure that `certs` is outlived by,
-    /// or in other words stays alive as long as, `'a`.
-    pub(crate) unsafe fn new(certs: *const bssl_sys::stack_st_CRYPTO_BUFFER) -> Self {
-        let len = if certs.is_null() {
-            0
-        } else {
-            unsafe {
-                // Safety: `certs` is valid now.
-                bssl_sys::sk_CRYPTO_BUFFER_num(certs)
-            }
-        };
+impl<T: CryptoBufferWrapper> CryptoBufferIterator<'_, T> {
+    /// Safety: caller must ensure that `sk` outlives `'a`.
+    pub(crate) unsafe fn new(sk: *const bssl_sys::stack_st_CRYPTO_BUFFER) -> Self {
         Self {
-            certs,
-            len,
-            curr: 0,
+            inner: unsafe {
+                // Safety: `sk` outlives `'a` per pre-condition.
+                StackIterator::new(sk)
+            },
             _p: PhantomData,
         }
     }
 }
 
-impl<'a> Iterator for CertificateChainIterator<'a> {
-    type Item = Certificate;
+impl<T: CryptoBufferWrapper> Iterator for CryptoBufferIterator<'_, T> {
+    type Item = T;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if self.curr >= self.len {
-            return None;
-        }
-        let cert = unsafe {
-            // Safety: `self.certs` is still valid now and `self.curr` is within the bound.
-            bssl_sys::sk_CRYPTO_BUFFER_value(self.certs, self.curr)
-        };
-        self.curr += 1;
-        let Some(cert) = NonNull::new(cert) else {
-            // Fuse the iterator.
-            self.curr = self.len;
-            return None;
-        };
-        unsafe {
-            // Safety: `cert` is valid here.
-            bssl_sys::CRYPTO_BUFFER_up_ref(cert.as_ptr());
-        }
-        Some(Certificate(cert))
+        self.inner
+            .next()
+            .map(|buf| unsafe {
+                // Safety: we are only bumping the ref-count
+                bssl_sys::CRYPTO_BUFFER_dup_ref(buf)
+            })
+            .and_then(|ptr| NonNull::new(ptr as *mut _))
+            .map(|buf| {
+                unsafe {
+                    // Safety: `buf` is now exclusively owned.
+                    T::from_crypto_buffer(buf)
+                }
+            })
     }
 }
 
-impl ExactSizeIterator for CertificateChainIterator<'_> {
+impl<T: CryptoBufferWrapper> ExactSizeIterator for CryptoBufferIterator<'_, T> {
     fn len(&self) -> usize {
-        self.len - self.curr
+        self.inner.len()
     }
 }
 
-impl FusedIterator for CertificateChainIterator<'_> {}
+impl<T: CryptoBufferWrapper> FusedIterator for CryptoBufferIterator<'_, T> {}
 
 /// Safety: this callback stub must be installed with a context object allocated
 /// as a `Box<dyn VerifyCertificate>`.

diff --git a/rust/bssl-tls/src/ffi.rs b/rust/bssl-tls/src/ffi.rs
index c462fd9..55d4b02 100644
--- a/rust/bssl-tls/src/ffi.rs
+++ b/rust/bssl-tls/src/ffi.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use core::{
+    iter::FusedIterator,
     marker::PhantomData,
     mem::{
         MaybeUninit,
@@ -275,6 +276,11 @@
     }
 }
 
+pub(crate) trait CryptoBufferWrapper {
+    /// Safety: `buf` must be exclusively owned.
+    unsafe fn from_crypto_buffer(buf: core::ptr::NonNull<::bssl_sys::CRYPTO_BUFFER>) -> Self;
+}
+
 pub(crate) unsafe trait BsslStack: Sized {
     type Element: StackElement;
 
@@ -283,6 +289,9 @@
     /// Safety: `this` handle must be a live `stack_st_*` handle.
     unsafe fn size(this: *const Self) -> usize;
 
+    /// Safety: `this` handle must be live and `idx` must be in bounds.
+    unsafe fn index(this: *const Self, idx: usize) -> *const Self::Element;
+
     /// Safety: both `this` and `elem` cannot be aliased.
     unsafe fn push(this: *mut Self, elem: *mut Self::Element);
 
@@ -313,6 +322,13 @@
         }
     }
 
+    unsafe fn index(this: *const Self, idx: usize) -> *const Self::Element {
+        unsafe {
+            // Safety: `this` is valid and live
+            bssl_sys::sk_CRYPTO_BUFFER_value(this, idx)
+        }
+    }
+
     unsafe fn push(this: *mut Self, elem: *mut bssl_sys::CRYPTO_BUFFER) {
         let rc = unsafe {
             // Safety: `this` and `elem` are exclusively owned and valid.
@@ -377,3 +393,60 @@
         ptr
     }
 }
+
+#[derive(Clone, Copy)]
+pub(crate) struct StackIterator<'a, T: StackElement> {
+    sk: *const T::Stack,
+    len: usize,
+    curr: usize,
+    _p: PhantomData<&'a fn() -> T>,
+}
+
+impl<'a, T: StackElement> StackIterator<'a, T> {
+    /// Safety: caller must ensure that `sk` outlives `'a`.
+    pub(crate) unsafe fn new(sk: *const T::Stack) -> Self {
+        let len = if sk.is_null() {
+            0
+        } else {
+            unsafe {
+                // Safety: `sk` is valid now.
+                T::Stack::size(sk)
+            }
+        };
+        Self {
+            sk,
+            len,
+            curr: 0,
+            _p: PhantomData,
+        }
+    }
+}
+
+impl<'a, T: StackElement> Iterator for StackIterator<'a, T> {
+    type Item = *const T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.curr >= self.len {
+            return None;
+        }
+        let elem = unsafe {
+            // Safety: `self.sk` is still valid now and `self.curr` is within the bound.
+            T::Stack::index(self.sk, self.curr)
+        };
+        self.curr += 1;
+        if elem.is_null() {
+            // Fuse the iterator.
+            self.curr = self.len;
+            return None;
+        }
+        Some(elem)
+    }
+}
+
+impl<T: StackElement> ExactSizeIterator for StackIterator<'_, T> {
+    fn len(&self) -> usize {
+        self.len - self.curr
+    }
+}
+
+impl<T: StackElement> FusedIterator for StackIterator<'_, T> {}

diff --git a/rust/bssl-tls/src/macros.rs b/rust/bssl-tls/src/macros.rs
index eac9622..d9fb5ae 100644
--- a/rust/bssl-tls/src/macros.rs
+++ b/rust/bssl-tls/src/macros.rs

@@ -77,6 +77,12 @@
             }
         }
 
+        impl $crate::ffi::CryptoBufferWrapper for $name {
+            unsafe fn from_crypto_buffer(buf: ::core::ptr::NonNull<::bssl_sys::CRYPTO_BUFFER>) -> Self {
+                Self(buf)
+            }
+        }
+
         impl ::core::ops::Deref for $name {
             type Target = [u8];
             fn deref(&self) -> &Self::Target {