randomx_rust_wrapper/
vm.rs

1/*
2 * Copyright 2024 Fluence Labs Limited
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use crate::bindings::vm::*;
18use crate::cache::Cache;
19use crate::cache::CacheRawAPI;
20use crate::dataset::Dataset;
21use crate::dataset::DatasetRawAPI;
22use crate::errors::VmCreationError;
23use crate::flags::RandomXFlags;
24use crate::result_hash::ResultHash;
25use crate::try_alloc;
26use crate::RResult;
27
28#[derive(Debug)]
29pub struct RandomXVM<T> {
30    vm: *mut randomx_vm,
31    // too ensure that state outlives VM
32    _state: T,
33}
34
35impl<T> RandomXVM<T>
36where
37    T: CacheRawAPI,
38{
39    pub fn light(cache: T, flags: RandomXFlags) -> RResult<Self> {
40        if !flags.is_light_mode() {
41            return Err(VmCreationError::IncorrectLightModeFlag { flags })?;
42        }
43
44        let vm = try_alloc! { randomx_create_vm(flags.bits(), cache.raw(), std::ptr::null_mut()), VmCreationError::AllocationFailed {flags} };
45
46        let vm = RandomXVM { vm, _state: cache };
47        Ok(vm)
48    }
49
50    /// (from RandomX doc) Reinitializes a virtual machine with a new Cache.
51    /// This function should be called anytime the Cache is reinitialized with a new key.
52    /// Does nothing if called with a Cache containing the same key value as already set.
53    pub fn set_new_cache(&mut self, cache: &'_ Cache) {
54        unsafe { randomx_vm_set_cache(self.vm, cache.raw()) }
55    }
56}
57
58impl<T> RandomXVM<T>
59where
60    T: DatasetRawAPI,
61{
62    pub fn fast(dataset: T, flags: RandomXFlags) -> RResult<Self> {
63        if !flags.is_fast_mode() {
64            return Err(VmCreationError::IncorrectFastModeFlag { flags })?;
65        }
66
67        let vm = try_alloc! { randomx_create_vm(flags.bits(), std::ptr::null_mut(), dataset.raw()), VmCreationError::AllocationFailed {flags} };
68
69        let vm = RandomXVM {
70            vm,
71            _state: dataset,
72        };
73
74        Ok(vm)
75    }
76
77    /// Reinitializes a virtual machine with a new Dataset.
78    pub fn set_new_dataset(&mut self, dataset: &'_ Dataset) {
79        unsafe { randomx_vm_set_dataset(self.vm, dataset.raw()) }
80    }
81}
82
83impl<T> Drop for RandomXVM<T> {
84    fn drop(&mut self) {
85        unsafe { randomx_destroy_vm(self.vm) }
86    }
87}
88
89impl<T> RandomXVM<T> {
90    /// Calculates a RandomX hash value.
91    pub fn hash(&self, local_nonce: &[u8]) -> ResultHash {
92        let mut hash = ResultHash::empty();
93
94        unsafe {
95            randomx_calculate_hash(
96                self.vm,
97                local_nonce.as_ptr() as *const std::ffi::c_void,
98                local_nonce.len(),
99                hash.as_raw_mut(),
100            )
101        };
102
103        hash
104    }
105
106    /// Begins a RandomX hash calculation.
107    pub fn hash_first(&self, local_nonce: &[u8]) {
108        unsafe {
109            randomx_calculate_hash_first(
110                self.vm,
111                local_nonce.as_ptr() as *const std::ffi::c_void,
112                local_nonce.len(),
113            )
114        };
115    }
116
117    /// Output the hash value of the previous input
118    /// and begin the calculation of the next hash.
119    pub fn hash_next(&self, local_nonce: &[u8]) -> ResultHash {
120        let mut hash = ResultHash::empty();
121
122        unsafe {
123            randomx_calculate_hash_next(
124                self.vm,
125                local_nonce.as_ptr() as *const std::ffi::c_void,
126                local_nonce.len(),
127                hash.as_raw_mut(),
128            )
129        };
130
131        hash
132    }
133
134    /// Output the hash value of the previous input.
135    pub fn hash_last(&self) -> ResultHash {
136        let mut hash = ResultHash::empty();
137
138        unsafe { randomx_calculate_hash_last(self.vm, hash.as_raw_mut()) };
139
140        hash
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use crate::{Cache, Dataset, RandomXFlags, RandomXVM};
147
148    #[test]
149    fn light_no_creates_with_full_mem() {
150        let flags = RandomXFlags::recommended_full_mem();
151        let cache = Cache::new(&[0, 1], flags).unwrap();
152        let vm = RandomXVM::light(cache.handle(), flags);
153
154        assert!(vm.is_err());
155    }
156
157    #[test]
158    fn fast_no_creates_without_full_mem() {
159        let flags = RandomXFlags::recommended();
160        let dataset = Dataset::new(&[0, 1], flags).unwrap();
161        let vm = RandomXVM::fast(dataset.handle(), flags);
162
163        assert!(vm.is_err());
164    }
165}