fortrangoingonforty/afs-ld / d608670

Browse files

Parallelize code signature hashing

Authored by mfwolffe <wolffemf@dukes.jmu.edu>
SHA
d60867098ebf0a2599ccfed021f9780185bd210a
Parents
bb7169d
Tree
a9dd134

2 changed files

StatusFile+-
M src/macho/writer.rs 1 1
M src/synth/code_sig.rs 73 2
src/macho/writer.rsmodified
@@ -427,7 +427,7 @@ pub fn write_finalized_with_linkedit(
427
     out[stroff..end].copy_from_slice(&linkedit_plan.strtab_bytes);
427
     out[stroff..end].copy_from_slice(&linkedit_plan.strtab_bytes);
428
     if let Some(code_signature) = &linkedit_plan.code_signature {
428
     if let Some(code_signature) = &linkedit_plan.code_signature {
429
         let start = code_signature.dataoff as usize;
429
         let start = code_signature.dataoff as usize;
430
-        let bytes = code_signature.build(&out[..start]);
430
+        let bytes = code_signature.build_with_jobs(&out[..start], opts.parallel_jobs());
431
         let end = start + bytes.len();
431
         let end = start + bytes.len();
432
         out[start..end].copy_from_slice(&bytes);
432
         out[start..end].copy_from_slice(&bytes);
433
     }
433
     }
src/synth/code_sig.rsmodified
@@ -1,3 +1,5 @@ 
1
+use std::thread;
2
+
1
 use crate::layout::Layout;
3
 use crate::layout::Layout;
2
 use crate::section::is_executable;
4
 use crate::section::is_executable;
3
 use crate::LinkOptions;
5
 use crate::LinkOptions;
@@ -52,6 +54,10 @@ impl CodeSignaturePlan {
52
     }
54
     }
53
 
55
 
54
     pub fn build(&self, signed_prefix: &[u8]) -> Vec<u8> {
56
     pub fn build(&self, signed_prefix: &[u8]) -> Vec<u8> {
57
+        self.build_with_jobs(signed_prefix, 1)
58
+    }
59
+
60
+    pub fn build_with_jobs(&self, signed_prefix: &[u8], parallel_jobs: usize) -> Vec<u8> {
55
         debug_assert_eq!(signed_prefix.len(), self.code_limit as usize);
61
         debug_assert_eq!(signed_prefix.len(), self.code_limit as usize);
56
 
62
 
57
         let code_slots = code_slots(self.code_limit as usize);
63
         let code_slots = code_slots(self.code_limit as usize);
@@ -92,8 +98,8 @@ impl CodeSignaturePlan {
92
 
98
 
93
         out.extend_from_slice(self.identifier.as_bytes());
99
         out.extend_from_slice(self.identifier.as_bytes());
94
         out.push(0);
100
         out.push(0);
95
-        for page in signed_prefix.chunks(PAGE_SIZE) {
101
+        for hash in page_hashes(signed_prefix, parallel_jobs) {
96
-            out.extend_from_slice(&sha256(page));
102
+            out.extend_from_slice(&hash);
97
         }
103
         }
98
         out.resize(padded_len, 0);
104
         out.resize(padded_len, 0);
99
         out
105
         out
@@ -149,6 +155,33 @@ fn code_slots(code_limit: usize) -> usize {
149
     }
155
     }
150
 }
156
 }
151
 
157
 
158
+fn page_hashes(data: &[u8], parallel_jobs: usize) -> Vec<[u8; 32]> {
159
+    let page_count = code_slots(data.len());
160
+    if page_count == 0 {
161
+        return Vec::new();
162
+    }
163
+    let parallel_jobs = parallel_jobs.max(1).min(page_count);
164
+    if parallel_jobs == 1 || page_count < 2 {
165
+        return data.chunks(PAGE_SIZE).map(sha256).collect();
166
+    }
167
+
168
+    let chunk_pages = page_count.div_ceil(parallel_jobs);
169
+    let chunk_bytes = PAGE_SIZE * chunk_pages;
170
+    thread::scope(|scope| {
171
+        let mut handles = Vec::new();
172
+        for chunk in data.chunks(chunk_bytes) {
173
+            handles
174
+                .push(scope.spawn(move || chunk.chunks(PAGE_SIZE).map(sha256).collect::<Vec<_>>()));
175
+        }
176
+
177
+        let mut hashes = Vec::with_capacity(page_count);
178
+        for handle in handles {
179
+            hashes.extend(handle.join().expect("code-signature hash worker panicked"));
180
+        }
181
+        hashes
182
+    })
183
+}
184
+
152
 fn push_be_u32(out: &mut Vec<u8>, value: u32) {
185
 fn push_be_u32(out: &mut Vec<u8>, value: u32) {
153
     out.extend_from_slice(&value.to_be_bytes());
186
     out.extend_from_slice(&value.to_be_bytes());
154
 }
187
 }
@@ -336,4 +369,42 @@ mod tests {
336
         assert_eq!(read_be_u32(&blob, 52), 16_512);
369
         assert_eq!(read_be_u32(&blob, 52), 16_512);
337
         assert_eq!(&blob[108..114], b"apple\0");
370
         assert_eq!(&blob[108..114], b"apple\0");
338
     }
371
     }
372
+
373
+    #[test]
374
+    fn parallel_page_hashes_preserve_serial_order() {
375
+        let mut bytes = Vec::with_capacity(PAGE_SIZE * 9 + 123);
376
+        for index in 0..PAGE_SIZE * 9 + 123 {
377
+            bytes.push((index.wrapping_mul(37).wrapping_add(19) & 0xff) as u8);
378
+        }
379
+
380
+        let serial = page_hashes(&bytes, 1);
381
+        let parallel = page_hashes(&bytes, 4);
382
+        assert_eq!(parallel, serial);
383
+        assert_eq!(parallel.len(), 10);
384
+    }
385
+
386
+    #[test]
387
+    fn parallel_code_signature_matches_single_worker() {
388
+        let opts = LinkOptions {
389
+            output: Some("parallel".into()),
390
+            ..LinkOptions::default()
391
+        };
392
+        let code_limit = PAGE_SIZE * 11 + 777;
393
+        let plan = CodeSignaturePlan::new(
394
+            &Layout::empty(crate::OutputKind::Executable, 0),
395
+            &opts,
396
+            code_limit as u64,
397
+            true,
398
+        )
399
+        .unwrap();
400
+        let mut signed_prefix = Vec::with_capacity(code_limit);
401
+        for index in 0..code_limit {
402
+            signed_prefix.push((index.wrapping_mul(13).wrapping_add(index / 7) & 0xff) as u8);
403
+        }
404
+
405
+        assert_eq!(
406
+            plan.build_with_jobs(&signed_prefix, 8),
407
+            plan.build_with_jobs(&signed_prefix, 1)
408
+        );
409
+    }
339
 }
410
 }