diff --git a/detectors/Cargo.lock b/detectors/Cargo.lock index a324cab0..1118edf2 100644 --- a/detectors/Cargo.lock +++ b/detectors/Cargo.lock @@ -2901,6 +2901,16 @@ dependencies = [ "tinyvec", ] +[[package]] +name = "unsafe-unwrap" +version = "0.1.0" +dependencies = [ + "clippy_utils", + "if_chain", + "scout-audit-clippy-wrappers-ink", + "scout-audit-dylint-linting", +] + [[package]] name = "untrusted" version = "0.9.0" diff --git a/detectors/unsafe-unwrap/Cargo.toml b/detectors/unsafe-unwrap/Cargo.toml new file mode 100644 index 00000000..32cec230 --- /dev/null +++ b/detectors/unsafe-unwrap/Cargo.toml @@ -0,0 +1,16 @@ +[package] +edition = "2021" +name = "unsafe-unwrap" +version = "0.1.0" + +[lib] +crate-type = ["cdylib"] + +[dependencies] +clippy_utils = { workspace = true } +clippy_wrappers = { workspace = true } +if_chain = { workspace = true } +scout-audit-dylint-linting = { workspace = true } + +[package.metadata.rust-analyzer] +rustc_private = true diff --git a/detectors/unsafe-unwrap/src/lib.rs b/detectors/unsafe-unwrap/src/lib.rs new file mode 100644 index 00000000..45b4ccc7 --- /dev/null +++ b/detectors/unsafe-unwrap/src/lib.rs @@ -0,0 +1,364 @@ +#![feature(rustc_private)] +#![allow(clippy::enum_variant_names)] + +extern crate rustc_ast; +extern crate rustc_hir; +extern crate rustc_span; + +use clippy_utils::higher; +use clippy_wrappers::span_lint_and_help; +use if_chain::if_chain; +use rustc_hir::{ + def::Res, + def_id::LocalDefId, + intravisit::{walk_expr, FnKind, Visitor}, + BinOpKind, Body, Expr, ExprKind, FnDecl, HirId, PathSegment, QPath, UnOp, +}; +use rustc_lint::{LateContext, LateLintPass}; +use rustc_span::{sym, Span, Symbol}; +use std::{collections::HashSet, hash::Hash}; + +const LINT_MESSAGE: &str = "Unsafe usage of `unwrap`"; +const PANIC_INDUCING_FUNCTIONS: [&str; 2] = ["panic", "bail"]; + +scout_audit_dylint_linting::declare_late_lint! { + /// ### What it does + /// Checks for usage of `unwrap` + /// + /// ### Why is this bad? + /// `unwrap` might panic if the result value is an error or `None`. + /// + /// ### Example + /// ```rust + /// // example code where a warning is issued + /// fn main() { + /// let result = result_fn().unwrap("error"); + /// } + /// + /// fn result_fn() -> Result { + /// Err(Error::new(ErrorKind::Other, "error")) + /// } + /// ``` + /// Use instead: + /// ```rust + /// // example code that does not raise a warning + /// fn main() { + /// let result = if let Ok(result) = result_fn() { + /// result + /// } + /// } + /// + /// fn result_fn() -> Result { + /// Err(Error::new(ErrorKind::Other, "error")) + /// } + /// ``` + pub UNSAFE_UNWRAP, + Warn, + LINT_MESSAGE, + { + name: "Unsafe Unwrap", + long_message: "This vulnerability class pertains to the inappropriate usage of the unwrap method in Rust, which is commonly employed for error handling. The unwrap method retrieves the inner value of an Option or Result, but if an error or None occurs, it triggers a panic and crashes the program. ", + severity: "Medium", + help: "https://coinfabrik.github.io/scout-soroban/docs/detectors/unsafe-unwrap", + vulnerability_class: "Validations and error handling", + } +} + +/// Represents the type of check performed on method call expressions to determine their safety or behavior. +#[derive(Clone, Copy, Hash, Eq, PartialEq)] +enum CheckType { + IsSome, + IsNone, + IsOk, + IsErr, +} + +impl CheckType { + fn from_method_name(name: Symbol) -> Option { + match name.as_str() { + "is_some" => Some(Self::IsSome), + "is_none" => Some(Self::IsNone), + "is_ok" => Some(Self::IsOk), + "is_err" => Some(Self::IsErr), + _ => None, + } + } + + fn inverse(self) -> Self { + match self { + Self::IsSome => Self::IsNone, + Self::IsNone => Self::IsSome, + Self::IsOk => Self::IsErr, + Self::IsErr => Self::IsOk, + } + } + + /// Determines if the check type implies execution should halt, such as in error conditions. + fn should_halt_execution(self) -> bool { + matches!(self, Self::IsNone | Self::IsErr) + } + + /// Determines if it is safe to unwrap the value without further checks, i.e., the value is present. + fn is_safe_to_unwrap(self) -> bool { + matches!(self, Self::IsSome | Self::IsOk) + } +} + +/// Represents a conditional checker that is used to analyze `if` or `if let` expressions. +#[derive(Clone, Copy, Hash, Eq, PartialEq)] +struct ConditionalChecker { + check_type: CheckType, + checked_expr_hir_id: HirId, +} + +impl ConditionalChecker { + /// Handle te condition of the `if` or `if let` expression. + fn handle_condition(condition: &Expr<'_>, inverse: bool) -> HashSet { + if_chain! { + if let ExprKind::MethodCall(path_segment, receiver, _, _) = condition.kind; + if let Some(check_type) = CheckType::from_method_name(path_segment.ident.name); + if let ExprKind::Path(QPath::Resolved(_, checked_expr_path)) = receiver.kind; + if let Res::Local(checked_expr_hir_id) = checked_expr_path.res; + then { + let check_type = if inverse { check_type.inverse() } else { check_type }; + return std::iter::once(Self { check_type, checked_expr_hir_id }).collect(); + } + } + HashSet::new() + } + + /// Constructs a ConditionalChecker from an expression if it matches a method call with a valid CheckType. + fn from_expression(condition: &Expr<'_>) -> HashSet { + match condition.kind { + // Single `not` expressions are supported + ExprKind::Unary(op, condition) => Self::handle_condition(condition, op == UnOp::Not), + // Multiple `or` expressions are supported + ExprKind::Binary(op, left_condition, right_condition) if op.node == BinOpKind::Or => { + let mut result = Self::from_expression(left_condition); + result.extend(Self::from_expression(right_condition)); + result + } + ExprKind::MethodCall(..) => Self::handle_condition(condition, false), + _ => HashSet::new(), + } + } +} + +/// Main unsafe-unwrap visitor +struct UnsafeUnwrapVisitor<'a, 'tcx> { + cx: &'a LateContext<'tcx>, + conditional_checker: HashSet, + checked_exprs: HashSet, +} + +impl UnsafeUnwrapVisitor<'_, '_> { + fn is_panic_inducing_call(&self, func: &Expr<'_>) -> bool { + if let ExprKind::Path(QPath::Resolved(_, path)) = &func.kind { + return PANIC_INDUCING_FUNCTIONS.iter().any(|&func| { + path.segments + .iter() + .any(|segment| segment.ident.name.as_str().contains(func)) + }); + } + false + } + + fn get_unwrap_info(&self, receiver: &Expr<'_>) -> Option { + if_chain! { + if let ExprKind::Path(QPath::Resolved(_, path)) = &receiver.kind; + if let Res::Local(hir_id) = path.res; + then { + return Some(hir_id); + } + } + None + } + + fn set_conditional_checker(&mut self, conditional_checkers: &HashSet) { + for checker in conditional_checkers { + self.conditional_checker.insert(*checker); + if checker.check_type.is_safe_to_unwrap() { + self.checked_exprs.insert(checker.checked_expr_hir_id); + } + } + } + + fn reset_conditional_checker(&mut self, conditional_checkers: HashSet) { + for checker in conditional_checkers { + if checker.check_type.is_safe_to_unwrap() { + self.checked_exprs.remove(&checker.checked_expr_hir_id); + } + self.conditional_checker.remove(&checker); + } + } + + /// Process conditional expressions to determine if they should halt execution. + fn handle_if_expressions(&mut self) { + self.conditional_checker.iter().for_each(|checker| { + if checker.check_type.should_halt_execution() { + self.checked_exprs.insert(checker.checked_expr_hir_id); + } + }); + } + + fn is_literal_or_composed_of_literals(&self, expr: &Expr<'_>) -> bool { + let mut stack = vec![expr]; + + while let Some(current_expr) = stack.pop() { + match current_expr.kind { + ExprKind::Lit(_) => continue, // A literal is fine, continue processing. + ExprKind::Tup(elements) | ExprKind::Array(elements) => { + stack.extend(elements); + } + ExprKind::Struct(_, fields, _) => { + for field in fields { + stack.push(field.expr); + } + } + ExprKind::Repeat(element, _) => { + stack.push(element); + } + _ => return false, // If any element is not a literal or a compound of literals, return false. + } + } + + true // If the stack is emptied without finding a non-literal, all elements are literals. + } + + fn is_method_call_unsafe( + &self, + path_segment: &PathSegment, + receiver: &Expr, + args: &[Expr], + ) -> bool { + if path_segment.ident.name == sym::unwrap { + return self + .get_unwrap_info(receiver) + .map_or(true, |id| !self.checked_exprs.contains(&id)); + } + + args.iter().any(|arg| self.contains_unsafe_method_call(arg)) + || self.contains_unsafe_method_call(receiver) + } + + fn contains_unsafe_method_call(&self, expr: &Expr) -> bool { + match &expr.kind { + ExprKind::MethodCall(path_segment, receiver, args, _) => { + self.is_method_call_unsafe(path_segment, receiver, args) + } + _ => false, + } + } +} + +impl<'a, 'tcx> Visitor<'tcx> for UnsafeUnwrapVisitor<'a, 'tcx> { + fn visit_local(&mut self, local: &'tcx rustc_hir::Local<'tcx>) { + if let Some(init) = local.init { + match init.kind { + ExprKind::MethodCall(path_segment, receiver, args, _) => { + if self.is_method_call_unsafe(path_segment, receiver, args) { + span_lint_and_help( + self.cx, + UNSAFE_UNWRAP, + local.span, + LINT_MESSAGE, + None, + "Please, use a custom error instead of `unwrap`", + ); + } + } + ExprKind::Call(func, args) => { + if let ExprKind::Path(QPath::Resolved(_, path)) = func.kind { + let is_some_or_ok = path + .segments + .iter() + .any(|segment| matches!(segment.ident.name, sym::Some | sym::Ok)); + let all_literals = args + .iter() + .all(|arg| self.is_literal_or_composed_of_literals(arg)); + if is_some_or_ok && all_literals { + self.checked_exprs.insert(local.pat.hir_id); + } + } + } + _ => {} + } + } + } + + fn visit_expr(&mut self, expr: &'tcx Expr<'_>) { + // Skip macro-generated code early + if expr.span.from_expansion() { + return; + } + + // If we are inside an `if` or `if let` expression, we analyze its body + if !self.conditional_checker.is_empty() { + match &expr.kind { + ExprKind::Ret(..) => self.handle_if_expressions(), + ExprKind::Call(func, _) if self.is_panic_inducing_call(func) => { + self.handle_if_expressions() + } + _ => {} + } + } + + // Find `if` or `if let` expressions + if let Some(higher::IfOrIfLet { + cond, + then: if_expr, + r#else: _, + }) = higher::IfOrIfLet::hir(expr) + { + // If we are interested in the condition (if it is a CheckType) we traverse the body. + let conditional_checker = ConditionalChecker::from_expression(cond); + self.set_conditional_checker(&conditional_checker); + walk_expr(self, if_expr); + self.reset_conditional_checker(conditional_checker); + return; + } + + // If we find an unsafe `unwrap`, we raise a warning + if_chain! { + if let ExprKind::MethodCall(path_segment, receiver, _, _) = &expr.kind; + if path_segment.ident.name == sym::unwrap; + then { + let receiver_hir_id = self.get_unwrap_info(receiver); + // If the receiver is `None`, then we asume that the `unwrap` is unsafe + let is_checked_safe = receiver_hir_id.map_or(false, |id| self.checked_exprs.contains(&id)); + if !is_checked_safe { + span_lint_and_help( + self.cx, + UNSAFE_UNWRAP, + expr.span, + LINT_MESSAGE, + None, + "Please, use a custom error instead of `unwrap`", + ); + } + } + } + + walk_expr(self, expr); + } +} + +impl<'tcx> LateLintPass<'tcx> for UnsafeUnwrap { + fn check_fn( + &mut self, + cx: &LateContext<'tcx>, + _: FnKind<'tcx>, + _: &'tcx FnDecl<'tcx>, + body: &'tcx Body<'tcx>, + _: Span, + _: LocalDefId, + ) { + let mut visitor = UnsafeUnwrapVisitor { + cx, + checked_exprs: HashSet::new(), + conditional_checker: HashSet::new(), + }; + + walk_expr(&mut visitor, body.value); + } +} diff --git a/test-cases/Cargo.lock b/test-cases/Cargo.lock index e9947675..081dc5ed 100644 --- a/test-cases/Cargo.lock +++ b/test-cases/Cargo.lock @@ -1669,6 +1669,70 @@ dependencies = [ "sp-runtime", ] +[[package]] +name = "pallet-example-basic-remediated-1" +version = "27.0.0" +dependencies = [ + "frame-benchmarking", + "frame-support", + "frame-system", + "log", + "pallet-balances", + "parity-scale-codec", + "scale-info", + "sp-core", + "sp-io", + "sp-runtime", +] + +[[package]] +name = "pallet-example-basic-remediated-2" +version = "27.0.0" +dependencies = [ + "frame-benchmarking", + "frame-support", + "frame-system", + "log", + "pallet-balances", + "parity-scale-codec", + "scale-info", + "sp-core", + "sp-io", + "sp-runtime", +] + +[[package]] +name = "pallet-example-basic-remediated-3" +version = "27.0.0" +dependencies = [ + "frame-benchmarking", + "frame-support", + "frame-system", + "log", + "pallet-balances", + "parity-scale-codec", + "scale-info", + "sp-core", + "sp-io", + "sp-runtime", +] + +[[package]] +name = "pallet-example-basic-vulnerable-1" +version = "27.0.0" +dependencies = [ + "frame-benchmarking", + "frame-support", + "frame-system", + "log", + "pallet-balances", + "parity-scale-codec", + "scale-info", + "sp-core", + "sp-io", + "sp-runtime", +] + [[package]] name = "pallet-integer-overflow-or-underflow-remediated-1" version = "27.0.0" diff --git a/test-cases/Cargo.toml b/test-cases/Cargo.toml index 845a7ea3..5b10129a 100644 --- a/test-cases/Cargo.toml +++ b/test-cases/Cargo.toml @@ -38,3 +38,8 @@ strip = "symbols" [profile.release-with-logs] debug-assertions = true inherits = "release" + +[workspace.metadata.dylint] +libraries = [ + { path = "/Users/josegarcia/Desktop/coinfabrik/scout-substrate/detectors/unsafe-unwrap" }, +] diff --git a/test-cases/unsafe-unwrap/remediated/remediated-1/Cargo.toml b/test-cases/unsafe-unwrap/remediated/remediated-1/Cargo.toml new file mode 100644 index 00000000..c455a7c2 --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-1/Cargo.toml @@ -0,0 +1,41 @@ +[package] +name = "pallet-example-basic-remediated-1" +version = "27.0.0" +authors.workspace = true +edition.workspace = true +license = "MIT-0" +description = "FRAME example pallet" +readme = "README.md" +publish = false + +[package.metadata.docs.rs] +targets = ["x86_64-unknown-linux-gnu"] + +[dependencies] +codec = { workspace = true } +log = { workspace = true } +scale-info = { features = ["derive"], workspace = true } +frame-benchmarking = { optional = true, workspace = true } +frame-support = { workspace = true } +frame-system = { workspace = true } +pallet-balances = { workspace = true } +sp-io = { workspace = true } +sp-runtime = { workspace = true } + +[dev-dependencies] +sp-core = { workspace = true } + +[features] +default = ["std"] +std = [ + "codec/std", + "frame-benchmarking?/std", + "frame-support/std", + "frame-system/std", + "log/std", + "pallet-balances/std", + "scale-info/std", + "sp-core/std", + "sp-io/std", + "sp-runtime/std", +] diff --git a/test-cases/unsafe-unwrap/remediated/remediated-1/src/lib.rs b/test-cases/unsafe-unwrap/remediated/remediated-1/src/lib.rs new file mode 100644 index 00000000..a6dd635c --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-1/src/lib.rs @@ -0,0 +1,105 @@ +// lib.rs +#![cfg_attr(not(feature = "std"), no_std)] + +extern crate alloc; + +#[cfg(test)] +mod tests; + +pub use pallet::*; + +pub mod weights; +pub use weights::*; + +#[frame_support::pallet] +pub mod pallet { + use super::*; + use frame_support::{pallet_prelude::*, traits::BuildGenesisConfig}; + use frame_system::pallet_prelude::*; + + #[pallet::config] + pub trait Config: pallet_balances::Config + frame_system::Config { + type RuntimeEvent: From> + IsType<::RuntimeEvent>; + /// Weight information for extrinsics in this pallet. + type WeightInfo: WeightInfo; + } + + #[pallet::pallet] + pub struct Pallet(_); + + #[pallet::storage] + #[pallet::getter(fn example_storage)] + pub type ExampleStorage = StorageValue<_, u32>; + + #[pallet::hooks] + impl Hooks> for Pallet { + // Add on_initialize if we need any per-block logic + fn on_initialize(_n: BlockNumberFor) -> Weight { + Weight::zero() + } + } + + #[pallet::call(weight(::WeightInfo))] + impl Pallet { + #[pallet::call_index(0)] + pub fn unsafe_get_storage(origin: OriginFor) -> DispatchResult { + let who = ensure_signed(origin)?; + let value = ExampleStorage::::get().unwrap_or_default(); + Self::deposit_event(Event::UnsafeGetStorage { who, value }); + Ok(()) + } + + #[pallet::call_index(1)] + pub fn set_storage(origin: OriginFor, new_value: u32) -> DispatchResult { + let who = ensure_signed(origin)?; + ExampleStorage::::put(new_value); + Self::deposit_event(Event::StorageSet { + who, + value: new_value, + }); + Ok(()) + } + + #[pallet::call_index(2)] + pub fn clear_storage(origin: OriginFor) -> DispatchResult { + let who = ensure_signed(origin)?; + ExampleStorage::::kill(); + Self::deposit_event(Event::StorageCleared { who }); + Ok(()) + } + } + + #[pallet::event] + #[pallet::generate_deposit(pub(super) fn deposit_event)] + pub enum Event { + /// Storage was accessed through unsafe getter + UnsafeGetStorage { who: T::AccountId, value: u32 }, + /// Storage value was set + StorageSet { who: T::AccountId, value: u32 }, + /// Storage was cleared + StorageCleared { who: T::AccountId }, + } + + #[pallet::error] + pub enum Error { + /// Value is not initialized + NotInitialized, + } + + #[pallet::genesis_config] + #[derive(frame_support::DefaultNoBound)] + pub struct GenesisConfig { + pub initial_value: Option, + #[serde(skip)] + pub _phantom: PhantomData, + } + + #[pallet::genesis_build] + impl BuildGenesisConfig for GenesisConfig { + fn build(&self) { + if let Some(value) = self.initial_value { + ExampleStorage::::put(value); + } + } + } +} diff --git a/test-cases/unsafe-unwrap/remediated/remediated-1/src/tests.rs b/test-cases/unsafe-unwrap/remediated/remediated-1/src/tests.rs new file mode 100644 index 00000000..302383da --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-1/src/tests.rs @@ -0,0 +1,106 @@ +use crate as pallet_example_basic; +use crate::*; +use frame_support::{assert_ok, derive_impl, dispatch::GetDispatchInfo}; +use sp_core::H256; +use sp_runtime::{ + traits::{BlakeTwo256, IdentityLookup}, + BuildStorage, +}; + +type Block = frame_system::mocking::MockBlock; + +frame_support::construct_runtime!( + pub enum Test + { + System: frame_system, + Balances: pallet_balances, + Example: pallet_example_basic, + } +); + +#[derive_impl(frame_system::config_preludes::TestDefaultConfig)] +impl frame_system::Config for Test { + type BaseCallFilter = frame_support::traits::Everything; + type BlockWeights = (); + type BlockLength = (); + type DbWeight = (); + type RuntimeOrigin = RuntimeOrigin; + type Nonce = u64; + type Hash = H256; + type RuntimeCall = RuntimeCall; + type Hashing = BlakeTwo256; + type AccountId = u64; + type Lookup = IdentityLookup; + type Block = Block; + type RuntimeEvent = RuntimeEvent; + type Version = (); + type PalletInfo = PalletInfo; + type AccountData = pallet_balances::AccountData; + type OnNewAccount = (); + type OnKilledAccount = (); + type SystemWeightInfo = (); + type SS58Prefix = (); + type OnSetCode = (); + type MaxConsumers = frame_support::traits::ConstU32<16>; +} + +#[derive_impl(pallet_balances::config_preludes::TestDefaultConfig)] +impl pallet_balances::Config for Test { + type AccountStore = System; +} + +impl Config for Test { + type RuntimeEvent = RuntimeEvent; + type WeightInfo = (); +} + +pub fn new_test_ext() -> sp_io::TestExternalities { + let t = RuntimeGenesisConfig { + system: Default::default(), + balances: Default::default(), + example: pallet_example_basic::GenesisConfig { + initial_value: Some(0), + _phantom: Default::default(), + }, + } + .build_storage() + .unwrap(); + t.into() +} + +#[test] +fn weights_work() { + let call = pallet_example_basic::Call::::unsafe_get_storage {}; + let info = call.get_dispatch_info(); + assert!(info.weight.ref_time() > 0); + assert_eq!( + info.weight, + ::WeightInfo::unsafe_get_storage() + ); +} + +#[test] +fn unsafe_get_storage_works() { + new_test_ext().execute_with(|| { + ExampleStorage::::put(42u32); + assert_ok!(Example::unsafe_get_storage(RuntimeOrigin::signed(1))); + }); +} + +#[test] +fn set_storage_works() { + new_test_ext().execute_with(|| { + let test_val = 133u32; + assert_ok!(Example::set_storage(RuntimeOrigin::signed(1), test_val)); + assert_eq!(ExampleStorage::::get(), Some(test_val)); + }); +} + +#[test] +fn clear_storage_works() { + new_test_ext().execute_with(|| { + ExampleStorage::::put(42u32); + assert_ok!(Example::clear_storage(RuntimeOrigin::signed(1))); + assert_eq!(ExampleStorage::::get(), None); + }); +} diff --git a/test-cases/unsafe-unwrap/remediated/remediated-1/src/weights.rs b/test-cases/unsafe-unwrap/remediated/remediated-1/src/weights.rs new file mode 100644 index 00000000..e9608480 --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-1/src/weights.rs @@ -0,0 +1,43 @@ +#![cfg_attr(rustfmt, rustfmt_skip)] +#![allow(unused_parens)] +#![allow(unused_imports)] + +use frame_support::{traits::Get, weights::{Weight, constants::RocksDbWeight}}; +use core::marker::PhantomData; + +pub trait WeightInfo { + fn unsafe_get_storage() -> Weight; + fn set_storage() -> Weight; + fn clear_storage() -> Weight; +} + +pub struct SubstrateWeight(PhantomData); +impl WeightInfo for SubstrateWeight { + fn unsafe_get_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads(1_u64)) + } + fn set_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads_writes(1_u64, 1_u64)) + } + fn clear_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads_writes(1_u64, 1_u64)) + } +} + +impl WeightInfo for () { + fn unsafe_get_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads(1_u64)) + } + fn set_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads_writes(1_u64, 1_u64)) + } + fn clear_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads_writes(1_u64, 1_u64)) + } +} diff --git a/test-cases/unsafe-unwrap/remediated/remediated-2/Cargo.toml b/test-cases/unsafe-unwrap/remediated/remediated-2/Cargo.toml new file mode 100644 index 00000000..2305b69f --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-2/Cargo.toml @@ -0,0 +1,41 @@ +[package] +name = "pallet-example-basic-remediated-2" +version = "27.0.0" +authors.workspace = true +edition.workspace = true +license = "MIT-0" +description = "FRAME example pallet" +readme = "README.md" +publish = false + +[package.metadata.docs.rs] +targets = ["x86_64-unknown-linux-gnu"] + +[dependencies] +codec = { workspace = true } +log = { workspace = true } +scale-info = { features = ["derive"], workspace = true } +frame-benchmarking = { optional = true, workspace = true } +frame-support = { workspace = true } +frame-system = { workspace = true } +pallet-balances = { workspace = true } +sp-io = { workspace = true } +sp-runtime = { workspace = true } + +[dev-dependencies] +sp-core = { workspace = true } + +[features] +default = ["std"] +std = [ + "codec/std", + "frame-benchmarking?/std", + "frame-support/std", + "frame-system/std", + "log/std", + "pallet-balances/std", + "scale-info/std", + "sp-core/std", + "sp-io/std", + "sp-runtime/std", +] diff --git a/test-cases/unsafe-unwrap/remediated/remediated-2/src/lib.rs b/test-cases/unsafe-unwrap/remediated/remediated-2/src/lib.rs new file mode 100644 index 00000000..d0f9e022 --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-2/src/lib.rs @@ -0,0 +1,108 @@ +// lib.rs +#![cfg_attr(not(feature = "std"), no_std)] + +extern crate alloc; + +#[cfg(test)] +mod tests; + +pub use pallet::*; + +pub mod weights; +pub use weights::*; + +#[frame_support::pallet] +pub mod pallet { + use super::*; + use frame_support::{pallet_prelude::*, traits::BuildGenesisConfig}; + use frame_system::pallet_prelude::*; + + #[pallet::config] + pub trait Config: pallet_balances::Config + frame_system::Config { + type RuntimeEvent: From> + IsType<::RuntimeEvent>; + /// Weight information for extrinsics in this pallet. + type WeightInfo: WeightInfo; + } + + #[pallet::pallet] + pub struct Pallet(_); + + #[pallet::storage] + #[pallet::getter(fn example_storage)] + pub type ExampleStorage = StorageValue<_, u32>; + + #[pallet::hooks] + impl Hooks> for Pallet { + // Add on_initialize if we need any per-block logic + fn on_initialize(_n: BlockNumberFor) -> Weight { + Weight::zero() + } + } + + #[pallet::call(weight(::WeightInfo))] + impl Pallet { + #[pallet::call_index(0)] + pub fn unsafe_get_storage(origin: OriginFor) -> DispatchResult { + let who = ensure_signed(origin)?; + let example_storage = ExampleStorage::::get(); + if example_storage.is_some() { + let value = example_storage.unwrap(); + Self::deposit_event(Event::UnsafeGetStorage { who, value }); + } + Ok(()) + } + + #[pallet::call_index(1)] + pub fn set_storage(origin: OriginFor, new_value: u32) -> DispatchResult { + let who = ensure_signed(origin)?; + ExampleStorage::::put(new_value); + Self::deposit_event(Event::StorageSet { + who, + value: new_value, + }); + Ok(()) + } + + #[pallet::call_index(2)] + pub fn clear_storage(origin: OriginFor) -> DispatchResult { + let who = ensure_signed(origin)?; + ExampleStorage::::kill(); + Self::deposit_event(Event::StorageCleared { who }); + Ok(()) + } + } + + #[pallet::event] + #[pallet::generate_deposit(pub(super) fn deposit_event)] + pub enum Event { + /// Storage was accessed through unsafe getter + UnsafeGetStorage { who: T::AccountId, value: u32 }, + /// Storage value was set + StorageSet { who: T::AccountId, value: u32 }, + /// Storage was cleared + StorageCleared { who: T::AccountId }, + } + + #[pallet::error] + pub enum Error { + /// Value is not initialized + NotInitialized, + } + + #[pallet::genesis_config] + #[derive(frame_support::DefaultNoBound)] + pub struct GenesisConfig { + pub initial_value: Option, + #[serde(skip)] + pub _phantom: PhantomData, + } + + #[pallet::genesis_build] + impl BuildGenesisConfig for GenesisConfig { + fn build(&self) { + if let Some(value) = self.initial_value { + ExampleStorage::::put(value); + } + } + } +} diff --git a/test-cases/unsafe-unwrap/remediated/remediated-2/src/tests.rs b/test-cases/unsafe-unwrap/remediated/remediated-2/src/tests.rs new file mode 100644 index 00000000..302383da --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-2/src/tests.rs @@ -0,0 +1,106 @@ +use crate as pallet_example_basic; +use crate::*; +use frame_support::{assert_ok, derive_impl, dispatch::GetDispatchInfo}; +use sp_core::H256; +use sp_runtime::{ + traits::{BlakeTwo256, IdentityLookup}, + BuildStorage, +}; + +type Block = frame_system::mocking::MockBlock; + +frame_support::construct_runtime!( + pub enum Test + { + System: frame_system, + Balances: pallet_balances, + Example: pallet_example_basic, + } +); + +#[derive_impl(frame_system::config_preludes::TestDefaultConfig)] +impl frame_system::Config for Test { + type BaseCallFilter = frame_support::traits::Everything; + type BlockWeights = (); + type BlockLength = (); + type DbWeight = (); + type RuntimeOrigin = RuntimeOrigin; + type Nonce = u64; + type Hash = H256; + type RuntimeCall = RuntimeCall; + type Hashing = BlakeTwo256; + type AccountId = u64; + type Lookup = IdentityLookup; + type Block = Block; + type RuntimeEvent = RuntimeEvent; + type Version = (); + type PalletInfo = PalletInfo; + type AccountData = pallet_balances::AccountData; + type OnNewAccount = (); + type OnKilledAccount = (); + type SystemWeightInfo = (); + type SS58Prefix = (); + type OnSetCode = (); + type MaxConsumers = frame_support::traits::ConstU32<16>; +} + +#[derive_impl(pallet_balances::config_preludes::TestDefaultConfig)] +impl pallet_balances::Config for Test { + type AccountStore = System; +} + +impl Config for Test { + type RuntimeEvent = RuntimeEvent; + type WeightInfo = (); +} + +pub fn new_test_ext() -> sp_io::TestExternalities { + let t = RuntimeGenesisConfig { + system: Default::default(), + balances: Default::default(), + example: pallet_example_basic::GenesisConfig { + initial_value: Some(0), + _phantom: Default::default(), + }, + } + .build_storage() + .unwrap(); + t.into() +} + +#[test] +fn weights_work() { + let call = pallet_example_basic::Call::::unsafe_get_storage {}; + let info = call.get_dispatch_info(); + assert!(info.weight.ref_time() > 0); + assert_eq!( + info.weight, + ::WeightInfo::unsafe_get_storage() + ); +} + +#[test] +fn unsafe_get_storage_works() { + new_test_ext().execute_with(|| { + ExampleStorage::::put(42u32); + assert_ok!(Example::unsafe_get_storage(RuntimeOrigin::signed(1))); + }); +} + +#[test] +fn set_storage_works() { + new_test_ext().execute_with(|| { + let test_val = 133u32; + assert_ok!(Example::set_storage(RuntimeOrigin::signed(1), test_val)); + assert_eq!(ExampleStorage::::get(), Some(test_val)); + }); +} + +#[test] +fn clear_storage_works() { + new_test_ext().execute_with(|| { + ExampleStorage::::put(42u32); + assert_ok!(Example::clear_storage(RuntimeOrigin::signed(1))); + assert_eq!(ExampleStorage::::get(), None); + }); +} diff --git a/test-cases/unsafe-unwrap/remediated/remediated-2/src/weights.rs b/test-cases/unsafe-unwrap/remediated/remediated-2/src/weights.rs new file mode 100644 index 00000000..e9608480 --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-2/src/weights.rs @@ -0,0 +1,43 @@ +#![cfg_attr(rustfmt, rustfmt_skip)] +#![allow(unused_parens)] +#![allow(unused_imports)] + +use frame_support::{traits::Get, weights::{Weight, constants::RocksDbWeight}}; +use core::marker::PhantomData; + +pub trait WeightInfo { + fn unsafe_get_storage() -> Weight; + fn set_storage() -> Weight; + fn clear_storage() -> Weight; +} + +pub struct SubstrateWeight(PhantomData); +impl WeightInfo for SubstrateWeight { + fn unsafe_get_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads(1_u64)) + } + fn set_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads_writes(1_u64, 1_u64)) + } + fn clear_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads_writes(1_u64, 1_u64)) + } +} + +impl WeightInfo for () { + fn unsafe_get_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads(1_u64)) + } + fn set_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads_writes(1_u64, 1_u64)) + } + fn clear_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads_writes(1_u64, 1_u64)) + } +} diff --git a/test-cases/unsafe-unwrap/remediated/remediated-3/Cargo.toml b/test-cases/unsafe-unwrap/remediated/remediated-3/Cargo.toml new file mode 100644 index 00000000..85c9779b --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-3/Cargo.toml @@ -0,0 +1,41 @@ +[package] +name = "pallet-example-basic-remediated-3" +version = "27.0.0" +authors.workspace = true +edition.workspace = true +license = "MIT-0" +description = "FRAME example pallet" +readme = "README.md" +publish = false + +[package.metadata.docs.rs] +targets = ["x86_64-unknown-linux-gnu"] + +[dependencies] +codec = { workspace = true } +log = { workspace = true } +scale-info = { features = ["derive"], workspace = true } +frame-benchmarking = { optional = true, workspace = true } +frame-support = { workspace = true } +frame-system = { workspace = true } +pallet-balances = { workspace = true } +sp-io = { workspace = true } +sp-runtime = { workspace = true } + +[dev-dependencies] +sp-core = { workspace = true } + +[features] +default = ["std"] +std = [ + "codec/std", + "frame-benchmarking?/std", + "frame-support/std", + "frame-system/std", + "log/std", + "pallet-balances/std", + "scale-info/std", + "sp-core/std", + "sp-io/std", + "sp-runtime/std", +] diff --git a/test-cases/unsafe-unwrap/remediated/remediated-3/src/lib.rs b/test-cases/unsafe-unwrap/remediated/remediated-3/src/lib.rs new file mode 100644 index 00000000..ca1c89ac --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-3/src/lib.rs @@ -0,0 +1,109 @@ +// lib.rs +#![cfg_attr(not(feature = "std"), no_std)] + +extern crate alloc; + +#[cfg(test)] +mod tests; + +pub use pallet::*; + +pub mod weights; +pub use weights::*; + +#[frame_support::pallet] +pub mod pallet { + use super::*; + use frame_support::{pallet_prelude::*, traits::BuildGenesisConfig}; + use frame_system::pallet_prelude::*; + + #[pallet::config] + pub trait Config: pallet_balances::Config + frame_system::Config { + type RuntimeEvent: From> + IsType<::RuntimeEvent>; + /// Weight information for extrinsics in this pallet. + type WeightInfo: WeightInfo; + } + + #[pallet::pallet] + pub struct Pallet(_); + + #[pallet::storage] + #[pallet::getter(fn example_storage)] + pub type ExampleStorage = StorageValue<_, u32>; + + #[pallet::hooks] + impl Hooks> for Pallet { + // Add on_initialize if we need any per-block logic + fn on_initialize(_n: BlockNumberFor) -> Weight { + Weight::zero() + } + } + + #[pallet::call(weight(::WeightInfo))] + impl Pallet { + #[pallet::call_index(0)] + pub fn unsafe_get_storage(origin: OriginFor) -> DispatchResult { + let who = ensure_signed(origin)?; + let example_storage = ExampleStorage::::get(); + if example_storage.is_none() { + return Err(Error::::NotInitialized.into()); + } + let value = example_storage.unwrap(); + Self::deposit_event(Event::UnsafeGetStorage { who, value }); + Ok(()) + } + + #[pallet::call_index(1)] + pub fn set_storage(origin: OriginFor, new_value: u32) -> DispatchResult { + let who = ensure_signed(origin)?; + ExampleStorage::::put(new_value); + Self::deposit_event(Event::StorageSet { + who, + value: new_value, + }); + Ok(()) + } + + #[pallet::call_index(2)] + pub fn clear_storage(origin: OriginFor) -> DispatchResult { + let who = ensure_signed(origin)?; + ExampleStorage::::kill(); + Self::deposit_event(Event::StorageCleared { who }); + Ok(()) + } + } + + #[pallet::event] + #[pallet::generate_deposit(pub(super) fn deposit_event)] + pub enum Event { + /// Storage was accessed through unsafe getter + UnsafeGetStorage { who: T::AccountId, value: u32 }, + /// Storage value was set + StorageSet { who: T::AccountId, value: u32 }, + /// Storage was cleared + StorageCleared { who: T::AccountId }, + } + + #[pallet::error] + pub enum Error { + /// Value is not initialized + NotInitialized, + } + + #[pallet::genesis_config] + #[derive(frame_support::DefaultNoBound)] + pub struct GenesisConfig { + pub initial_value: Option, + #[serde(skip)] + pub _phantom: PhantomData, + } + + #[pallet::genesis_build] + impl BuildGenesisConfig for GenesisConfig { + fn build(&self) { + if let Some(value) = self.initial_value { + ExampleStorage::::put(value); + } + } + } +} diff --git a/test-cases/unsafe-unwrap/remediated/remediated-3/src/tests.rs b/test-cases/unsafe-unwrap/remediated/remediated-3/src/tests.rs new file mode 100644 index 00000000..302383da --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-3/src/tests.rs @@ -0,0 +1,106 @@ +use crate as pallet_example_basic; +use crate::*; +use frame_support::{assert_ok, derive_impl, dispatch::GetDispatchInfo}; +use sp_core::H256; +use sp_runtime::{ + traits::{BlakeTwo256, IdentityLookup}, + BuildStorage, +}; + +type Block = frame_system::mocking::MockBlock; + +frame_support::construct_runtime!( + pub enum Test + { + System: frame_system, + Balances: pallet_balances, + Example: pallet_example_basic, + } +); + +#[derive_impl(frame_system::config_preludes::TestDefaultConfig)] +impl frame_system::Config for Test { + type BaseCallFilter = frame_support::traits::Everything; + type BlockWeights = (); + type BlockLength = (); + type DbWeight = (); + type RuntimeOrigin = RuntimeOrigin; + type Nonce = u64; + type Hash = H256; + type RuntimeCall = RuntimeCall; + type Hashing = BlakeTwo256; + type AccountId = u64; + type Lookup = IdentityLookup; + type Block = Block; + type RuntimeEvent = RuntimeEvent; + type Version = (); + type PalletInfo = PalletInfo; + type AccountData = pallet_balances::AccountData; + type OnNewAccount = (); + type OnKilledAccount = (); + type SystemWeightInfo = (); + type SS58Prefix = (); + type OnSetCode = (); + type MaxConsumers = frame_support::traits::ConstU32<16>; +} + +#[derive_impl(pallet_balances::config_preludes::TestDefaultConfig)] +impl pallet_balances::Config for Test { + type AccountStore = System; +} + +impl Config for Test { + type RuntimeEvent = RuntimeEvent; + type WeightInfo = (); +} + +pub fn new_test_ext() -> sp_io::TestExternalities { + let t = RuntimeGenesisConfig { + system: Default::default(), + balances: Default::default(), + example: pallet_example_basic::GenesisConfig { + initial_value: Some(0), + _phantom: Default::default(), + }, + } + .build_storage() + .unwrap(); + t.into() +} + +#[test] +fn weights_work() { + let call = pallet_example_basic::Call::::unsafe_get_storage {}; + let info = call.get_dispatch_info(); + assert!(info.weight.ref_time() > 0); + assert_eq!( + info.weight, + ::WeightInfo::unsafe_get_storage() + ); +} + +#[test] +fn unsafe_get_storage_works() { + new_test_ext().execute_with(|| { + ExampleStorage::::put(42u32); + assert_ok!(Example::unsafe_get_storage(RuntimeOrigin::signed(1))); + }); +} + +#[test] +fn set_storage_works() { + new_test_ext().execute_with(|| { + let test_val = 133u32; + assert_ok!(Example::set_storage(RuntimeOrigin::signed(1), test_val)); + assert_eq!(ExampleStorage::::get(), Some(test_val)); + }); +} + +#[test] +fn clear_storage_works() { + new_test_ext().execute_with(|| { + ExampleStorage::::put(42u32); + assert_ok!(Example::clear_storage(RuntimeOrigin::signed(1))); + assert_eq!(ExampleStorage::::get(), None); + }); +} diff --git a/test-cases/unsafe-unwrap/remediated/remediated-3/src/weights.rs b/test-cases/unsafe-unwrap/remediated/remediated-3/src/weights.rs new file mode 100644 index 00000000..e9608480 --- /dev/null +++ b/test-cases/unsafe-unwrap/remediated/remediated-3/src/weights.rs @@ -0,0 +1,43 @@ +#![cfg_attr(rustfmt, rustfmt_skip)] +#![allow(unused_parens)] +#![allow(unused_imports)] + +use frame_support::{traits::Get, weights::{Weight, constants::RocksDbWeight}}; +use core::marker::PhantomData; + +pub trait WeightInfo { + fn unsafe_get_storage() -> Weight; + fn set_storage() -> Weight; + fn clear_storage() -> Weight; +} + +pub struct SubstrateWeight(PhantomData); +impl WeightInfo for SubstrateWeight { + fn unsafe_get_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads(1_u64)) + } + fn set_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads_writes(1_u64, 1_u64)) + } + fn clear_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads_writes(1_u64, 1_u64)) + } +} + +impl WeightInfo for () { + fn unsafe_get_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads(1_u64)) + } + fn set_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads_writes(1_u64, 1_u64)) + } + fn clear_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads_writes(1_u64, 1_u64)) + } +} diff --git a/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/Cargo.toml b/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/Cargo.toml new file mode 100644 index 00000000..7cfd1368 --- /dev/null +++ b/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/Cargo.toml @@ -0,0 +1,41 @@ +[package] +name = "pallet-example-basic-vulnerable-1" +version = "27.0.0" +authors.workspace = true +edition.workspace = true +license = "MIT-0" +description = "FRAME example pallet" +readme = "README.md" +publish = false + +[package.metadata.docs.rs] +targets = ["x86_64-unknown-linux-gnu"] + +[dependencies] +codec = { workspace = true } +log = { workspace = true } +scale-info = { features = ["derive"], workspace = true } +frame-benchmarking = { optional = true, workspace = true } +frame-support = { workspace = true } +frame-system = { workspace = true } +pallet-balances = { workspace = true } +sp-io = { workspace = true } +sp-runtime = { workspace = true } + +[dev-dependencies] +sp-core = { workspace = true } + +[features] +default = ["std"] +std = [ + "codec/std", + "frame-benchmarking?/std", + "frame-support/std", + "frame-system/std", + "log/std", + "pallet-balances/std", + "scale-info/std", + "sp-core/std", + "sp-io/std", + "sp-runtime/std", +] diff --git a/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/src/lib.rs b/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/src/lib.rs new file mode 100644 index 00000000..af30390e --- /dev/null +++ b/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/src/lib.rs @@ -0,0 +1,105 @@ +// lib.rs +#![cfg_attr(not(feature = "std"), no_std)] + +extern crate alloc; + +#[cfg(test)] +mod tests; + +pub use pallet::*; + +pub mod weights; +pub use weights::*; + +#[frame_support::pallet] +pub mod pallet { + use super::*; + use frame_support::{pallet_prelude::*, traits::BuildGenesisConfig}; + use frame_system::pallet_prelude::*; + + #[pallet::config] + pub trait Config: pallet_balances::Config + frame_system::Config { + type RuntimeEvent: From> + IsType<::RuntimeEvent>; + /// Weight information for extrinsics in this pallet. + type WeightInfo: WeightInfo; + } + + #[pallet::pallet] + pub struct Pallet(_); + + #[pallet::storage] + #[pallet::getter(fn example_storage)] + pub type ExampleStorage = StorageValue<_, u32>; + + #[pallet::hooks] + impl Hooks> for Pallet { + // Add on_initialize if we need any per-block logic + fn on_initialize(_n: BlockNumberFor) -> Weight { + Weight::zero() + } + } + + #[pallet::call(weight(::WeightInfo))] + impl Pallet { + #[pallet::call_index(0)] + pub fn unsafe_get_storage(origin: OriginFor) -> DispatchResult { + let who = ensure_signed(origin)?; + let value = ExampleStorage::::get().unwrap(); + Self::deposit_event(Event::UnsafeGetStorage { who, value }); + Ok(()) + } + + #[pallet::call_index(1)] + pub fn set_storage(origin: OriginFor, new_value: u32) -> DispatchResult { + let who = ensure_signed(origin)?; + ExampleStorage::::put(new_value); + Self::deposit_event(Event::StorageSet { + who, + value: new_value, + }); + Ok(()) + } + + #[pallet::call_index(2)] + pub fn clear_storage(origin: OriginFor) -> DispatchResult { + let who = ensure_signed(origin)?; + ExampleStorage::::kill(); + Self::deposit_event(Event::StorageCleared { who }); + Ok(()) + } + } + + #[pallet::event] + #[pallet::generate_deposit(pub(super) fn deposit_event)] + pub enum Event { + /// Storage was accessed through unsafe getter + UnsafeGetStorage { who: T::AccountId, value: u32 }, + /// Storage value was set + StorageSet { who: T::AccountId, value: u32 }, + /// Storage was cleared + StorageCleared { who: T::AccountId }, + } + + #[pallet::error] + pub enum Error { + /// Value is not initialized + NotInitialized, + } + + #[pallet::genesis_config] + #[derive(frame_support::DefaultNoBound)] + pub struct GenesisConfig { + pub initial_value: Option, + #[serde(skip)] + pub _phantom: PhantomData, + } + + #[pallet::genesis_build] + impl BuildGenesisConfig for GenesisConfig { + fn build(&self) { + if let Some(value) = self.initial_value { + ExampleStorage::::put(value); + } + } + } +} diff --git a/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/src/tests.rs b/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/src/tests.rs new file mode 100644 index 00000000..302383da --- /dev/null +++ b/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/src/tests.rs @@ -0,0 +1,106 @@ +use crate as pallet_example_basic; +use crate::*; +use frame_support::{assert_ok, derive_impl, dispatch::GetDispatchInfo}; +use sp_core::H256; +use sp_runtime::{ + traits::{BlakeTwo256, IdentityLookup}, + BuildStorage, +}; + +type Block = frame_system::mocking::MockBlock; + +frame_support::construct_runtime!( + pub enum Test + { + System: frame_system, + Balances: pallet_balances, + Example: pallet_example_basic, + } +); + +#[derive_impl(frame_system::config_preludes::TestDefaultConfig)] +impl frame_system::Config for Test { + type BaseCallFilter = frame_support::traits::Everything; + type BlockWeights = (); + type BlockLength = (); + type DbWeight = (); + type RuntimeOrigin = RuntimeOrigin; + type Nonce = u64; + type Hash = H256; + type RuntimeCall = RuntimeCall; + type Hashing = BlakeTwo256; + type AccountId = u64; + type Lookup = IdentityLookup; + type Block = Block; + type RuntimeEvent = RuntimeEvent; + type Version = (); + type PalletInfo = PalletInfo; + type AccountData = pallet_balances::AccountData; + type OnNewAccount = (); + type OnKilledAccount = (); + type SystemWeightInfo = (); + type SS58Prefix = (); + type OnSetCode = (); + type MaxConsumers = frame_support::traits::ConstU32<16>; +} + +#[derive_impl(pallet_balances::config_preludes::TestDefaultConfig)] +impl pallet_balances::Config for Test { + type AccountStore = System; +} + +impl Config for Test { + type RuntimeEvent = RuntimeEvent; + type WeightInfo = (); +} + +pub fn new_test_ext() -> sp_io::TestExternalities { + let t = RuntimeGenesisConfig { + system: Default::default(), + balances: Default::default(), + example: pallet_example_basic::GenesisConfig { + initial_value: Some(0), + _phantom: Default::default(), + }, + } + .build_storage() + .unwrap(); + t.into() +} + +#[test] +fn weights_work() { + let call = pallet_example_basic::Call::::unsafe_get_storage {}; + let info = call.get_dispatch_info(); + assert!(info.weight.ref_time() > 0); + assert_eq!( + info.weight, + ::WeightInfo::unsafe_get_storage() + ); +} + +#[test] +fn unsafe_get_storage_works() { + new_test_ext().execute_with(|| { + ExampleStorage::::put(42u32); + assert_ok!(Example::unsafe_get_storage(RuntimeOrigin::signed(1))); + }); +} + +#[test] +fn set_storage_works() { + new_test_ext().execute_with(|| { + let test_val = 133u32; + assert_ok!(Example::set_storage(RuntimeOrigin::signed(1), test_val)); + assert_eq!(ExampleStorage::::get(), Some(test_val)); + }); +} + +#[test] +fn clear_storage_works() { + new_test_ext().execute_with(|| { + ExampleStorage::::put(42u32); + assert_ok!(Example::clear_storage(RuntimeOrigin::signed(1))); + assert_eq!(ExampleStorage::::get(), None); + }); +} diff --git a/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/src/weights.rs b/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/src/weights.rs new file mode 100644 index 00000000..e9608480 --- /dev/null +++ b/test-cases/unsafe-unwrap/vulnerable/vulnerable-1/src/weights.rs @@ -0,0 +1,43 @@ +#![cfg_attr(rustfmt, rustfmt_skip)] +#![allow(unused_parens)] +#![allow(unused_imports)] + +use frame_support::{traits::Get, weights::{Weight, constants::RocksDbWeight}}; +use core::marker::PhantomData; + +pub trait WeightInfo { + fn unsafe_get_storage() -> Weight; + fn set_storage() -> Weight; + fn clear_storage() -> Weight; +} + +pub struct SubstrateWeight(PhantomData); +impl WeightInfo for SubstrateWeight { + fn unsafe_get_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads(1_u64)) + } + fn set_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads_writes(1_u64, 1_u64)) + } + fn clear_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(T::DbWeight::get().reads_writes(1_u64, 1_u64)) + } +} + +impl WeightInfo for () { + fn unsafe_get_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads(1_u64)) + } + fn set_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads_writes(1_u64, 1_u64)) + } + fn clear_storage() -> Weight { + Weight::from_parts(10_000_000_u64, 0) + .saturating_add(RocksDbWeight::get().reads_writes(1_u64, 1_u64)) + } +}