Skip to main content

randomx_rs/
lib.rs

1// Copyright 2019. The Tari Project
2//
3// Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
4// following conditions are met:
5//
6// 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
7// disclaimer.
8//
9// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the
10// following disclaimer in the documentation and/or other materials provided with the distribution.
11//
12// 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
13// products derived from this software without specific prior written permission.
14//
15// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
16// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
17// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
18// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
19// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
20// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
21// USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
22
23//! # RandomX
24//!
25//! The `randomx-rs` crate provides bindings to the RandomX proof-of-work (PoW) system.
26//!
27//! From the [RandomX github repo]:
28//!
29//! "RandomX is a proof-of-work (PoW) algorithm that is optimized for general-purpose CPUs. RandomX uses random code
30//! execution together with several memory-hard techniques to minimize the efficiency advantage of specialized
31//! hardware."
32//!
33//! Read more about how RandomX works in the [design document].
34//!
35//! [RandomX github repo]: <https://github.com/tevador/RandomX>
36//! [design document]: <https://github.com/tevador/RandomX/blob/master/doc/design.md>
37mod bindings;
38/// Test utilities for fuzzing
39pub mod test_utils;
40
41use std::{convert::TryFrom, num::TryFromIntError, ptr, sync::Arc};
42
43use bindings::{
44    randomx_alloc_cache,
45    randomx_alloc_dataset,
46    randomx_cache,
47    randomx_calculate_hash,
48    randomx_create_vm,
49    randomx_dataset,
50    randomx_dataset_item_count,
51    randomx_destroy_vm,
52    randomx_get_dataset_memory,
53    randomx_init_cache,
54    randomx_init_dataset,
55    randomx_release_cache,
56    randomx_release_dataset,
57    randomx_vm,
58    randomx_vm_set_cache,
59    randomx_vm_set_dataset,
60    RANDOMX_HASH_SIZE,
61};
62use bitflags::bitflags;
63use libc::{c_ulong, c_void};
64use thiserror::Error;
65
66/// The size, in bytes, of a single RandomX dataset item. The full dataset returned by
67/// [`RandomXDataset::get_data`] is `RandomXDataset::count()` items of this size.
68pub use crate::bindings::RANDOMX_DATASET_ITEM_SIZE;
69use crate::bindings::{
70    randomx_calculate_hash_first,
71    randomx_calculate_hash_last,
72    randomx_calculate_hash_next,
73    randomx_get_flags,
74};
75
76bitflags! {
77    /// RandomX Flags are used to configure the library.
78    pub struct RandomXFlag: u32 {
79        /// No flags set. Works on all platforms, but is the slowest.
80        const FLAG_DEFAULT      = 0b0000_0000;
81        /// Allocate memory in large pages.
82        const FLAG_LARGE_PAGES  = 0b0000_0001;
83        /// Use hardware accelerated AES.
84        const FLAG_HARD_AES     = 0b0000_0010;
85        /// Use the full dataset.
86        const FLAG_FULL_MEM     = 0b0000_0100;
87        /// Use JIT compilation support.
88        const FLAG_JIT          = 0b0000_1000;
89        /// When combined with FLAG_JIT, the JIT pages are never writable and executable at the
90        /// same time.
91        const FLAG_SECURE       = 0b0001_0000;
92        /// Optimize Argon2 for CPUs with the SSSE3 instruction set.
93        const FLAG_ARGON2_SSSE3 = 0b0010_0000;
94        /// Optimize Argon2 for CPUs with the AVX2 instruction set.
95        const FLAG_ARGON2_AVX2  = 0b0100_0000;
96        /// Optimize Argon2 for CPUs without the AVX2 or SSSE3 instruction sets.
97        const FLAG_ARGON2       = 0b0110_0000;
98    }
99}
100
101impl RandomXFlag {
102    /// Returns the recommended flags to be used.
103    ///
104    /// Does not include:
105    /// * FLAG_LARGE_PAGES
106    /// * FLAG_FULL_MEM
107    /// * FLAG_SECURE
108    ///
109    /// The above flags need to be set manually, if required.
110    pub fn get_recommended_flags() -> RandomXFlag {
111        RandomXFlag {
112            bits: unsafe { randomx_get_flags() },
113        }
114    }
115}
116
117impl Default for RandomXFlag {
118    /// Default value for RandomXFlag
119    fn default() -> RandomXFlag {
120        RandomXFlag::FLAG_DEFAULT
121    }
122}
123
124#[derive(Debug, Clone, Error)]
125/// This enum specifies the possible errors that may occur.
126pub enum RandomXError {
127    #[error("Problem creating the RandomX object: {0}")]
128    CreationError(String),
129    #[error("Problem with configuration flags: {0}")]
130    FlagConfigError(String),
131    #[error("Problem with parameters supplied: {0}")]
132    ParameterError(String),
133    #[error("Failed to convert Int to usize")]
134    TryFromIntError(#[from] TryFromIntError),
135    #[error("Unknown problem running RandomX: {0}")]
136    Other(String),
137}
138
139#[derive(Debug)]
140struct RandomXCacheInner {
141    cache_ptr: *mut randomx_cache,
142}
143
144impl Drop for RandomXCacheInner {
145    /// De-allocates memory for the `cache` object
146    fn drop(&mut self) {
147        unsafe {
148            randomx_release_cache(self.cache_ptr);
149        }
150    }
151}
152
153#[derive(Debug, Clone)]
154/// The Cache is used for light verification and Dataset construction.
155pub struct RandomXCache {
156    inner: Arc<RandomXCacheInner>,
157}
158
159impl RandomXCache {
160    /// Creates and alllcates memory for a new cache object, and initializes it with
161    /// the key value.
162    ///
163    /// `flags` is any combination of the following two flags:
164    /// * FLAG_LARGE_PAGES
165    /// * FLAG_JIT
166    ///
167    /// and (optionally) one of the following flags (depending on instruction set supported):
168    /// * FLAG_ARGON2_SSSE3
169    /// * FLAG_ARGON2_AVX2
170    ///
171    /// `key` is a sequence of u8 used to initialize SuperScalarHash.
172    pub fn new(flags: RandomXFlag, key: &[u8]) -> Result<RandomXCache, RandomXError> {
173        if key.is_empty() {
174            Err(RandomXError::ParameterError("key is empty".to_string()))
175        } else {
176            let cache_ptr = unsafe { randomx_alloc_cache(flags.bits) };
177            if cache_ptr.is_null() {
178                Err(RandomXError::CreationError("Could not allocate cache".to_string()))
179            } else {
180                let inner = RandomXCacheInner { cache_ptr };
181                let result = RandomXCache { inner: Arc::new(inner) };
182                let key_ptr = key.as_ptr() as *mut c_void;
183                let key_size = key.len();
184                unsafe {
185                    randomx_init_cache(result.inner.cache_ptr, key_ptr, key_size);
186                }
187                Ok(result)
188            }
189        }
190    }
191}
192
193#[derive(Debug)]
194struct RandomXDatasetInner {
195    dataset_ptr: *mut randomx_dataset,
196    dataset_count: u32,
197    #[allow(dead_code)]
198    cache: RandomXCache,
199}
200
201impl Drop for RandomXDatasetInner {
202    /// De-allocates memory for the `dataset` object.
203    fn drop(&mut self) {
204        unsafe {
205            randomx_release_dataset(self.dataset_ptr);
206        }
207    }
208}
209
210#[derive(Debug, Clone)]
211/// The Dataset is a read-only memory structure that is used during VM program execution.
212pub struct RandomXDataset {
213    inner: Arc<RandomXDatasetInner>,
214}
215
216impl RandomXDataset {
217    /// Creates a new dataset object, allocates memory to the `dataset` object and initializes it.
218    ///
219    /// `flags` is one of the following:
220    /// * FLAG_DEFAULT
221    /// * FLAG_LARGE_PAGES
222    ///
223    /// `cache` is a cache object.
224    ///
225    /// `start` is the item number where initialization should start. **Pass 0.** The items
226    /// `[start, RandomXDataset::count())` are initialized by the RandomX library; the leading `start` items are
227    /// zeroed, since the library never writes them.
228    ///
229    /// # Warning
230    ///
231    /// The RandomX API requires that *every* item from `0` to `RandomXDataset::count() - 1` is initialized before a
232    /// dataset may be used (see the note on `randomx_init_dataset` in `randomx.h`). A non-zero `start` therefore
233    /// produces a dataset that does **not** satisfy that precondition and must **not** be passed to
234    /// [`RandomXVM::new`] or [`RandomXVM::reinit_dataset`]. Doing so is not detected or reported: the VM will read
235    /// the zeroed leading items as though they were real dataset items and silently compute hashes that disagree
236    /// with every other RandomX implementation.
237    ///
238    /// The only legitimate use of a non-zero `start` in the upstream API is splitting the initialization of a single
239    /// *shared* dataset across several threads, each initializing a different item range. This wrapper cannot express
240    /// that, because `new` always allocates its own dataset, so there is no correct value other than 0.
241    // Conversions may be lossy on Windows or Linux
242    #[allow(clippy::useless_conversion)]
243    pub fn new(flags: RandomXFlag, cache: RandomXCache, start: u32) -> Result<RandomXDataset, RandomXError> {
244        let item_count = RandomXDataset::count()
245            .map_err(|e| RandomXError::CreationError(format!("Could not get dataset count: {e:?}")))?;
246
247        let test = unsafe { randomx_alloc_dataset(flags.bits) };
248        if test.is_null() {
249            Err(RandomXError::CreationError("Could not allocate dataset".to_string()))
250        } else {
251            let inner = RandomXDatasetInner {
252                dataset_ptr: test,
253                dataset_count: item_count,
254                cache,
255            };
256            let result = RandomXDataset { inner: Arc::new(inner) };
257
258            if start < item_count {
259                // `randomx_init_dataset` initialises the items `[start, start + count)`, so the count passed to it
260                // must be the number of *remaining* items. Passing the full `item_count` with a non-zero `start`
261                // writes `start` items past the end of the allocation (a heap buffer overflow, silent in release
262                // builds because the library's assertions are compiled out by `NDEBUG`).
263                let remaining = item_count.saturating_sub(start);
264                // `randomx_alloc_dataset` hands back uninitialised memory (only `FLAG_LARGE_PAGES` gets zero pages),
265                // and the call below only writes from `start` onwards, so zero the leading `start` items. Without
266                // this the first `start * RANDOMX_DATASET_ITEM_SIZE` bytes stay uninitialised and reading them in
267                // `get_data` would be undefined behaviour as well as an information leak.
268                if start > 0 {
269                    let memory = unsafe { randomx_get_dataset_memory(result.inner.dataset_ptr) };
270                    if memory.is_null() {
271                        return Err(RandomXError::CreationError(
272                            "Could not get dataset memory to zero the uninitialised prefix".to_string(),
273                        ));
274                    }
275                    let prefix_len = usize::try_from(start)?
276                        .checked_mul(RANDOMX_DATASET_ITEM_SIZE)
277                        .ok_or_else(|| {
278                            RandomXError::CreationError(format!("Dataset prefix size overflows: {start}"))
279                        })?;
280                    // SAFETY: `memory` is the non-null start of the dataset buffer, which the library allocated with
281                    // room for `item_count` items of `RANDOMX_DATASET_ITEM_SIZE` bytes. `start < item_count`, so
282                    // `prefix_len` bytes lie inside that allocation. `u8` is always valid for any bit pattern and has
283                    // an alignment of 1, and nothing else refers to the buffer yet.
284                    unsafe {
285                        ptr::write_bytes(memory.cast::<u8>(), 0, prefix_len);
286                    }
287                }
288                unsafe {
289                    randomx_init_dataset(
290                        result.inner.dataset_ptr,
291                        result.inner.cache.inner.cache_ptr,
292                        c_ulong::from(start),
293                        c_ulong::from(remaining),
294                    );
295                }
296                Ok(result)
297            } else {
298                Err(RandomXError::CreationError(format!(
299                    "start must be less than item_count: start: {start}, item_count: {item_count}",
300                )))
301            }
302        }
303    }
304
305    /// Returns the number of items in the `dataset` or an error on failure.
306    pub fn count() -> Result<u32, RandomXError> {
307        match unsafe { randomx_dataset_item_count() } {
308            0 => Err(RandomXError::Other("Dataset item count was 0".to_string())),
309            x => {
310                // This weirdness brought to you by c_ulong being different on Windows and Linux
311                #[cfg(target_os = "windows")]
312                return Ok(x);
313                #[cfg(not(target_os = "windows"))]
314                return Ok(u32::try_from(x)?);
315            },
316        }
317    }
318
319    /// Returns a copy of the *entire* internal memory buffer of the `dataset`, or an error on failure.
320    ///
321    /// The returned buffer is `RandomXDataset::count()` items of [`RANDOMX_DATASET_ITEM_SIZE`] (64) bytes each, i.e.
322    /// approximately 2.03 GB with the default RandomX configuration. This is an expensive, fully allocating copy of
323    /// the dataset, so avoid calling it on a hot path. The allocation is fallible: an out-of-memory condition is
324    /// reported as a [`RandomXError`] instead of aborting the process.
325    ///
326    /// If the dataset was created with a non-zero `start`, the first `start` items were never initialised by the
327    /// RandomX library; [`RandomXDataset::new`] zeroes them, so they are returned here as zero bytes.
328    pub fn get_data(&self) -> Result<Vec<u8>, RandomXError> {
329        let memory = unsafe { randomx_get_dataset_memory(self.inner.dataset_ptr) };
330        if memory.is_null() {
331            return Err(RandomXError::Other("Could not get dataset memory".into()));
332        }
333        // `dataset_count` is an *item* count, not a byte count; each item is `RANDOMX_DATASET_ITEM_SIZE` bytes.
334        let item_count = usize::try_from(self.inner.dataset_count)?;
335        let size_in_bytes = item_count.checked_mul(RANDOMX_DATASET_ITEM_SIZE).ok_or_else(|| {
336            RandomXError::Other(format!(
337                "Dataset size overflows usize: {item_count} items of {RANDOMX_DATASET_ITEM_SIZE} bytes each",
338            ))
339        })?;
340        // SAFETY: `memory` is a non-null pointer to the dataset buffer owned by the RandomX library. The library
341        // allocated that buffer with room for `randomx_dataset_item_count()` items of `RANDOMX_DATASET_ITEM_SIZE`
342        // bytes each, and `dataset_count` was set from that same call, so exactly `size_in_bytes` bytes are inside
343        // the allocation. Every one of those bytes is initialised: `RandomXDataset::new` has the library write the
344        // items `[start, count)` and zeroes the `[0, start)` prefix that the library leaves untouched. `u8` has an
345        // alignment of 1, so the pointer is trivially aligned. The buffer outlives the slice: `&self` keeps the
346        // `Arc<RandomXDatasetInner>` (and hence the dataset allocation) alive, the dataset is read-only once
347        // initialised, and the slice is copied into an owned `Vec` before this function returns.
348        let data = unsafe { std::slice::from_raw_parts(memory.cast::<u8>(), size_in_bytes) };
349        // Allocate fallibly: a plain `to_vec` of ~2 GB would call `handle_alloc_error` and abort the whole process
350        // on failure, which is not an acceptable outcome for a library that returns a `Result`.
351        let mut result = Vec::new();
352        result.try_reserve_exact(size_in_bytes).map_err(|e| {
353            RandomXError::Other(format!(
354                "Could not allocate {size_in_bytes} bytes for the dataset copy: {e}"
355            ))
356        })?;
357        result.extend_from_slice(data);
358        Ok(result)
359    }
360}
361
362#[derive(Debug)]
363/// The RandomX Virtual Machine (VM) is a complex instruction set computer that executes generated programs.
364pub struct RandomXVM {
365    flags: RandomXFlag,
366    vm: *mut randomx_vm,
367    linked_cache: Option<RandomXCache>,
368    linked_dataset: Option<RandomXDataset>,
369}
370
371impl Drop for RandomXVM {
372    /// De-allocates memory for the `VM` object.
373    fn drop(&mut self) {
374        unsafe {
375            randomx_destroy_vm(self.vm);
376        }
377    }
378}
379
380impl RandomXVM {
381    /// Creates a new `VM` and initializes it, error on failure.
382    ///
383    /// `flags` is any combination of the following 5 flags:
384    /// * FLAG_LARGE_PAGES
385    /// * FLAG_HARD_AES
386    /// * FLAG_FULL_MEM
387    /// * FLAG_JIT
388    /// * FLAG_SECURE
389    ///
390    /// Or
391    ///
392    /// * FLAG_DEFAULT
393    ///
394    /// `cache` is a cache object, optional if FLAG_FULL_MEM is set.
395    ///
396    /// `dataset` is a dataset object, optional if FLAG_FULL_MEM is not set.
397    pub fn new(
398        flags: RandomXFlag,
399        cache: Option<RandomXCache>,
400        dataset: Option<RandomXDataset>,
401    ) -> Result<RandomXVM, RandomXError> {
402        let is_full_mem = flags.contains(RandomXFlag::FLAG_FULL_MEM);
403        match (cache, dataset) {
404            (None, None) => Err(RandomXError::CreationError("Failed to allocate VM".to_string())),
405            (None, _) if !is_full_mem => Err(RandomXError::FlagConfigError(
406                "No cache and FLAG_FULL_MEM not set".to_string(),
407            )),
408            (_, None) if is_full_mem => Err(RandomXError::FlagConfigError(
409                "No dataset and FLAG_FULL_MEM set".to_string(),
410            )),
411            (cache, dataset) => {
412                let cache_ptr = cache
413                    .as_ref()
414                    .map(|stash| stash.inner.cache_ptr)
415                    .unwrap_or_else(ptr::null_mut);
416                let dataset_ptr = dataset
417                    .as_ref()
418                    .map(|data| data.inner.dataset_ptr)
419                    .unwrap_or_else(ptr::null_mut);
420                let vm = unsafe { randomx_create_vm(flags.bits, cache_ptr, dataset_ptr) };
421                Ok(RandomXVM {
422                    vm,
423                    flags,
424                    linked_cache: cache,
425                    linked_dataset: dataset,
426                })
427            },
428        }
429    }
430
431    /// Re-initializes the `VM` with a new cache that was initialised without
432    /// RandomXFlag::FLAG_FULL_MEM.
433    pub fn reinit_cache(&mut self, cache: RandomXCache) -> Result<(), RandomXError> {
434        if self.flags.contains(RandomXFlag::FLAG_FULL_MEM) {
435            Err(RandomXError::FlagConfigError(
436                "Cannot reinit cache with FLAG_FULL_MEM set".to_string(),
437            ))
438        } else {
439            unsafe {
440                randomx_vm_set_cache(self.vm, cache.inner.cache_ptr);
441            }
442            self.linked_cache = Some(cache);
443            Ok(())
444        }
445    }
446
447    /// Re-initializes the `VM` with a new dataset that was initialised with
448    /// RandomXFlag::FLAG_FULL_MEM.
449    pub fn reinit_dataset(&mut self, dataset: RandomXDataset) -> Result<(), RandomXError> {
450        if self.flags.contains(RandomXFlag::FLAG_FULL_MEM) {
451            unsafe {
452                randomx_vm_set_dataset(self.vm, dataset.inner.dataset_ptr);
453            }
454            self.linked_dataset = Some(dataset);
455            Ok(())
456        } else {
457            Err(RandomXError::FlagConfigError(
458                "Cannot reinit dataset without FLAG_FULL_MEM set".to_string(),
459            ))
460        }
461    }
462
463    /// Calculates a RandomX hash value and returns it, error on failure.
464    ///
465    /// `input` is a sequence of u8 to be hashed.
466    pub fn calculate_hash(&self, input: &[u8]) -> Result<Vec<u8>, RandomXError> {
467        if input.is_empty() {
468            Err(RandomXError::ParameterError("input was empty".to_string()))
469        } else {
470            let size_input = input.len();
471            let input_ptr = input.as_ptr() as *mut c_void;
472            let arr = [0; RANDOMX_HASH_SIZE as usize];
473            let output_ptr = arr.as_ptr() as *mut c_void;
474            unsafe {
475                randomx_calculate_hash(self.vm, input_ptr, size_input, output_ptr);
476            }
477            // if this failed, arr should still be empty
478            if arr == [0; RANDOMX_HASH_SIZE as usize] {
479                Err(RandomXError::Other("RandomX calculated hash was empty".to_string()))
480            } else {
481                let result = arr.to_vec();
482                Ok(result)
483            }
484        }
485    }
486
487    /// Calculates hashes from a set of inputs.
488    ///
489    /// `input` is an array of a sequence of u8 to be hashed.
490    #[allow(clippy::needless_range_loop)] // Range loop is not only for indexing `input`
491    pub fn calculate_hash_set(&self, input: &[&[u8]]) -> Result<Vec<Vec<u8>>, RandomXError> {
492        if input.is_empty() {
493            // Empty set
494            return Err(RandomXError::ParameterError("input was empty".to_string()));
495        }
496
497        let mut result = Vec::new();
498        // For single input
499        if input.len() == 1 {
500            let hash = self.calculate_hash(input[0])?;
501            result.push(hash);
502            return Ok(result);
503        }
504
505        // For multiple inputs
506        let mut output_ptr: *mut c_void = ptr::null_mut();
507        let arr = [0; RANDOMX_HASH_SIZE as usize];
508
509        // Not len() as last iteration assigns final hash
510        let iterations = input.len() + 1;
511        for i in 0..iterations {
512            if i == iterations - 1 {
513                // For last iteration
514                unsafe {
515                    randomx_calculate_hash_last(self.vm, output_ptr);
516                }
517            } else {
518                if input[i].is_empty() {
519                    // Stop calculations
520                    if arr != [0; RANDOMX_HASH_SIZE as usize] {
521                        // Complete what was started
522                        unsafe {
523                            randomx_calculate_hash_last(self.vm, output_ptr);
524                        }
525                    }
526                    return Err(RandomXError::ParameterError("input was empty".to_string()));
527                };
528                let size_input = input[i].len();
529                let input_ptr = input[i].as_ptr() as *mut c_void;
530                output_ptr = arr.as_ptr() as *mut c_void;
531                if i == 0 {
532                    // For first iteration
533                    unsafe {
534                        randomx_calculate_hash_first(self.vm, input_ptr, size_input);
535                    }
536                } else {
537                    unsafe {
538                        // For every other iteration
539                        randomx_calculate_hash_next(self.vm, input_ptr, size_input, output_ptr);
540                    }
541                }
542            }
543
544            if i != 0 {
545                // First hash is only available in 2nd iteration
546                if arr == [0; RANDOMX_HASH_SIZE as usize] {
547                    return Err(RandomXError::Other("RandomX hash was zero".to_string()));
548                }
549                let output: Vec<u8> = arr.to_vec();
550                result.push(output);
551            }
552        }
553        Ok(result)
554    }
555}
556
557#[cfg(test)]
558mod tests {
559    use std::{convert::TryFrom, ptr, sync::Arc};
560
561    use crate::{
562        bindings::randomx_get_dataset_memory,
563        RandomXCache,
564        RandomXCacheInner,
565        RandomXDataset,
566        RandomXDatasetInner,
567        RandomXFlag,
568        RandomXVM,
569        RANDOMX_DATASET_ITEM_SIZE,
570    };
571
572    #[test]
573    fn lib_alloc_cache() {
574        let flags = RandomXFlag::default();
575        let key = "Key";
576        let cache = RandomXCache::new(flags, key.as_bytes()).expect("Failed to allocate cache");
577        drop(cache);
578    }
579
580    #[test]
581    fn lib_alloc_dataset() {
582        let flags = RandomXFlag::default();
583        let key = "Key";
584        let cache = RandomXCache::new(flags, key.as_bytes()).unwrap();
585        let dataset = RandomXDataset::new(flags, cache.clone(), 0).expect("Failed to allocate dataset");
586        drop(dataset);
587        drop(cache);
588    }
589
590    #[test]
591    fn lib_alloc_vm() {
592        let flags = RandomXFlag::default();
593        let key = "Key";
594        let cache = RandomXCache::new(flags, key.as_bytes()).unwrap();
595        let mut vm = RandomXVM::new(flags, Some(cache.clone()), None).expect("Failed to allocate VM");
596        drop(vm);
597        let dataset = RandomXDataset::new(flags, cache.clone(), 0).unwrap();
598        vm = RandomXVM::new(flags, Some(cache.clone()), Some(dataset.clone())).expect("Failed to allocate VM");
599        drop(dataset);
600        drop(cache);
601        drop(vm);
602    }
603
604    #[test]
605    fn lib_dataset_memory() {
606        let flags = RandomXFlag::default();
607        let key = "Key";
608        let cache = RandomXCache::new(flags, key.as_bytes()).unwrap();
609        let dataset = RandomXDataset::new(flags, cache.clone(), 0).unwrap();
610        let item_count = usize::try_from(RandomXDataset::count().unwrap()).unwrap();
611        let memory = dataset.get_data().expect("Failed to get dataset memory");
612        // `get_data` must return the *whole* dataset: one item is `RANDOMX_DATASET_ITEM_SIZE` bytes.
613        assert_eq!(
614            memory.len(),
615            item_count * RANDOMX_DATASET_ITEM_SIZE,
616            "get_data did not return the full dataset"
617        );
618        // Check the *last* item, which is the part of the range the length fix actually extended, rather than
619        // scanning from the front (which would short-circuit on the very first byte and prove nothing).
620        assert!(
621            memory[memory.len() - RANDOMX_DATASET_ITEM_SIZE..]
622                .iter()
623                .any(|&b| b != 0),
624            "The last dataset item was all zeroes"
625        );
626        drop(memory);
627        drop(dataset);
628        drop(cache);
629    }
630
631    #[test]
632    fn lib_dataset_non_zero_start() {
633        const START: u32 = 2;
634        let flags = RandomXFlag::default();
635        let key = "Key";
636        let cache = RandomXCache::new(flags, key.as_bytes()).unwrap();
637        let item_count = RandomXDataset::count().unwrap();
638        let dataset = RandomXDataset::new(flags, cache.clone(), START).expect("Failed to allocate dataset");
639
640        // Read the library's buffer directly rather than through `get_data`, so this test does not allocate a
641        // second ~2 GB copy on top of the dataset itself.
642        let memory = unsafe { randomx_get_dataset_memory(dataset.inner.dataset_ptr) };
643        assert!(!memory.is_null());
644        let len = usize::try_from(item_count).unwrap() * RANDOMX_DATASET_ITEM_SIZE;
645        // SAFETY: `memory` is the non-null dataset buffer of `item_count` items, fully initialised by
646        // `RandomXDataset::new`, and it is kept alive by `dataset` for the duration of the borrow.
647        let data = unsafe { std::slice::from_raw_parts(memory.cast::<u8>(), len) };
648
649        // The library never writes the first `START` items, so `new` must have zeroed them.
650        let prefix = usize::try_from(START).unwrap() * RANDOMX_DATASET_ITEM_SIZE;
651        assert!(
652            data[..prefix].iter().all(|&b| b == 0),
653            "The uninitialised prefix was not zeroed"
654        );
655        // The final item must be initialised, i.e. the range is not *under*-initialised. Note that this assertion
656        // does not by itself catch a reintroduced overflow: the last item is written both when the correct
657        // remaining count is passed and when the full item count is passed. What catches that regression is the
658        // RandomX library's own `assert(startItem + itemCount <= DatasetItemCount)` in `randomx.cpp`, which is live
659        // in debug builds (the `cmake` crate maps a debug Rust profile to `CMAKE_BUILD_TYPE=Debug`) and aborts the
660        // test binary. CI runs the suite in debug, so a regression here fails the build.
661        assert!(
662            data[len - RANDOMX_DATASET_ITEM_SIZE..].iter().any(|&b| b != 0),
663            "The last dataset item was not initialised"
664        );
665
666        // `start` must stay inside the dataset.
667        assert!(RandomXDataset::new(flags, cache.clone(), item_count).is_err());
668
669        drop(dataset);
670        drop(cache);
671    }
672
673    #[test]
674    fn test_null_assignments() {
675        let flags = RandomXFlag::get_recommended_flags();
676        if let Ok(mut vm) = RandomXVM::new(flags, None, None) {
677            let cache = RandomXCache {
678                inner: Arc::new(RandomXCacheInner {
679                    cache_ptr: ptr::null_mut(),
680                }),
681            };
682            assert!(vm.reinit_cache(cache.clone()).is_err());
683            let dataset = RandomXDataset {
684                inner: Arc::new(RandomXDatasetInner {
685                    dataset_ptr: ptr::null_mut(),
686                    dataset_count: 0,
687                    cache,
688                }),
689            };
690            assert!(vm.reinit_dataset(dataset.clone()).is_err());
691        }
692    }
693
694    #[test]
695    fn lib_calculate_hash() {
696        let flags = RandomXFlag::get_recommended_flags();
697        let flags2 = flags | RandomXFlag::FLAG_FULL_MEM;
698        let key = "Key";
699        let input = "Input";
700        let cache1 = RandomXCache::new(flags, key.as_bytes()).unwrap();
701        let mut vm1 = RandomXVM::new(flags, Some(cache1.clone()), None).unwrap();
702        let hash1 = vm1.calculate_hash(input.as_bytes()).expect("no data");
703        let vec = vec![0u8; hash1.len()];
704        assert_ne!(hash1, vec);
705        let reinit_cache = vm1.reinit_cache(cache1.clone());
706        assert!(reinit_cache.is_ok());
707        let hash2 = vm1.calculate_hash(input.as_bytes()).expect("no data");
708        assert_ne!(hash2, vec);
709        assert_eq!(hash1, hash2);
710
711        let cache2 = RandomXCache::new(flags, key.as_bytes()).unwrap();
712        let vm2 = RandomXVM::new(flags, Some(cache2.clone()), None).unwrap();
713        let hash3 = vm2.calculate_hash(input.as_bytes()).expect("no data");
714        assert_eq!(hash2, hash3);
715
716        let cache3 = RandomXCache::new(flags, key.as_bytes()).unwrap();
717        let dataset3 = RandomXDataset::new(flags, cache3.clone(), 0).unwrap();
718        let mut vm3 = RandomXVM::new(flags2, None, Some(dataset3.clone())).unwrap();
719        let hash4 = vm3.calculate_hash(input.as_bytes()).expect("no data");
720        assert_ne!(hash3, vec);
721        let reinit_dataset = vm3.reinit_dataset(dataset3.clone());
722        assert!(reinit_dataset.is_ok());
723        let hash5 = vm3.calculate_hash(input.as_bytes()).expect("no data");
724        assert_ne!(hash4, vec);
725        assert_eq!(hash4, hash5);
726
727        let cache4 = RandomXCache::new(flags, key.as_bytes()).unwrap();
728        let dataset4 = RandomXDataset::new(flags, cache4.clone(), 0).unwrap();
729        let vm4 = RandomXVM::new(flags2, Some(cache4), Some(dataset4.clone())).unwrap();
730        let hash6 = vm3.calculate_hash(input.as_bytes()).expect("no data");
731        assert_eq!(hash5, hash6);
732
733        drop(dataset3);
734        drop(dataset4);
735        drop(cache1);
736        drop(cache2);
737        drop(cache3);
738        drop(vm1);
739        drop(vm2);
740        drop(vm3);
741        drop(vm4);
742    }
743
744    #[test]
745    fn lib_calculate_hash_set() {
746        let flags = RandomXFlag::default();
747        let key = "Key";
748        let inputs = vec!["Input".as_bytes(), "Input 2".as_bytes(), "Inputs 3".as_bytes()];
749        let cache = RandomXCache::new(flags, key.as_bytes()).unwrap();
750        let vm = RandomXVM::new(flags, Some(cache.clone()), None).unwrap();
751        let hashes = vm.calculate_hash_set(inputs.as_slice()).expect("no data");
752        assert_eq!(inputs.len(), hashes.len());
753        let mut prev_hash = Vec::new();
754        for (i, hash) in hashes.into_iter().enumerate() {
755            let vec = vec![0u8; hash.len()];
756            assert_ne!(hash, vec);
757            assert_ne!(hash, prev_hash);
758            let compare = vm.calculate_hash(inputs[i]).unwrap(); // sanity check
759            assert_eq!(hash, compare);
760            prev_hash = hash;
761        }
762        drop(cache);
763        drop(vm);
764    }
765
766    #[test]
767    fn lib_calculate_hash_is_consistent() {
768        let flags = RandomXFlag::get_recommended_flags();
769        let key = "Key";
770        let input = "Input";
771        let cache = RandomXCache::new(flags, key.as_bytes()).unwrap();
772        let dataset = RandomXDataset::new(flags, cache.clone(), 0).unwrap();
773        let vm = RandomXVM::new(flags, Some(cache.clone()), Some(dataset.clone())).unwrap();
774        let hash = vm.calculate_hash(input.as_bytes()).expect("no data");
775        assert_eq!(hash, [
776            114, 81, 192, 5, 165, 242, 107, 100, 184, 77, 37, 129, 52, 203, 217, 227, 65, 83, 215, 213, 59, 71, 32,
777            172, 253, 155, 204, 111, 183, 213, 157, 155
778        ]);
779        drop(vm);
780        drop(dataset);
781        drop(cache);
782
783        let cache1 = RandomXCache::new(flags, key.as_bytes()).unwrap();
784        let dataset1 = RandomXDataset::new(flags, cache1.clone(), 0).unwrap();
785        let vm1 = RandomXVM::new(flags, Some(cache1.clone()), Some(dataset1.clone())).unwrap();
786        let hash1 = vm1.calculate_hash(input.as_bytes()).expect("no data");
787        assert_eq!(hash1, [
788            114, 81, 192, 5, 165, 242, 107, 100, 184, 77, 37, 129, 52, 203, 217, 227, 65, 83, 215, 213, 59, 71, 32,
789            172, 253, 155, 204, 111, 183, 213, 157, 155
790        ]);
791        drop(vm1);
792        drop(dataset1);
793        drop(cache1);
794    }
795
796    #[test]
797    fn lib_check_cache_and_dataset_lifetimes() {
798        let flags = RandomXFlag::get_recommended_flags();
799        let key = "Key";
800        let input = "Input";
801        let cache = RandomXCache::new(flags, key.as_bytes()).unwrap();
802        let dataset = RandomXDataset::new(flags, cache.clone(), 0).unwrap();
803        let vm = RandomXVM::new(flags, Some(cache.clone()), Some(dataset.clone())).unwrap();
804        drop(dataset);
805        drop(cache);
806        let hash = vm.calculate_hash(input.as_bytes()).expect("no data");
807        assert_eq!(hash, [
808            114, 81, 192, 5, 165, 242, 107, 100, 184, 77, 37, 129, 52, 203, 217, 227, 65, 83, 215, 213, 59, 71, 32,
809            172, 253, 155, 204, 111, 183, 213, 157, 155
810        ]);
811        drop(vm);
812
813        let cache1 = RandomXCache::new(flags, key.as_bytes()).unwrap();
814        let dataset1 = RandomXDataset::new(flags, cache1.clone(), 0).unwrap();
815        let vm1 = RandomXVM::new(flags, Some(cache1.clone()), Some(dataset1.clone())).unwrap();
816        drop(dataset1);
817        drop(cache1);
818        let hash1 = vm1.calculate_hash(input.as_bytes()).expect("no data");
819        assert_eq!(hash1, [
820            114, 81, 192, 5, 165, 242, 107, 100, 184, 77, 37, 129, 52, 203, 217, 227, 65, 83, 215, 213, 59, 71, 32,
821            172, 253, 155, 204, 111, 183, 213, 157, 155
822        ]);
823        drop(vm1);
824    }
825
826    #[test]
827    fn randomx_hash_fast_vs_light() {
828        let input = b"input";
829        let key = b"key";
830
831        let flags = RandomXFlag::get_recommended_flags() | RandomXFlag::FLAG_FULL_MEM;
832        let cache = RandomXCache::new(flags, key).unwrap();
833        let dataset = RandomXDataset::new(flags, cache, 0).unwrap();
834        let fast_vm = RandomXVM::new(flags, None, Some(dataset)).unwrap();
835
836        let flags = RandomXFlag::get_recommended_flags();
837        let cache = RandomXCache::new(flags, key).unwrap();
838        let light_vm = RandomXVM::new(flags, Some(cache), None).unwrap();
839
840        let fast = fast_vm.calculate_hash(input).unwrap();
841        let light = light_vm.calculate_hash(input).unwrap();
842        assert_eq!(fast, light);
843    }
844
845    #[test]
846    fn test_vectors_fast_mode() {
847        // test vectors from https://github.com/tevador/RandomX/blob/040f4500a6e79d54d84a668013a94507045e786f/src/tests/tests.cpp#L963-L979
848        let key = b"test key 000";
849        let vectors = [
850            (
851                b"This is a test".as_slice(),
852                "639183aae1bf4c9a35884cb46b09cad9175f04efd7684e7262a0ac1c2f0b4e3f",
853            ),
854            (
855                b"Lorem ipsum dolor sit amet".as_slice(),
856                "300a0adb47603dedb42228ccb2b211104f4da45af709cd7547cd049e9489c969",
857            ),
858            (
859                b"sed do eiusmod tempor incididunt ut labore et dolore magna aliqua".as_slice(),
860                "c36d4ed4191e617309867ed66a443be4075014e2b061bcdaf9ce7b721d2b77a8",
861            ),
862        ];
863
864        let flags = RandomXFlag::get_recommended_flags() | RandomXFlag::FLAG_FULL_MEM;
865        let cache = RandomXCache::new(flags, key).unwrap();
866        let dataset = RandomXDataset::new(flags, cache, 0).unwrap();
867        let vm = RandomXVM::new(flags, None, Some(dataset)).unwrap();
868
869        for (input, expected) in vectors {
870            let hash = vm.calculate_hash(input).unwrap();
871            assert_eq!(hex::decode(expected).unwrap(), hash);
872        }
873    }
874
875    #[test]
876    fn test_vectors_light_mode() {
877        // test vectors from https://github.com/tevador/RandomX/blob/040f4500a6e79d54d84a668013a94507045e786f/src/tests/tests.cpp#L963-L985
878        let vectors = [
879            (
880                b"test key 000",
881                b"This is a test".as_slice(),
882                "639183aae1bf4c9a35884cb46b09cad9175f04efd7684e7262a0ac1c2f0b4e3f",
883            ),
884            (
885                b"test key 000",
886                b"Lorem ipsum dolor sit amet".as_slice(),
887                "300a0adb47603dedb42228ccb2b211104f4da45af709cd7547cd049e9489c969",
888            ),
889            (
890                b"test key 000",
891                b"sed do eiusmod tempor incididunt ut labore et dolore magna aliqua".as_slice(),
892                "c36d4ed4191e617309867ed66a443be4075014e2b061bcdaf9ce7b721d2b77a8",
893            ),
894            (
895                b"test key 001",
896                b"sed do eiusmod tempor incididunt ut labore et dolore magna aliqua".as_slice(),
897                "e9ff4503201c0c2cca26d285c93ae883f9b1d30c9eb240b820756f2d5a7905fc",
898            ),
899        ];
900
901        let flags = RandomXFlag::get_recommended_flags();
902        for (key, input, expected) in vectors {
903            let cache = RandomXCache::new(flags, key).unwrap();
904            let vm = RandomXVM::new(flags, Some(cache), None).unwrap();
905            let hash = vm.calculate_hash(input).unwrap();
906            assert_eq!(hex::decode(expected).unwrap(), hash);
907        }
908    }
909}