1use {
2 core::ffi::c_void,
3 dlopen2::symbor::{Container, SymBorApi, Symbol},
4 log::*,
5 std::{
6 env,
7 ffi::OsStr,
8 fs,
9 os::raw::{c_int, c_uint},
10 path::{Path, PathBuf},
11 sync::Once,
12 },
13};
14
15#[repr(C)]
16pub struct Elems {
17 pub elems: *const u8,
18 pub num: u32,
19}
20
21#[derive(SymBorApi)]
22pub struct Api<'a> {
23 pub ed25519_init: Symbol<'a, unsafe extern "C" fn() -> bool>,
24 pub ed25519_set_verbose: Symbol<'a, unsafe extern "C" fn(val: bool)>,
25
26 #[allow(clippy::type_complexity)]
27 pub ed25519_verify_many: Symbol<
28 'a,
29 unsafe extern "C" fn(
30 vecs: *const Elems,
31 num: u32, message_size: u32, total_packets: u32,
34 total_signatures: u32,
35 message_lens: *const u32,
36 pubkey_offsets: *const u32,
37 signature_offsets: *const u32,
38 signed_message_offsets: *const u32,
39 out: *mut u8, use_non_default_stream: u8,
41 ) -> u32,
42 >,
43
44 #[allow(clippy::type_complexity)]
45 pub ed25519_sign_many: Symbol<
46 'a,
47 unsafe extern "C" fn(
48 vecs: *mut Elems,
49 num: u32, message_size: u32, total_packets: u32,
52 total_signatures: u32,
53 message_lens: *const u32,
54 pubkey_offsets: *const u32,
55 privkey_offsets: *const u32,
56 signed_message_offsets: *const u32,
57 sgnatures_out: *mut u8, use_non_default_stream: u8,
59 ) -> u32,
60 >,
61
62 pub poh_verify_many: Symbol<
63 'a,
64 unsafe extern "C" fn(
65 hashes: *mut u8,
66 num_hashes_arr: *const u64,
67 num_elems: usize,
68 use_non_default_stream: u8,
69 ) -> c_int,
70 >,
71
72 pub cuda_host_register:
73 Symbol<'a, unsafe extern "C" fn(ptr: *mut c_void, size: usize, flags: c_uint) -> c_int>,
74
75 pub cuda_host_unregister: Symbol<'a, unsafe extern "C" fn(ptr: *mut c_void) -> c_int>,
76
77 pub ed25519_get_checked_scalar:
78 Symbol<'a, unsafe extern "C" fn(out_scalar: *mut u8, in_scalar: *const u8) -> c_int>,
79
80 pub ed25519_check_packed_ge_small_order:
81 Symbol<'a, unsafe extern "C" fn(packed_ge: *const u8) -> c_int>,
82}
83
84static mut API: Option<Container<Api>> = None;
85
86fn init(name: &OsStr) {
87 static INIT_HOOK: Once = Once::new();
88
89 info!("Loading {:?}", name);
90 unsafe {
91 INIT_HOOK.call_once(|| {
92 API = Some(Container::load(name).unwrap_or_else(|err| {
93 error!("Unable to load {:?}: {}", name, err);
94 std::process::exit(1);
95 }));
96 })
97 }
98}
99
100pub fn locate_perf_libs() -> Option<PathBuf> {
101 let exe = env::current_exe().expect("Unable to get executable path");
102 let perf_libs = exe.parent().unwrap().join("perf-libs");
103 if perf_libs.is_dir() {
104 info!("perf-libs found at {:?}", perf_libs);
105 return Some(perf_libs);
106 }
107 warn!("{:?} does not exist", perf_libs);
108 None
109}
110
111fn find_cuda_home(perf_libs_path: &Path) -> Option<PathBuf> {
112 if let Ok(cuda_home) = env::var("CUDA_HOME") {
113 let path = PathBuf::from(cuda_home);
114 if path.is_dir() {
115 info!("Using CUDA_HOME: {:?}", path);
116 return Some(path);
117 }
118 warn!("Ignoring CUDA_HOME, not a path: {:?}", path);
119 }
120
121 for entry in fs::read_dir(perf_libs_path).unwrap().flatten() {
123 let path = entry.path();
124 if !path.is_dir() {
125 continue;
126 }
127 let dir_name = path.file_name().unwrap().to_str().unwrap_or("");
128 if !dir_name.starts_with("cuda-") {
129 continue;
130 }
131
132 let cuda_home: PathBuf = ["/", "usr", "local", dir_name].iter().collect();
133 if !cuda_home.is_dir() {
134 continue;
135 }
136
137 info!("CUDA installation found at {:?}", cuda_home);
138 return Some(cuda_home);
139 }
140 None
141}
142
143pub fn append_to_ld_library_path(mut ld_library_path: String) {
144 if let Ok(env_value) = env::var("LD_LIBRARY_PATH") {
145 ld_library_path.push(':');
146 ld_library_path.push_str(&env_value);
147 }
148 info!("setting ld_library_path to: {:?}", ld_library_path);
149 env::set_var("LD_LIBRARY_PATH", ld_library_path);
150}
151
152pub fn init_cuda() {
153 if let Some(perf_libs_path) = locate_perf_libs() {
154 if let Some(cuda_home) = find_cuda_home(&perf_libs_path) {
155 let cuda_lib64_dir = cuda_home.join("lib64");
156 if cuda_lib64_dir.is_dir() {
157 append_to_ld_library_path(cuda_lib64_dir.to_str().unwrap_or("").to_string())
160 } else {
161 warn!("CUDA lib64 directory does not exist: {:?}", cuda_lib64_dir);
162 }
163
164 let libcuda_crypt = perf_libs_path
165 .join(cuda_home.file_name().unwrap())
166 .join("libcuda-crypt.so");
167 return init(libcuda_crypt.as_os_str());
168 } else {
169 warn!("CUDA installation not found");
170 }
171 }
172
173 init(OsStr::new("libcuda-crypt.so"))
175}
176
177pub fn api() -> Option<&'static Container<Api<'static>>> {
178 {
179 static INIT_HOOK: Once = Once::new();
180 INIT_HOOK.call_once(|| {
181 if std::env::var("TEST_PERF_LIBS_CUDA").is_ok() {
182 init_cuda();
183 }
184 })
185 }
186
187 unsafe { API.as_ref() }
188}