`d608670`

Parallelize code signature hashing

Authored by mfwolffe <wolffemf@dukes.jmu.edu> 1 week ago

SHA: d60867098ebf0a2599ccfed021f9780185bd210a
Parents: bb7169d
Tree: a9dd134

2 changed files

Status	File	+	-
M	`src/macho/writer.rs`	1	1
M	`src/synth/code_sig.rs`	73	2

src/macho/writer.rsmodified

      out[stroff..end].copy_from_slice(&linkedit_plan.strtab_bytes);
      if let Some(code_signature) = &linkedit_plan.code_signature {
          let start = code_signature.dataoff as usize;
 -        let bytes = code_signature.build(&out[..start]);
 +        let bytes = code_signature.build_with_jobs(&out[..start], opts.parallel_jobs());
          let end = start + bytes.len();
          out[start..end].copy_from_slice(&bytes);
+     }

src/synth/code_sig.rsmodified

 +use std::thread;
++
  use crate::layout::Layout;
  use crate::section::is_executable;
  use crate::LinkOptions;
+     }
      pub fn build(&self, signed_prefix: &[u8]) -> Vec<u8> {
 +        self.build_with_jobs(signed_prefix, 1)
 +    }
++
 +    pub fn build_with_jobs(&self, signed_prefix: &[u8], parallel_jobs: usize) -> Vec<u8> {
          debug_assert_eq!(signed_prefix.len(), self.code_limit as usize);
          let code_slots = code_slots(self.code_limit as usize);
          out.extend_from_slice(self.identifier.as_bytes());
          out.push(0);
 -        for page in signed_prefix.chunks(PAGE_SIZE) {
 -            out.extend_from_slice(&sha256(page));
 +        for hash in page_hashes(signed_prefix, parallel_jobs) {
 +            out.extend_from_slice(&hash);
+         }
          out.resize(padded_len, 0);
          out
+     }
+ }
 +fn page_hashes(data: &[u8], parallel_jobs: usize) -> Vec<[u8; 32]> {
 +    let page_count = code_slots(data.len());
 +    if page_count == 0 {
 +        return Vec::new();
 +    }
 +    let parallel_jobs = parallel_jobs.max(1).min(page_count);
 +    if parallel_jobs == 1 || page_count < 2 {
 +        return data.chunks(PAGE_SIZE).map(sha256).collect();
 +    }
++
 +    let chunk_pages = page_count.div_ceil(parallel_jobs);
 +    let chunk_bytes = PAGE_SIZE * chunk_pages;
 +    thread::scope(|scope| {
 +        let mut handles = Vec::new();
 +        for chunk in data.chunks(chunk_bytes) {
 +            handles
 +                .push(scope.spawn(move || chunk.chunks(PAGE_SIZE).map(sha256).collect::<Vec<_>>()));
 +        }
++
 +        let mut hashes = Vec::with_capacity(page_count);
 +        for handle in handles {
 +            hashes.extend(handle.join().expect("code-signature hash worker panicked"));
 +        }
 +        hashes
 +    })
 +}
++
  fn push_be_u32(out: &mut Vec<u8>, value: u32) {
      out.extend_from_slice(&value.to_be_bytes());
+ }
          assert_eq!(read_be_u32(&blob, 52), 16_512);
          assert_eq!(&blob[108..114], b"apple\0");
+     }
++
 +    #[test]
 +    fn parallel_page_hashes_preserve_serial_order() {
 +        let mut bytes = Vec::with_capacity(PAGE_SIZE * 9 + 123);
 +        for index in 0..PAGE_SIZE * 9 + 123 {
 +            bytes.push((index.wrapping_mul(37).wrapping_add(19) & 0xff) as u8);
 +        }
++
 +        let serial = page_hashes(&bytes, 1);
 +        let parallel = page_hashes(&bytes, 4);
 +        assert_eq!(parallel, serial);
 +        assert_eq!(parallel.len(), 10);
 +    }
++
 +    #[test]
 +    fn parallel_code_signature_matches_single_worker() {
 +        let opts = LinkOptions {
 +            output: Some("parallel".into()),
 +            ..LinkOptions::default()
 +        };
 +        let code_limit = PAGE_SIZE * 11 + 777;
 +        let plan = CodeSignaturePlan::new(
 +            &Layout::empty(crate::OutputKind::Executable, 0),
 +            &opts,
 +            code_limit as u64,
 +            true,
 +        )
 +        .unwrap();
 +        let mut signed_prefix = Vec::with_capacity(code_limit);
 +        for index in 0..code_limit {
 +            signed_prefix.push((index.wrapping_mul(13).wrapping_add(index / 7) & 0xff) as u8);
 +        }
++
 +        assert_eq!(
 +            plan.build_with_jobs(&signed_prefix, 8),
 +            plan.build_with_jobs(&signed_prefix, 1)
 +        );
 +    }
+ }