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