Merge pull request #63 from gendx/add-debug-allocations

Add feature to track allocations in libtock-rs

.github/workflows/cargo_check.yml

@@ -46,12 +46,30 @@ jobs:
           command: check
           args: --target thumbv7em-none-eabi --release --features debug_ctap
 
+      - name: Check OpenSK panic_console
+        uses: actions-rs/cargo@v1
+        with:
+          command: check
+          args: --target thumbv7em-none-eabi --release --features panic_console
+
+      - name: Check OpenSK debug_allocations
+        uses: actions-rs/cargo@v1
+        with:
+          command: check
+          args: --target thumbv7em-none-eabi --release --features debug_allocations
+
       - name: Check OpenSK debug_ctap,with_ctap1
         uses: actions-rs/cargo@v1
         with:
           command: check
           args: --target thumbv7em-none-eabi --release --features debug_ctap,with_ctap1
 
+      - name: Check OpenSK debug_ctap,with_ctap1,panic_console,debug_allocations
+        uses: actions-rs/cargo@v1
+        with:
+          command: check
+          args: --target thumbv7em-none-eabi --release --features debug_ctap,with_ctap1,panic_console,debug_allocations
+
       - name: Check examples
         uses: actions-rs/cargo@v1
         with:

Cargo.toml

@@ -22,6 +22,7 @@ std = ["cbor/std", "crypto/std", "crypto/derive_debug"]
 debug_ctap = ["crypto/derive_debug"]
 with_ctap1 = ["crypto/with_ctap1"]
 panic_console = ["libtock/panic_console"]
+debug_allocations = ["libtock/debug_allocations"]
 
 [dev-dependencies]
 elf2tab = "0.4.0"

deploy.py

@@ -370,6 +370,14 @@ if __name__ == "__main__":
             "output messages before starting blinking the LEDs on the "
             "board."),
   )
+  app_commands.add_argument(
+      "--debug-allocations",
+      action="append_const",
+      const="debug_allocations",
+      dest="features",
+      help=("The console will be used to output allocator statistics every "
+            "time an allocation/deallocation happens."),
+  )
   app_commands.add_argument(
       "--no-u2f",
       action=RemoveConstAction,

patches/libtock-rs/07-debug_allocations.patch

@@ -0,0 +1,103 @@
+diff --git a/Cargo.toml b/Cargo.toml
+index 386a9ed..af3c5db 100644
+--- a/Cargo.toml
++++ b/Cargo.toml
+@@ -10,6 +10,7 @@ linked_list_allocator = { version = "0.6.6", default-features = false }
+ 
+ [features]
+ panic_console = []
++debug_allocations = []
+ 
+ [dev-dependencies]
+ corepack = { version = "0.4.0", default-features = false, features = ["alloc"] }
+diff --git a/src/entry_point.rs b/src/entry_point.rs
+index 2fe5c40..545d163 100644
+--- a/src/entry_point.rs
++++ b/src/entry_point.rs
+@@ -368,22 +368,82 @@ pub unsafe extern "C" fn rust_start(app_start: usize, stacktop: usize, app_heap_
+     }
+ }
+ 
++#[cfg(feature = "debug_allocations")]
++use core::fmt::Write;
++#[cfg(feature = "debug_allocations")]
++use core::sync::atomic;
++#[cfg(feature = "debug_allocations")]
++use core::sync::atomic::AtomicUsize;
++
+ #[cfg(any(target_arch = "arm", target_arch = "riscv32"))]
+ #[global_allocator]
+-static ALLOCATOR: TockAllocator = TockAllocator;
++static ALLOCATOR: TockAllocator = TockAllocator::new();
+ 
+ static mut HEAP: Heap = Heap::empty();
+ 
+-struct TockAllocator;
++// With the "debug_allocations" feature, we use `AtomicUsize` to store the
++// statistics because:
++// - it is `Sync`, so we can use it in a static object (the allocator),
++// - it implements interior mutability, so we can use it in the allocator
++//   methods (that take an immutable `&self` reference).
++struct TockAllocator {
++    #[cfg(feature = "debug_allocations")]
++    count: AtomicUsize,
++    #[cfg(feature = "debug_allocations")]
++    size: AtomicUsize,
++}
++
++impl TockAllocator {
++    const fn new() -> TockAllocator {
++        TockAllocator {
++            #[cfg(feature = "debug_allocations")]
++            count: AtomicUsize::new(0),
++            #[cfg(feature = "debug_allocations")]
++            size: AtomicUsize::new(0),
++        }
++    }
++}
+ 
+ unsafe impl GlobalAlloc for TockAllocator {
+     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+-        HEAP.allocate_first_fit(layout)
++        let ptr = HEAP
++            .allocate_first_fit(layout)
+             .ok()
+-            .map_or(0 as *mut u8, |allocation| allocation.as_ptr())
++            .map_or(ptr::null_mut(), NonNull::as_ptr);
++        #[cfg(feature = "debug_allocations")]
++        {
++            self.count.fetch_add(1, atomic::Ordering::SeqCst);
++            self.size.fetch_add(layout.size(), atomic::Ordering::SeqCst);
++            writeln!(
++                crate::console::Console::new(),
++                "alloc[{}, {}] = {:?} ({} ptrs, {} bytes)",
++                layout.size(),
++                layout.align(),
++                ptr,
++                self.count.load(atomic::Ordering::SeqCst),
++                self.size.load(atomic::Ordering::SeqCst)
++            )
++            .unwrap();
++        }
++        ptr
+     }
+ 
+     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
++        #[cfg(feature = "debug_allocations")]
++        {
++            self.count.fetch_sub(1, atomic::Ordering::SeqCst);
++            self.size.fetch_sub(layout.size(), atomic::Ordering::SeqCst);
++            writeln!(
++                crate::console::Console::new(),
++                "dealloc[{}, {}] = {:?} ({} ptrs, {} bytes)",
++                layout.size(),
++                layout.align(),
++                ptr,
++                self.count.load(atomic::Ordering::SeqCst),
++                self.size.load(atomic::Ordering::SeqCst)
++            )
++            .unwrap();
++        }
+         HEAP.deallocate(NonNull::new_unchecked(ptr), layout)
+     }
+ }

run_desktop_tests.sh

@@ -28,7 +28,10 @@ echo "Checking that CTAP2 builds properly..."
 cargo check --release --target=thumbv7em-none-eabi
 cargo check --release --target=thumbv7em-none-eabi --features with_ctap1
 cargo check --release --target=thumbv7em-none-eabi --features debug_ctap
+cargo check --release --target=thumbv7em-none-eabi --features panic_console
+cargo check --release --target=thumbv7em-none-eabi --features debug_allocations
 cargo check --release --target=thumbv7em-none-eabi --features debug_ctap,with_ctap1
+cargo check --release --target=thumbv7em-none-eabi --features debug_ctap,with_ctap1,panic_console,debug_allocations
 
 echo "Checking that examples build properly..."
 cargo check --release --target=thumbv7em-none-eabi --examples