<commit_before>extern crate gnuplot;

use std::collections::HashMap;
use std::env;

use std::path::Path;
use std::fs;
use std::u64;
use std::fs::DirEntry;

use gnuplot::*;

fn get_dir_size(direntry: &DirEntry) -> u64 {
  // Initialize with the size of the current directory
  let mut size = direntry.metadata().unwrap().len();
  // Sum up the current file children
  for subfile in get_subfiles(direntry).iter() {
    let subfile_size = subfile.metadata().unwrap().len();
    size += subfile_size;
  }
  // And recurse into subdirectories
  for subdir in get_subdirs(direntry).iter() {
    size += get_dir_size(subdir); }
  size
}

fn is_dir(direntry : &DirEntry) -> bool {
  let metadata = direntry.metadata();
  return match metadata {
    Err(_) => false,
    Ok(metadata) => metadata.is_dir()
  }
}

fn get_subdirs(direntry : &DirEntry) -> Vec<DirEntry> {
  get_path_subdirs(direntry.path().as_path())
}
fn get_subfiles(direntry : &DirEntry) -> Vec<DirEntry> {
  get_path_subfiles(direntry.path().as_path())
}

fn get_path_subdirs(path : &Path) -> Vec<DirEntry> {
  let children = fs::read_dir(path).unwrap();
  let children_entries = children.map(|c| c.unwrap()).collect::<Vec<_>>();
  children_entries.into_iter().filter(|c| is_dir(c)).collect::<Vec<_>>()
}
fn get_path_subfiles(path : &Path) -> Vec<DirEntry> {
  let children = fs::read_dir(path).unwrap();
  let children_entries = children.map(|c| c.unwrap()).collect::<Vec<_>>();
  children_entries.into_iter().filter(|c| !is_dir(c)).collect::<Vec<_>>()
}

fn main() {
  let args: Vec<_> = env::args().collect();
  let mut min_size = u64::MAX;
  let mut max_size = 0;
  let arxiv_root = &Path::new(&args[1]);
  let mut arxiv_size_frequencies: HashMap<u64, u64> = HashMap::new();
  let mut arxiv_monthly_sizes: HashMap<String, u64> = HashMap::new();
  let mut arxiv_counts: HashMap<String, u64> = HashMap::new();
  let arxiv_month_dirs = get_path_subdirs(arxiv_root);

  for month_dir in arxiv_month_dirs.iter() {
    let mut monthly_size = 0;
    let month_papers = get_subdirs(month_dir);
    // Record how many papers were submitted in that month
    arxiv_counts.insert(month_dir.file_name().to_str().unwrap().to_string(), month_papers.len() as u64);

    for paper_dir in month_papers.iter() {
      // Record the size of each paper directory
      let paper_size = get_dir_size(paper_dir) / 1024; // In KB
      if paper_size > max_size {
        max_size = paper_size;
      }
      if paper_size < min_size {
        min_size = paper_size;
      }
      monthly_size += paper_size;

      let size_frequency = arxiv_size_frequencies.entry(paper_size).or_insert(0);
      *size_frequency += 1;
    }
    arxiv_monthly_sizes.insert(month_dir.file_name().to_str().unwrap().to_string(), monthly_size);
  }

  // Print the recorded stats
  println!("Counts: {:?}", arxiv_counts);
  println!("Sizes: {:?}", arxiv_monthly_sizes);

  // arXiv months in order
  // TODO: Only current as of 1505, you'll have to extend manually next time this is run
  let ordered_months = [
    "9107", "9108", "9109", "9110", "9111", "9112", "9201", "9202", "9203", "9204", "9205", "9206", "9207", "9208", "9209",
    "9210", "9211", "9212", "9301", "9302", "9303", "9304", "9305", "9306", "9307", "9308", "9309", "9310", "9311", "9312",
    "9401", "9402", "9403", "9404", "9405", "9406", "9407", "9408", "9409", "9410", "9411", "9412", "9501", "9502", "9503",
    "9504", "9505", "9506", "9507", "9508", "9509", "9510", "9511", "9512", "9601", "9602", "9603", "9604", "9605", "9606",
    "9607", "9608", "9609", "9610", "9611", "9612", "9701", "9702", "9703", "9704", "9705", "9706", "9707", "9708", "9709",
    "9710", "9711", "9712", "9801", "9802", "9803", "9804", "9805", "9806", "9807", "9808", "9809", "9810", "9811", "9812",
    "9901", "9902", "9903", "9904", "9905", "9906", "9907", "9908", "9909", "9910", "9911", "9912", "0001", "0002", "0003",
    "0004", "0005", "0006", "0007", "0008", "0009", "0010", "0011", "0012", "0101", "0102", "0103", "0104", "0105", "0106",
    "0107", "0108", "0109", "0110", "0111", "0112", "0201", "0202", "0203", "0204", "0205", "0206", "0207", "0208", "0209",
    "0210", "0211", "0212", "0301", "0302", "0303", "0304", "0305", "0306", "0307", "0308", "0309", "0310", "0311", "0312",
    "0401", "0402", "0403", "0404", "0405", "0406", "0407", "0408", "0409", "0410", "0411", "0412", "0501", "0502", "0503",
    "0504", "0505", "0506", "0507", "0508", "0509", "0510", "0511", "0512", "0601", "0602", "0603", "0604", "0605", "0606",
    "0607", "0608", "0609", "0610", "0611", "0612", "0701", "0702", "0703", "0704", "0705", "0706", "0707", "0708", "0709",
    "0710", "0711", "0712", "0801", "0802", "0803", "0804", "0805", "0806", "0807", "0808", "0809", "0810", "0811", "0812",
    "0901", "0902", "0903", "0904", "0905", "0906", "0907", "0908", "0909", "0910", "0911", "0912", "1001", "1002", "1003",
    "1004", "1005", "1006", "1007", "1008", "1009", "1010", "1011", "1012", "1101", "1102", "1103", "1104", "1105", "1106",
    "1107", "1108", "1109", "1110", "1111", "1112", "1201", "1202", "1203", "1204", "1205", "1206", "1207", "1208", "1209",
    "1210", "1211", "1212", "1301", "1302", "1303", "1304", "1305", "1306", "1307", "1308", "1309", "1310", "1311", "1312",
    "1401", "1402", "1403", "1404", "1405", "1406", "1407", "1408", "1409", "1410", "1411", "1412", "1501", "1502", "1503",
    "1504", "1505"];

  // Plot of paper counts by month:
  let zero = 0 as u64;
  let ordered_counts = ordered_months.iter().map(|m|
    match arxiv_counts.get(&m.to_string()) {
        Some(counts) => counts,
        None => &zero
    });
  let ordered_month_xcoords = 1..ordered_months.len()+1;

  let mut fg = Figure::new();
  fg.axes2d()
  .points(ordered_month_xcoords.clone(), ordered_counts, &[PointSymbol('D'), Color("#ffaa77"), PointSize(0.5)])
  .set_x_label("arXiv month", &[Rotate(45.0)])
  .set_y_label("Submitted papers", &[Rotate(90.0)])
  .set_title("arXiv TeX submission counts", &[]);

  fg.set_terminal("pngcairo", "arxiv_submission_counts.png");
  fg.show();

  // Plot of submission counts by month:
  let ordered_sizes = ordered_months.iter().map(|m|
    match arxiv_monthly_sizes.get(&m.to_string()) {
        Some(&counts) => counts,
        None => zero
    });
  fg = Figure::new();
  fg.axes2d()
  .points(ordered_month_xcoords.clone(), ordered_sizes, &[PointSymbol('D'), Color("#ffaa77"), PointSize(0.5)])
  .set_x_label("arXiv month", &[Rotate(45.0)])
  .set_y_label("Submission size in MB", &[Rotate(90.0)])
  .set_title("arXiv TeX submission sizes by month", &[]);

  fg.set_terminal("pngcairo", "arxiv_submission_sizes.png");
  fg.show();

  // Plot average paper size in KB
  // Plot of submission size by month:

  let freq_keys = arxiv_size_frequencies.clone().into_iter().map(|entry| entry.0);
  let freq_values = arxiv_size_frequencies.clone().into_iter().map(|entry| entry.1);

  fg = Figure::new();
  fg.axes2d()
  .points(freq_keys, freq_values, &[PointSymbol('D'), Color("#ffaa77"), PointSize(0.5)])
  .set_x_label("Paper size in KB", &[Rotate(45.0)])
  .set_y_label("Paper count", &[Rotate(90.0)])
  .set_title("arXiv TeX paper sizes", &[]);

  fg.set_terminal("pngcairo", "arxiv_paper_sizes.png");
  fg.show();

}<commit_msg>add a log message<commit_after>extern crate gnuplot;

use std::collections::HashMap;
use std::env;

use std::path::Path;
use std::fs;
use std::u64;
use std::fs::DirEntry;

use gnuplot::*;

fn get_dir_size(direntry: &DirEntry) -> u64 {
  // Initialize with the size of the current directory
  let mut size = direntry.metadata().unwrap().len();
  // Sum up the current file children
  for subfile in get_subfiles(direntry).iter() {
    let subfile_size = subfile.metadata().unwrap().len();
    size += subfile_size;
  }
  // And recurse into subdirectories
  for subdir in get_subdirs(direntry).iter() {
    size += get_dir_size(subdir); }
  size
}

fn is_dir(direntry : &DirEntry) -> bool {
  let metadata = direntry.metadata();
  return match metadata {
    Err(_) => false,
    Ok(metadata) => metadata.is_dir()
  }
}

fn get_subdirs(direntry : &DirEntry) -> Vec<DirEntry> {
  get_path_subdirs(direntry.path().as_path())
}
fn get_subfiles(direntry : &DirEntry) -> Vec<DirEntry> {
  get_path_subfiles(direntry.path().as_path())
}

fn get_path_subdirs(path : &Path) -> Vec<DirEntry> {
  let children = fs::read_dir(path).unwrap();
  let children_entries = children.map(|c| c.unwrap()).collect::<Vec<_>>();
  children_entries.into_iter().filter(|c| is_dir(c)).collect::<Vec<_>>()
}
fn get_path_subfiles(path : &Path) -> Vec<DirEntry> {
  let children = fs::read_dir(path).unwrap();
  let children_entries = children.map(|c| c.unwrap()).collect::<Vec<_>>();
  children_entries.into_iter().filter(|c| !is_dir(c)).collect::<Vec<_>>()
}

fn main() {
  let args: Vec<_> = env::args().collect();
  let mut min_size = u64::MAX;
  let mut max_size = 0;
  let arxiv_root = &Path::new(&args[1]);
  let mut arxiv_size_frequencies: HashMap<u64, u64> = HashMap::new();
  let mut arxiv_monthly_sizes: HashMap<String, u64> = HashMap::new();
  let mut arxiv_counts: HashMap<String, u64> = HashMap::new();
  let arxiv_month_dirs = get_path_subdirs(arxiv_root);

  for month_dir in arxiv_month_dirs.iter() {
    println!("-- Measuring {:?}", month_dir.file_name().to_str().unwrap());
    let mut monthly_size = 0;
    let month_papers = get_subdirs(month_dir);
    // Record how many papers were submitted in that month
    arxiv_counts.insert(month_dir.file_name().to_str().unwrap().to_string(), month_papers.len() as u64);

    for paper_dir in month_papers.iter() {
      // Record the size of each paper directory
      let paper_size = get_dir_size(paper_dir) / 1024; // In KB
      if paper_size > max_size {
        max_size = paper_size;
      }
      if paper_size < min_size {
        min_size = paper_size;
      }
      monthly_size += paper_size;

      let size_frequency = arxiv_size_frequencies.entry(paper_size).or_insert(0);
      *size_frequency += 1;
    }
    arxiv_monthly_sizes.insert(month_dir.file_name().to_str().unwrap().to_string(), monthly_size);
  }

  // Print the recorded stats
  println!("Counts: {:?}", arxiv_counts);
  println!("Sizes: {:?}", arxiv_monthly_sizes);

  // arXiv months in order
  // TODO: Only current as of 1505, you'll have to extend manually next time this is run
  let ordered_months = [
    "9107", "9108", "9109", "9110", "9111", "9112", "9201", "9202", "9203", "9204", "9205", "9206", "9207", "9208", "9209",
    "9210", "9211", "9212", "9301", "9302", "9303", "9304", "9305", "9306", "9307", "9308", "9309", "9310", "9311", "9312",
    "9401", "9402", "9403", "9404", "9405", "9406", "9407", "9408", "9409", "9410", "9411", "9412", "9501", "9502", "9503",
    "9504", "9505", "9506", "9507", "9508", "9509", "9510", "9511", "9512", "9601", "9602", "9603", "9604", "9605", "9606",
    "9607", "9608", "9609", "9610", "9611", "9612", "9701", "9702", "9703", "9704", "9705", "9706", "9707", "9708", "9709",
    "9710", "9711", "9712", "9801", "9802", "9803", "9804", "9805", "9806", "9807", "9808", "9809", "9810", "9811", "9812",
    "9901", "9902", "9903", "9904", "9905", "9906", "9907", "9908", "9909", "9910", "9911", "9912", "0001", "0002", "0003",
    "0004", "0005", "0006", "0007", "0008", "0009", "0010", "0011", "0012", "0101", "0102", "0103", "0104", "0105", "0106",
    "0107", "0108", "0109", "0110", "0111", "0112", "0201", "0202", "0203", "0204", "0205", "0206", "0207", "0208", "0209",
    "0210", "0211", "0212", "0301", "0302", "0303", "0304", "0305", "0306", "0307", "0308", "0309", "0310", "0311", "0312",
    "0401", "0402", "0403", "0404", "0405", "0406", "0407", "0408", "0409", "0410", "0411", "0412", "0501", "0502", "0503",
    "0504", "0505", "0506", "0507", "0508", "0509", "0510", "0511", "0512", "0601", "0602", "0603", "0604", "0605", "0606",
    "0607", "0608", "0609", "0610", "0611", "0612", "0701", "0702", "0703", "0704", "0705", "0706", "0707", "0708", "0709",
    "0710", "0711", "0712", "0801", "0802", "0803", "0804", "0805", "0806", "0807", "0808", "0809", "0810", "0811", "0812",
    "0901", "0902", "0903", "0904", "0905", "0906", "0907", "0908", "0909", "0910", "0911", "0912", "1001", "1002", "1003",
    "1004", "1005", "1006", "1007", "1008", "1009", "1010", "1011", "1012", "1101", "1102", "1103", "1104", "1105", "1106",
    "1107", "1108", "1109", "1110", "1111", "1112", "1201", "1202", "1203", "1204", "1205", "1206", "1207", "1208", "1209",
    "1210", "1211", "1212", "1301", "1302", "1303", "1304", "1305", "1306", "1307", "1308", "1309", "1310", "1311", "1312",
    "1401", "1402", "1403", "1404", "1405", "1406", "1407", "1408", "1409", "1410", "1411", "1412", "1501", "1502", "1503",
    "1504", "1505"];

  // Plot of paper counts by month:
  let zero = 0 as u64;
  let ordered_counts = ordered_months.iter().map(|m|
    match arxiv_counts.get(&m.to_string()) {
        Some(counts) => counts,
        None => &zero
    });
  let ordered_month_xcoords = 1..ordered_months.len()+1;

  let mut fg = Figure::new();
  fg.axes2d()
  .points(ordered_month_xcoords.clone(), ordered_counts, &[PointSymbol('D'), Color("#ffaa77"), PointSize(0.5)])
  .set_x_label("arXiv month", &[Rotate(45.0)])
  .set_y_label("Submitted papers", &[Rotate(90.0)])
  .set_title("arXiv TeX submission counts", &[]);

  fg.set_terminal("pngcairo", "arxiv_submission_counts.png");
  fg.show();

  // Plot of submission counts by month:
  let ordered_sizes = ordered_months.iter().map(|m|
    match arxiv_monthly_sizes.get(&m.to_string()) {
        Some(&counts) => counts,
        None => zero
    });
  fg = Figure::new();
  fg.axes2d()
  .points(ordered_month_xcoords.clone(), ordered_sizes, &[PointSymbol('D'), Color("#ffaa77"), PointSize(0.5)])
  .set_x_label("arXiv month", &[Rotate(45.0)])
  .set_y_label("Submission size in MB", &[Rotate(90.0)])
  .set_title("arXiv TeX submission sizes by month", &[]);

  fg.set_terminal("pngcairo", "arxiv_submission_sizes.png");
  fg.show();

  // Plot average paper size in KB
  // Plot of submission size by month:

  let freq_keys = arxiv_size_frequencies.clone().into_iter().map(|entry| entry.0);
  let freq_values = arxiv_size_frequencies.clone().into_iter().map(|entry| entry.1);

  fg = Figure::new();
  fg.axes2d()
  .points(freq_keys, freq_values, &[PointSymbol('D'), Color("#ffaa77"), PointSize(0.5)])
  .set_x_label("Paper size in KB", &[Rotate(45.0)])
  .set_y_label("Paper count", &[Rotate(90.0)])
  .set_title("arXiv TeX paper sizes", &[]);

  fg.set_terminal("pngcairo", "arxiv_paper_sizes.png");
  fg.show();

}<|endoftext|>


<commit_before>use rand::os::OsRng;
use rand::XorShiftRng;

pub struct FactionGen {
    rand: XorShiftRng,
    faction_amount: usize.
}

impl FactionGen {
	pub fn new(faction_amount: usize) {
		rand: XorShiftRng::rand(&mut OsRng::new().unwrap()),
		faction_amount: faction_amount,
	}

	pub fn generate(&mut self, level: &mut Level<Faction>) {
		for x in 0..level.get_width() {
			for y in 0..level.get_height() {
				match level.get_mut_tile(x, y) {
					Ok(tile) => {

					},
					Err(Error::IndexOutOfBounds) => {
						panic!("Generate method indexed out of bounds while simulating a step. This should never happen unless the programmer is not very bright.");
					}
				}
			}
		}
	}
}<commit_msg>basic logic for faction algorithm<commit_after>use rand::os::OsRng;
use rand::XorShiftRng;

pub struct FactionGen {
    rand: XorShiftRng,
    faction_amount: usize.
}

impl FactionGen {
	pub fn new(faction_amount: usize) {
		rand: XorShiftRng::rand(&mut OsRng::new().unwrap()),
		faction_amount: faction_amount,
	}

	pub fn generate(&mut self, level: &mut Level<Faction>) {
		for x in 0..level.get_width() {
			for y in 0..level.get_height() {
				match level.get_mut_tile(x, y) {
					Ok(tile) => {
						let mut deck = Vec::new();
						match *tile {
							Faction::Faction(_) => {
								deck.push(tile);
							},
							Faction::Void => continue,
							_ => {},
						}
						get_faction_neighbours(&mut deck, &mut level);
						self.rand.shuffle(deck);
						*tile = deck.pop();
					},
					Err(Error::IndexOutOfBounds) => {
						panic!("Generate method indexed out of bounds while simulating a step. This should never happen unless the programmer is not very bright.");
					}
				}
			}
		}
	}
}<|endoftext|>


<commit_before>const BAD_KEY: u64 = 0x5252_5252_5252_5252; //This encrypts to 0.

///Basic sanity tests of the cypto properties of aHash.
#[cfg(test)]
mod aes_tests {
    use std::hash::{Hash, Hasher};
    use crate::aes_hash::*;
    use crate::hash_quality_test::*;

    pub fn ahash<H: Hash>(b: H) -> u64 {
        let mut hasher = AesHasher::new_with_keys(BAD_KEY, BAD_KEY);
        b.hash(&mut hasher);
        hasher.finish()
    }

    #[test]
    fn test_single_bit_in_byte() {
        let mut hasher1= AesHasher::new_with_keys(64, 64);
        8_u32.hash(&mut hasher1);
        let mut hasher2= AesHasher::new_with_keys(64, 64);
        0_u32.hash(&mut hasher2);
        assert_sufficiently_different(hasher1.finish(), hasher2.finish());
    }

    #[test]
    fn test_single_bit_flip() {
        let size = 32;
        let compare_value = ahash(0u32);
        for pos in 0..size {
            let test_value = ahash(0 ^ (1u32 << pos));
            assert_sufficiently_different(compare_value, test_value);
        }
        let size = 64;
        let compare_value = ahash(0u64);
        for pos in 0..size {
            let test_value = ahash(0 ^ (1u64 << pos));
            assert_sufficiently_different(compare_value, test_value);
        }
        let size = 128;
        let compare_value = ahash(0u128);
        for pos in 0..size {
            let test_value = ahash(0 ^ (1u128 << pos));
            assert_sufficiently_different(compare_value, test_value);
        }
    }

    #[test]
    fn test_keys_change_output() {
        let mut a = AesHasher::new_with_keys(0, 0);
        let mut b = AesHasher::new_with_keys(0, 1);
        let mut c = AesHasher::new_with_keys(1, 0);
        let mut d = AesHasher::new_with_keys(1, 1);
        "test".hash(&mut a);
        "test".hash(&mut b);
        "test".hash(&mut c);
        "test".hash(&mut d);
        assert_sufficiently_different(a.finish(), b.finish());
        assert_sufficiently_different(a.finish(), c.finish());
        assert_sufficiently_different(a.finish(), d.finish());
        assert_sufficiently_different(b.finish(), c.finish());
        assert_sufficiently_different(b.finish(), d.finish());
        assert_sufficiently_different(c.finish(), d.finish());
    }

    #[test]
    fn test_finish_is_consistant() {
        let mut hasher = AesHasher::new_with_keys(1, 2);
        "Foo".hash(&mut hasher);
        let a = hasher.finish();
        let b = hasher.finish();
        assert_eq!(a, b);
    }

    #[test]
    fn test_single_key_bit_flip() {
        for bit in 0..64 {
            let mut a = AesHasher::new_with_keys(0, 0);
            let mut b = AesHasher::new_with_keys(0, 1 << bit);
            let mut c = AesHasher::new_with_keys(1 << bit, 0);
            "1234".hash(&mut a);
            "1234".hash(&mut b);
            "1234".hash(&mut c);
            assert_sufficiently_different(a.finish(), b.finish());
            assert_sufficiently_different(a.finish(), c.finish());
            assert_sufficiently_different(b.finish(), c.finish());
            let mut a = AesHasher::new_with_keys(0, 0);
            let mut b = AesHasher::new_with_keys(0, 1 << bit);
            let mut c = AesHasher::new_with_keys(1 << bit, 0);
            "12345678".hash(&mut a);
            "12345678".hash(&mut b);
            "12345678".hash(&mut c);
            assert_sufficiently_different(a.finish(), b.finish());
            assert_sufficiently_different(a.finish(), c.finish());
            assert_sufficiently_different(b.finish(), c.finish());
            let mut a = AesHasher::new_with_keys(0, 0);
            let mut b = AesHasher::new_with_keys(0, 1 << bit);
            let mut c = AesHasher::new_with_keys(1 << bit, 0);
            "1234567812345678".hash(&mut a);
            "1234567812345678".hash(&mut b);
            "1234567812345678".hash(&mut c);
            assert_sufficiently_different(a.finish(), b.finish());
            assert_sufficiently_different(a.finish(), c.finish());
            assert_sufficiently_different(b.finish(), c.finish());
        }
    }

    #[test]
    fn test_padding_doesnot_collide() {
        for c in 0..128u8 {
            for string in ["", "1234", "12345678", "1234567812345678"].into_iter() {
                let mut short = AesHasher::default();
                string.hash(&mut short);
                let value = short.finish();
                let mut string = string.to_string();
                for num in 1..=128 {
                    let mut long = AesHasher::default();
                    string.push(c as char);
                    string.hash(&mut long);
                    let same_bytes = count_same_bytes(value, long.finish());
                    assert!(same_bytes <= 2, format!("{} bytes of {} -> {:x} vs {:x}", num, c, value, long.finish()));
                    let flipped_bits = (value ^ long.finish()).count_ones();
                    assert!(flipped_bits > 10);
                }
            }
        }
    }
}

fn assert_sufficiently_different(a: u64, b: u64) {
    let same_byte_count = count_same_bytes(a, b);
    assert!(same_byte_count <= 1);
    let flipped_bits = (a ^ b).count_ones();
    assert!(flipped_bits > 18 && flipped_bits < 48, "{:x} and {:x}: {:}", a, b, flipped_bits);
    for rotate in 0..64 {
        let flipped_bits2 = (a ^ (b.rotate_left(rotate))).count_ones();
        assert!(flipped_bits2 > 14 && flipped_bits2 < 50, "{:x} and {:x}: {:}", a, b.rotate_left(rotate), flipped_bits2);
    }
}

fn count_same_bytes(a: u64, b: u64) -> i32 {
    let mut same_byte_count = 0;
    for byte in 0..8 {
        let ba = a >> 8 * byte as u8;
        let bb = b >> 8 * byte as u8;
        if ba == bb {
            same_byte_count += 1;
        }
    }
    same_byte_count
}

#[cfg(test)]
mod fallback_tests {
    use std::hash::{Hash, Hasher};
    use crate::fallback_hash::*;
    use crate::hash_quality_test::*;

    pub fn fallback<H: Hash>(b: H) -> u64 {
        let mut hasher = FallbackHasher::new_with_keys(BAD_KEY, BAD_KEY);
        b.hash(&mut hasher);
        hasher.finish()
    }

    #[test]
    fn test_single_bit_flip() {
        let size = 32;
        let compare_value = fallback(0u32);
        for pos in 0..size {
            let test_value = fallback(0 ^ (1u32 << pos));
            assert_sufficiently_different(compare_value, test_value);
        }
        let size = 64;
        let compare_value = fallback(0u64);
        for pos in 0..size {
            let test_value = fallback(0 ^ (1u64 << pos));
            assert_sufficiently_different(compare_value, test_value);
        }
        let size = 128;
        let compare_value = fallback(0u128);
        for pos in 0..size {
            let test_value = fallback(0 ^ (1u128 << pos));
            assert_sufficiently_different(compare_value, test_value);
        }
    }

    #[test]
    fn test_keys_change_output() {
        let mut a = FallbackHasher::new_with_keys(0, 0);
        let mut b = FallbackHasher::new_with_keys(0, 1);
        let mut c = FallbackHasher::new_with_keys(1, 0);
        let mut d = FallbackHasher::new_with_keys(1, 1);
        "test".hash(&mut a);
        "test".hash(&mut b);
        "test".hash(&mut c);
        "test".hash(&mut d);
        assert_sufficiently_different(a.finish(), b.finish());
        assert_sufficiently_different(a.finish(), c.finish());
        assert_sufficiently_different(a.finish(), d.finish());
        assert_sufficiently_different(b.finish(), c.finish());
        assert_sufficiently_different(b.finish(), d.finish());
        assert_sufficiently_different(c.finish(), d.finish());
    }

    #[test]
    fn test_finish_is_consistant() {
        let mut hasher = FallbackHasher::new_with_keys(1, 2);
        "Foo".hash(&mut hasher);
        let a = hasher.finish();
        let b = hasher.finish();
        assert_eq!(a, b);
    }

    #[test]
    fn test_single_key_bit_flip() {
        for bit in 0..64 {
            let mut a = FallbackHasher::new_with_keys(0, 0);
            let mut b = FallbackHasher::new_with_keys(0, 1 << bit);
            let mut c = FallbackHasher::new_with_keys(1 << bit, 0);
            "1234".hash(&mut a);
            "1234".hash(&mut b);
            "1234".hash(&mut c);
            assert_sufficiently_different(a.finish(), b.finish());
            assert_sufficiently_different(a.finish(), c.finish());
            assert_sufficiently_different(b.finish(), c.finish());
            let mut a = FallbackHasher::new_with_keys(0, 0);
            let mut b = FallbackHasher::new_with_keys(0, 1 << bit);
            let mut c = FallbackHasher::new_with_keys(1 << bit, 0);
            "12345678".hash(&mut a);
            "12345678".hash(&mut b);
            "12345678".hash(&mut c);
            assert_sufficiently_different(a.finish(), b.finish());
            assert_sufficiently_different(a.finish(), c.finish());
            assert_sufficiently_different(b.finish(), c.finish());
            let mut a = FallbackHasher::new_with_keys(0, 0);
            let mut b = FallbackHasher::new_with_keys(0, 1 << bit);
            let mut c = FallbackHasher::new_with_keys(1 << bit, 0);
            "1234567812345678".hash(&mut a);
            "1234567812345678".hash(&mut b);
            "1234567812345678".hash(&mut c);
            assert_sufficiently_different(a.finish(), b.finish());
            assert_sufficiently_different(a.finish(), c.finish());
            assert_sufficiently_different(b.finish(), c.finish());
        }
    }

    #[test]
    fn test_padding_doesnot_collide() {
        for c in 0..128u8 {
            for string in ["", "1234", "12345678", "1234567812345678"].into_iter() {
                let mut short = FallbackHasher::new_with_keys(BAD_KEY, BAD_KEY);
                string.hash(&mut short);
                let value = short.finish();
                let mut string = string.to_string();
                for num in 1..=128 {
                    let mut long = FallbackHasher::new_with_keys(BAD_KEY, BAD_KEY);
                    string.push(c as char);
                    string.hash(&mut long);
                    let same_bytes = count_same_bytes(value, long.finish());
                    assert!(same_bytes <= 2, format!("{} bytes of {} -> {:x} vs {:x}", num, c, value, long.finish()));
                    let flipped_bits = (value ^ long.finish()).count_ones();
                    assert!(flipped_bits > 10,
                            format!("Too few bits fliped {} -> {:x} vs {} -> {:x}", string, value, c, long.finish()));
                }
            }
        }
    }
}<commit_msg>Unify tests<commit_after>use std::hash::{Hash, Hasher};

fn assert_sufficiently_different(a: u64, b: u64) {
    let same_byte_count = count_same_bytes(a, b);
    assert!(same_byte_count <= 1);
    let flipped_bits = (a ^ b).count_ones();
    assert!(flipped_bits > 18 && flipped_bits < 48, "{:x} and {:x}: {:}", a, b, flipped_bits);
    for rotate in 0..64 {
        let flipped_bits2 = (a ^ (b.rotate_left(rotate))).count_ones();
        assert!(flipped_bits2 > 14 && flipped_bits2 < 50, "{:x} and {:x}: {:}", a, b.rotate_left(rotate), flipped_bits2);
    }
}

fn count_same_bytes(a: u64, b: u64) -> i32 {
    let mut same_byte_count = 0;
    for byte in 0..8 {
        let ba = a >> 8 * byte as u8;
        let bb = b >> 8 * byte as u8;
        if ba == bb {
            same_byte_count += 1;
        }
    }
    same_byte_count
}

fn test_keys_change_output<T:Hasher>(constructor: impl Fn(u64, u64)->T) {
    let mut a = constructor(0, 0);
    let mut b = constructor(0, 1);
    let mut c = constructor(1, 0);
    let mut d = constructor(1, 1);
    "test".hash(&mut a);
    "test".hash(&mut b);
    "test".hash(&mut c);
    "test".hash(&mut d);
    assert_sufficiently_different(a.finish(), b.finish());
    assert_sufficiently_different(a.finish(), c.finish());
    assert_sufficiently_different(a.finish(), d.finish());
    assert_sufficiently_different(b.finish(), c.finish());
    assert_sufficiently_different(b.finish(), d.finish());
    assert_sufficiently_different(c.finish(), d.finish());
}

fn test_finish_is_consistant<T:Hasher>(constructor: impl Fn(u64, u64)->T) {
    let mut hasher = constructor(1, 2);
    "Foo".hash(&mut hasher);
    let a = hasher.finish();
    let b = hasher.finish();
    assert_eq!(a, b);
}

fn test_single_key_bit_flip<T:Hasher>(constructor: impl Fn(u64, u64)->T) {
    for bit in 0..64 {
        let mut a = constructor(0, 0);
        let mut b = constructor(0, 1 << bit);
        let mut c = constructor(1 << bit, 0);
        "1234".hash(&mut a);
        "1234".hash(&mut b);
        "1234".hash(&mut c);
        assert_sufficiently_different(a.finish(), b.finish());
        assert_sufficiently_different(a.finish(), c.finish());
        assert_sufficiently_different(b.finish(), c.finish());
        let mut a = constructor(0, 0);
        let mut b = constructor(0, 1 << bit);
        let mut c = constructor(1 << bit, 0);
        "12345678".hash(&mut a);
        "12345678".hash(&mut b);
        "12345678".hash(&mut c);
        assert_sufficiently_different(a.finish(), b.finish());
        assert_sufficiently_different(a.finish(), c.finish());
        assert_sufficiently_different(b.finish(), c.finish());
        let mut a = constructor(0, 0);
        let mut b = constructor(0, 1 << bit);
        let mut c = constructor(1 << bit, 0);
        "1234567812345678".hash(&mut a);
        "1234567812345678".hash(&mut b);
        "1234567812345678".hash(&mut c);
        assert_sufficiently_different(a.finish(), b.finish());
        assert_sufficiently_different(a.finish(), c.finish());
        assert_sufficiently_different(b.finish(), c.finish());
    }
}

fn hash<T:Hasher>(b: impl Hash, hasher: &Fn()->T) -> u64 {
    let mut hasher = hasher();
    b.hash(&mut hasher);
    hasher.finish()
}

fn test_single_bit_flip<T:Hasher>(hasher: impl Fn()->T) {
    let size = 32;
    let compare_value = hash(0u32, &hasher);
    for pos in 0..size {
        let test_value = hash(0 ^ (1u32 << pos), &hasher);
        assert_sufficiently_different(compare_value, test_value);
    }
    let size = 64;
    let compare_value = hash(0u64, &hasher);
    for pos in 0..size {
        let test_value = hash(0 ^ (1u64 << pos), &hasher);
        assert_sufficiently_different(compare_value, test_value);
    }
    let size = 128;
    let compare_value = hash(0u128, &hasher);
    for pos in 0..size {
        let test_value = hash(0 ^ (1u128 << pos), &hasher);
        assert_sufficiently_different(compare_value, test_value);
    }
}

fn test_padding_doesnot_collide<T:Hasher>(hasher: impl Fn()->T) {
    for c in 0..128u8 {
        for string in ["", "1234", "12345678", "1234567812345678"].into_iter() {
            let mut short = hasher();
            string.hash(&mut short);
            let value = short.finish();
            let mut string = string.to_string();
            for num in 1..=128 {
                let mut long = hasher();
                string.push(c as char);
                string.hash(&mut long);
                let same_bytes = count_same_bytes(value, long.finish());
                assert!(same_bytes <= 2, format!("{} bytes of {} -> {:x} vs {:x}", num, c, value, long.finish()));
                let flipped_bits = (value ^ long.finish()).count_ones();
                assert!(flipped_bits > 10);
            }
        }
    }
}

#[cfg(test)]
mod fallback_tests {
    use std::hash::{Hash, Hasher};
    use crate::fallback_hash::*;
    use crate::hash_quality_test::*;

    #[test]
    fn fallback_single_bit_flip() {
        test_single_bit_flip(|| FallbackHasher::new_with_keys(0, 0))
    }

    #[test]
    fn fallback_single_key_bit_flip() {
        test_single_key_bit_flip(FallbackHasher::new_with_keys)
    }

    #[test]
    fn fallback_keys_change_output() {
        test_keys_change_output(FallbackHasher::new_with_keys);
    }

    #[test]
    fn fallback_finish_is_consistant() {
        test_finish_is_consistant(FallbackHasher::new_with_keys)
    }

    #[test]
    fn fallback_padding_doesnot_collide() {
        test_padding_doesnot_collide(|| FallbackHasher::new_with_keys(0, 0))
    }
}

///Basic sanity tests of the cypto properties of aHash.
#[cfg(test)]
mod aes_tests {
    use std::hash::{Hash, Hasher};
    use crate::aes_hash::*;
    use crate::hash_quality_test::*;

    const BAD_KEY: u64 = 0x5252_5252_5252_5252; //Thi   s encrypts to 0.

    #[test]
    fn test_single_bit_in_byte() {
        let mut hasher1= AesHasher::new_with_keys(64, 64);
        8_u32.hash(&mut hasher1);
        let mut hasher2= AesHasher::new_with_keys(64, 64);
        0_u32.hash(&mut hasher2);
        assert_sufficiently_different(hasher1.finish(), hasher2.finish());
    }

    #[test]
    fn aes_single_bit_flip() {
        test_single_bit_flip(|k1, k2| AesHasher::new_with_keys(k1,k2))
    }

    #[test]
    fn aes_single_key_bit_flip() {
        test_single_key_bit_flip(|k1, k2| AesHasher::new_with_keys(k1,k2))
    }

    #[test]
    fn aes_keys_change_output() {
        test_keys_change_output(AesHasher::new_with_keys);
    }

    #[test]
    fn aes_finish_is_consistant() {
        test_finish_is_consistant(AesHasher::new_with_keys)
    }

    #[test]
    fn aes_padding_doesnot_collide() {
        test_padding_doesnot_collide(|| AesHasher::new_with_keys(BAD_KEY,BAD_KEY))
    }
}
<|endoftext|>


<commit_before>use serde;

use std::fmt;
use std::io::Write;

use rmp::Marker;
use rmp::encode::{
    write_nil,
    write_bool,
    write_uint,
    write_sint_eff,
    write_f32,
    write_f64,
    write_str,
    write_array_len,
    write_map_len,
    write_bin_len,
    WriteError,
    FixedValueWriteError,
    ValueWriteError,
};

#[derive(Debug)]
pub enum Error {
    /// Failed to write MessagePack'ed single-byte value into the write.
    InvalidFixedValueWrite(WriteError),
    InvalidValueWrite(ValueWriteError),

    /// Failed to serialize struct, sequence or map, because its length is unknown.
    UnknownLength,

    /// Depth limit exceeded
    DepthLimitExceeded,
    Custom(String)
}

impl ::std::error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::InvalidFixedValueWrite(..) => "invalid fixed value write",
            Error::InvalidValueWrite(..) => "invalid value write",
            Error::UnknownLength => "attempt to serialize struct, sequence or map with unknown length",
            Error::DepthLimitExceeded => "depth limit exceeded",
            Error::Custom(..) => "custom message",
        }
    }

    fn cause(&self) -> Option<&::std::error::Error> {
        match *self {
            Error::InvalidFixedValueWrite(ref err) => Some(err),
            Error::InvalidValueWrite(ref err) => Some(err),
            Error::UnknownLength => None,
            Error::DepthLimitExceeded => None,
            Error::Custom(_) => None,
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        ::std::error::Error::description(self).fmt(f)
    }
}

impl From<FixedValueWriteError> for Error {
    fn from(err: FixedValueWriteError) -> Error {
        match err {
            FixedValueWriteError(err) => Error::InvalidFixedValueWrite(err)
        }
    }
}

impl From<ValueWriteError> for Error {
    fn from(err: ValueWriteError) -> Error {
        Error::InvalidValueWrite(err)
    }
}

impl serde::ser::Error for Error {
    /// Raised when there is general error when deserializing a type.
    fn custom<T: Into<String>>(msg: T) -> Error {
        Error::Custom(msg.into())
    }
}

pub trait VariantWriter {
    fn write_struct_len<W>(&self, wr: &mut W, len: u32) -> Result<Marker, ValueWriteError> where W: Write;
    fn write_field_name<W>(&self, wr: &mut W, _key: &str) -> Result<(), ValueWriteError> where W: Write;
}

/// Writes struct as MessagePack array with no field names
pub struct StructArrayWriter;

impl VariantWriter for StructArrayWriter {
    fn write_struct_len<W>(&self, wr: &mut W, len: u32) -> Result<Marker, ValueWriteError>
        where W: Write
    {
        write_array_len(wr, len)
    }

    /// This implementation does not write field names
    #[allow(unused_variables)]
    fn write_field_name<W>(&self, wr: &mut W, _key: &str) -> Result<(), ValueWriteError>
        where W: Write
    {
        Ok(())
    }
}

/// Represents MessagePack serialization implementation.
///
/// # Note
///
/// MessagePack has no specification about how to encode variant types. Thus we are free to do
/// whatever we want, so the given chose may be not ideal for you.
///
/// Every Rust variant value can be represented as a tuple of index and a value.
///
/// All instances of `ErrorKind::Interrupted` are handled by this function and the underlying
/// operation is retried.
// TODO: Docs. Examples.
pub struct Serializer<'a, W: VariantWriter> {
    wr: &'a mut Write,
    vw: W,
    depth: usize,
}

impl<'a, W: VariantWriter> Serializer<'a, W> {
    /// Changes the maximum nesting depth that is allowed
    pub fn set_max_depth(&mut self, depth: usize) {
        self.depth = depth;
    }
}

macro_rules! depth_count(
    ( $counter:expr, $expr:expr ) => {
        {
            $counter -= 1;
            if $counter == 0 {
                return Err(Error::DepthLimitExceeded)
            }
            let res = $expr;
            $counter += 1;
            res
        }
    }
);

impl<'a> Serializer<'a, StructArrayWriter> {
    /// Creates a new MessagePack encoder whose output will be written to the writer specified.
    pub fn new(wr: &'a mut Write) -> Serializer<'a, StructArrayWriter> {
        Serializer {
            wr: wr,
            vw: StructArrayWriter,
            depth: 1000,
        }
    }
}

impl<'a, W: VariantWriter> Serializer<'a, W> {
    /// Creates a new MessagePack encoder whose output will be written to the writer specified.
    pub fn with(wr: &'a mut Write, vw: W) -> Serializer<'a, W> {
        Serializer {
            wr: wr,
            vw: vw,
            depth: 1000,
        }
    }

    #[inline]
    fn serialize_variant<F>(&mut self, variant_index: usize, maybe_len: Option<usize>, mut visit: F) -> Result<(), Error>
        where F: FnMut(&mut Self) -> Result<Option<()>, <Self as serde::Serializer>::Error>
    {
        try!(write_array_len(&mut self.wr, 2));

        // Encode a value position...
        try!(serde::Serializer::serialize_usize(self, variant_index));

        let len = match maybe_len {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        // ... and its arguments length.
        try!(write_array_len(&mut self.wr, len as u32));

        while let Some(()) = try!(depth_count!(self.depth, visit(self))) { }

        Ok(())
    }
}

impl<'a, W: VariantWriter> serde::Serializer for Serializer<'a, W> {
    type Error = Error;

    fn serialize_unit(&mut self) -> Result<(), Error> {
        write_nil(&mut self.wr).map_err(From::from)
    }

    fn serialize_bool(&mut self, val: bool) -> Result<(), Error> {
        write_bool(&mut self.wr, val).map_err(From::from)
    }

    fn serialize_u8(&mut self, val: u8) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u16(&mut self, val: u16) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u32(&mut self, val: u32) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u64(&mut self, val: u64) -> Result<(), Error> {
        try!(write_uint(&mut self.wr, val));

        Ok(())
    }

    fn serialize_usize(&mut self, val: usize) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_i8(&mut self, val: i8) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i16(&mut self, val: i16) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i32(&mut self, val: i32) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i64(&mut self, val: i64) -> Result<(), Error> {
        try!(write_sint_eff(&mut self.wr, val));

        Ok(())
    }

    fn serialize_isize(&mut self, val: isize) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_f32(&mut self, val: f32) -> Result<(), Error> {
        write_f32(&mut self.wr, val).map_err(From::from)
    }

    fn serialize_f64(&mut self, val: f64) -> Result<(), Error> {
        write_f64(&mut self.wr, val).map_err(From::from)
    }

    // TODO: The implementation involves heap allocation and is unstable.
    fn serialize_char(&mut self, val: char) -> Result<(), Error> {
        let mut buf = String::new();
        buf.push(val);
        self.serialize_str(&buf)
    }

    fn serialize_str(&mut self, val: &str) -> Result<(), Error> {
        write_str(&mut self.wr, val).map_err(From::from)
    }

    fn serialize_unit_variant(&mut self,
                          _name: &str,
                          variant_index: usize,
                          _variant: &str) -> Result<(), Error>
    {
        // Mark that we want to encode a variant type.
        try!(write_array_len(&mut self.wr, 2));

        // Encode a value position...
        try!(self.serialize_usize(variant_index));

        // ... and its arguments length.
        try!(write_array_len(&mut self.wr, 0));

        Ok(())
    }

    /// Encodes and attempts to write the enum value into the Write.
    ///
    /// Currently we encode variant types as a tuple of id with array of args, like: [id, [args...]]
    fn serialize_tuple_variant<V>(&mut self,
                              _name: &str,
                              variant_index: usize,
                              _variant: &str,
                              mut visitor: V) -> Result<(), Error>
        where V: serde::ser::SeqVisitor,
    {
        self.serialize_variant(variant_index, visitor.len(), |s| visitor.visit(s))
    }

    fn serialize_struct_variant<V>(&mut self,
                               _name: &str,
                               variant_index: usize,
                               _variant: &str,
                               mut visitor: V) -> Result<(), Error>
        where V: serde::ser::MapVisitor,
    {
        self.serialize_variant(variant_index, visitor.len(), |s| visitor.visit(s))
    }

    fn serialize_none(&mut self) -> Result<(), Error> {
        self.serialize_unit()
    }

    fn serialize_some<T>(&mut self, v: T) -> Result<(), Error>
        where T: serde::Serialize,
    {
        depth_count!(self.depth, v.serialize(self))
    }

    // TODO: Check len, overflow is possible.
    fn serialize_seq<V>(&mut self, mut visitor: V) -> Result<(), Error>
        where V: serde::ser::SeqVisitor,
    {
        let len = match visitor.len() {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        try!(write_array_len(&mut self.wr, len as u32));

        while let Some(()) = try!(depth_count!(self.depth, visitor.visit(self))) { }

        Ok(())
    }

    fn serialize_seq_elt<V>(&mut self, value: V) -> Result<(), Error>
        where V: serde::Serialize,
    {
        value.serialize(self)
    }

    fn serialize_map<V>(&mut self, mut visitor: V) -> Result<(), Error>
        where V: serde::ser::MapVisitor,
    {
        let len = match visitor.len() {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        try!(write_map_len(&mut self.wr, len as u32));

        while let Some(()) = try!(depth_count!(self.depth, visitor.visit(self))) { }

        Ok(())
    }

    fn serialize_map_elt<K, V>(&mut self, key: K, value: V) -> Result<(), Error>
        where K: serde::Serialize,
              V: serde::Serialize,
    {
        try!(key.serialize(self));
        value.serialize(self)
    }

    fn serialize_unit_struct(&mut self, _name: &'static str) -> Result<(), Error> {
        try!(write_array_len(&mut self.wr, 0));

        Ok(())
    }

    fn serialize_struct<V>(&mut self, _name: &str, mut visitor: V) -> Result<(), Error>
        where V: serde::ser::MapVisitor,
    {
        let len = match visitor.len() {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        try!(self.vw.write_struct_len(&mut self.wr, len as u32));

        while let Some(()) = try!(depth_count!(self.depth, visitor.visit(self))) { }

        Ok(())
    }

    fn serialize_struct_elt<V>(&mut self, _key: &str, value: V) -> Result<(), Error>
        where V: serde::Serialize,
    {
        try!(self.vw.write_field_name(&mut self.wr, _key));
        value.serialize(self)
    }

    fn serialize_bytes(&mut self, value: &[u8]) -> Result<(), Error> {
        try!(write_bin_len(&mut self.wr, value.len() as u32));
        self.wr.write_all(value).map_err(|err| Error::InvalidValueWrite(ValueWriteError::InvalidDataWrite(WriteError(err))))
    }
}
<commit_msg>fix: switch unit structs to using the same serialization mechanism as other structs (#76)<commit_after>use serde;

use std::fmt;
use std::io::Write;

use rmp::Marker;
use rmp::encode::{
    write_nil,
    write_bool,
    write_uint,
    write_sint_eff,
    write_f32,
    write_f64,
    write_str,
    write_array_len,
    write_map_len,
    write_bin_len,
    WriteError,
    FixedValueWriteError,
    ValueWriteError,
};

#[derive(Debug)]
pub enum Error {
    /// Failed to write MessagePack'ed single-byte value into the write.
    InvalidFixedValueWrite(WriteError),
    InvalidValueWrite(ValueWriteError),

    /// Failed to serialize struct, sequence or map, because its length is unknown.
    UnknownLength,

    /// Depth limit exceeded
    DepthLimitExceeded,
    Custom(String)
}

impl ::std::error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::InvalidFixedValueWrite(..) => "invalid fixed value write",
            Error::InvalidValueWrite(..) => "invalid value write",
            Error::UnknownLength => "attempt to serialize struct, sequence or map with unknown length",
            Error::DepthLimitExceeded => "depth limit exceeded",
            Error::Custom(..) => "custom message",
        }
    }

    fn cause(&self) -> Option<&::std::error::Error> {
        match *self {
            Error::InvalidFixedValueWrite(ref err) => Some(err),
            Error::InvalidValueWrite(ref err) => Some(err),
            Error::UnknownLength => None,
            Error::DepthLimitExceeded => None,
            Error::Custom(_) => None,
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        ::std::error::Error::description(self).fmt(f)
    }
}

impl From<FixedValueWriteError> for Error {
    fn from(err: FixedValueWriteError) -> Error {
        match err {
            FixedValueWriteError(err) => Error::InvalidFixedValueWrite(err)
        }
    }
}

impl From<ValueWriteError> for Error {
    fn from(err: ValueWriteError) -> Error {
        Error::InvalidValueWrite(err)
    }
}

impl serde::ser::Error for Error {
    /// Raised when there is general error when deserializing a type.
    fn custom<T: Into<String>>(msg: T) -> Error {
        Error::Custom(msg.into())
    }
}

pub trait VariantWriter {
    fn write_struct_len<W>(&self, wr: &mut W, len: u32) -> Result<Marker, ValueWriteError> where W: Write;
    fn write_field_name<W>(&self, wr: &mut W, _key: &str) -> Result<(), ValueWriteError> where W: Write;
}

/// Writes struct as MessagePack array with no field names
pub struct StructArrayWriter;

impl VariantWriter for StructArrayWriter {
    fn write_struct_len<W>(&self, wr: &mut W, len: u32) -> Result<Marker, ValueWriteError>
        where W: Write
    {
        write_array_len(wr, len)
    }

    /// This implementation does not write field names
    #[allow(unused_variables)]
    fn write_field_name<W>(&self, wr: &mut W, _key: &str) -> Result<(), ValueWriteError>
        where W: Write
    {
        Ok(())
    }
}

/// Represents MessagePack serialization implementation.
///
/// # Note
///
/// MessagePack has no specification about how to encode variant types. Thus we are free to do
/// whatever we want, so the given chose may be not ideal for you.
///
/// Every Rust variant value can be represented as a tuple of index and a value.
///
/// All instances of `ErrorKind::Interrupted` are handled by this function and the underlying
/// operation is retried.
// TODO: Docs. Examples.
pub struct Serializer<'a, W: VariantWriter> {
    wr: &'a mut Write,
    vw: W,
    depth: usize,
}

impl<'a, W: VariantWriter> Serializer<'a, W> {
    /// Changes the maximum nesting depth that is allowed
    pub fn set_max_depth(&mut self, depth: usize) {
        self.depth = depth;
    }
}

macro_rules! depth_count(
    ( $counter:expr, $expr:expr ) => {
        {
            $counter -= 1;
            if $counter == 0 {
                return Err(Error::DepthLimitExceeded)
            }
            let res = $expr;
            $counter += 1;
            res
        }
    }
);

impl<'a> Serializer<'a, StructArrayWriter> {
    /// Creates a new MessagePack encoder whose output will be written to the writer specified.
    pub fn new(wr: &'a mut Write) -> Serializer<'a, StructArrayWriter> {
        Serializer {
            wr: wr,
            vw: StructArrayWriter,
            depth: 1000,
        }
    }
}

impl<'a, W: VariantWriter> Serializer<'a, W> {
    /// Creates a new MessagePack encoder whose output will be written to the writer specified.
    pub fn with(wr: &'a mut Write, vw: W) -> Serializer<'a, W> {
        Serializer {
            wr: wr,
            vw: vw,
            depth: 1000,
        }
    }

    #[inline]
    fn serialize_variant<F>(&mut self, variant_index: usize, maybe_len: Option<usize>, mut visit: F) -> Result<(), Error>
        where F: FnMut(&mut Self) -> Result<Option<()>, <Self as serde::Serializer>::Error>
    {
        try!(write_array_len(&mut self.wr, 2));

        // Encode a value position...
        try!(serde::Serializer::serialize_usize(self, variant_index));

        let len = match maybe_len {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        // ... and its arguments length.
        try!(write_array_len(&mut self.wr, len as u32));

        while let Some(()) = try!(depth_count!(self.depth, visit(self))) { }

        Ok(())
    }
}

impl<'a, W: VariantWriter> serde::Serializer for Serializer<'a, W> {
    type Error = Error;

    fn serialize_unit(&mut self) -> Result<(), Error> {
        write_nil(&mut self.wr).map_err(From::from)
    }

    fn serialize_bool(&mut self, val: bool) -> Result<(), Error> {
        write_bool(&mut self.wr, val).map_err(From::from)
    }

    fn serialize_u8(&mut self, val: u8) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u16(&mut self, val: u16) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u32(&mut self, val: u32) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_u64(&mut self, val: u64) -> Result<(), Error> {
        try!(write_uint(&mut self.wr, val));

        Ok(())
    }

    fn serialize_usize(&mut self, val: usize) -> Result<(), Error> {
        self.serialize_u64(val as u64)
    }

    fn serialize_i8(&mut self, val: i8) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i16(&mut self, val: i16) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i32(&mut self, val: i32) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_i64(&mut self, val: i64) -> Result<(), Error> {
        try!(write_sint_eff(&mut self.wr, val));

        Ok(())
    }

    fn serialize_isize(&mut self, val: isize) -> Result<(), Error> {
        self.serialize_i64(val as i64)
    }

    fn serialize_f32(&mut self, val: f32) -> Result<(), Error> {
        write_f32(&mut self.wr, val).map_err(From::from)
    }

    fn serialize_f64(&mut self, val: f64) -> Result<(), Error> {
        write_f64(&mut self.wr, val).map_err(From::from)
    }

    // TODO: The implementation involves heap allocation and is unstable.
    fn serialize_char(&mut self, val: char) -> Result<(), Error> {
        let mut buf = String::new();
        buf.push(val);
        self.serialize_str(&buf)
    }

    fn serialize_str(&mut self, val: &str) -> Result<(), Error> {
        write_str(&mut self.wr, val).map_err(From::from)
    }

    fn serialize_unit_variant(&mut self,
                          _name: &str,
                          variant_index: usize,
                          _variant: &str) -> Result<(), Error>
    {
        // Mark that we want to encode a variant type.
        try!(write_array_len(&mut self.wr, 2));

        // Encode a value position...
        try!(self.serialize_usize(variant_index));

        // ... and its arguments length.
        try!(write_array_len(&mut self.wr, 0));

        Ok(())
    }

    /// Encodes and attempts to write the enum value into the Write.
    ///
    /// Currently we encode variant types as a tuple of id with array of args, like: [id, [args...]]
    fn serialize_tuple_variant<V>(&mut self,
                              _name: &str,
                              variant_index: usize,
                              _variant: &str,
                              mut visitor: V) -> Result<(), Error>
        where V: serde::ser::SeqVisitor,
    {
        self.serialize_variant(variant_index, visitor.len(), |s| visitor.visit(s))
    }

    fn serialize_struct_variant<V>(&mut self,
                               _name: &str,
                               variant_index: usize,
                               _variant: &str,
                               mut visitor: V) -> Result<(), Error>
        where V: serde::ser::MapVisitor,
    {
        self.serialize_variant(variant_index, visitor.len(), |s| visitor.visit(s))
    }

    fn serialize_none(&mut self) -> Result<(), Error> {
        self.serialize_unit()
    }

    fn serialize_some<T>(&mut self, v: T) -> Result<(), Error>
        where T: serde::Serialize,
    {
        depth_count!(self.depth, v.serialize(self))
    }

    // TODO: Check len, overflow is possible.
    fn serialize_seq<V>(&mut self, mut visitor: V) -> Result<(), Error>
        where V: serde::ser::SeqVisitor,
    {
        let len = match visitor.len() {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        try!(write_array_len(&mut self.wr, len as u32));

        while let Some(()) = try!(depth_count!(self.depth, visitor.visit(self))) { }

        Ok(())
    }

    fn serialize_seq_elt<V>(&mut self, value: V) -> Result<(), Error>
        where V: serde::Serialize,
    {
        value.serialize(self)
    }

    fn serialize_map<V>(&mut self, mut visitor: V) -> Result<(), Error>
        where V: serde::ser::MapVisitor,
    {
        let len = match visitor.len() {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        try!(write_map_len(&mut self.wr, len as u32));

        while let Some(()) = try!(depth_count!(self.depth, visitor.visit(self))) { }

        Ok(())
    }

    fn serialize_map_elt<K, V>(&mut self, key: K, value: V) -> Result<(), Error>
        where K: serde::Serialize,
              V: serde::Serialize,
    {
        try!(key.serialize(self));
        value.serialize(self)
    }

    fn serialize_unit_struct(&mut self, _name: &'static str) -> Result<(), Error> {
        try!(self.vw.write_struct_len(&mut self.wr, 0));

        Ok(())
    }

    fn serialize_struct<V>(&mut self, _name: &str, mut visitor: V) -> Result<(), Error>
        where V: serde::ser::MapVisitor,
    {
        let len = match visitor.len() {
            Some(len) => len,
            None => return Err(Error::UnknownLength),
        };

        try!(self.vw.write_struct_len(&mut self.wr, len as u32));

        while let Some(()) = try!(depth_count!(self.depth, visitor.visit(self))) { }

        Ok(())
    }

    fn serialize_struct_elt<V>(&mut self, _key: &str, value: V) -> Result<(), Error>
        where V: serde::Serialize,
    {
        try!(self.vw.write_field_name(&mut self.wr, _key));
        value.serialize(self)
    }

    fn serialize_bytes(&mut self, value: &[u8]) -> Result<(), Error> {
        try!(write_bin_len(&mut self.wr, value.len() as u32));
        self.wr.write_all(value).map_err(|err| Error::InvalidValueWrite(ValueWriteError::InvalidDataWrite(WriteError(err))))
    }
}
<|endoftext|>


<commit_before>use crate::{Env, Error, NifResult, Term};

#[macro_use]
pub mod atom;
pub use crate::types::atom::Atom;

#[doc(hidden)]
pub mod binary;
pub use crate::types::binary::{Binary, OwnedBinary};

#[doc(hidden)]
pub mod list;
pub use crate::types::list::ListIterator;

#[doc(hidden)]
pub mod map;
pub use self::map::MapIterator;

#[doc(hidden)]
pub mod primitive;
#[doc(hidden)]
pub mod string;
pub mod tuple;

#[doc(hidden)]
pub mod local_pid;
pub use self::local_pid::LocalPid;

#[deprecated(since = "0.22.0", note = "Please use local_pid instead")]
pub mod pid {
    #[deprecated(since = "0.22.0", note = "Please use LocalPid instead")]
    pub use super::LocalPid as Pid;
}
#[deprecated(since = "0.22.0", note = "Please use LocalPid instead")]
pub use self::LocalPid as Pid;

pub mod truthy;

pub mod elixir_struct;

pub trait Encoder {
    fn encode<'a>(&self, env: Env<'a>) -> Term<'a>;
}
pub trait Decoder<'a>: Sized + 'a {
    fn decode(term: Term<'a>) -> NifResult<Self>;
}

impl<'a> Encoder for Term<'a> {
    fn encode<'b>(&self, env: Env<'b>) -> Term<'b> {
        self.in_env(env)
    }
}
impl<'a> Decoder<'a> for Term<'a> {
    fn decode(term: Term<'a>) -> NifResult<Self> {
        Ok(term)
    }
}

impl<'a, T> Encoder for &'a T
where
    T: Encoder,
{
    fn encode<'c>(&self, env: Env<'c>) -> Term<'c> {
        <T as Encoder>::encode(self, env)
    }
}

impl<T> Encoder for Option<T>
where
    T: Encoder,
{
    fn encode<'c>(&self, env: Env<'c>) -> Term<'c> {
        match *self {
            Some(ref value) => value.encode(env),
            None => atom::nil().encode(env),
        }
    }
}

impl<'a, T> Decoder<'a> for Option<T>
where
    T: Decoder<'a>,
{
    fn decode(term: Term<'a>) -> NifResult<Self> {
        if let Ok(term) = term.decode::<T>() {
            Ok(Some(term))
        } else {
            let decoded_atom: atom::Atom = term.decode()?;
            if decoded_atom == atom::nil() {
                Ok(None)
            } else {
                Err(Error::BadArg)
            }
        }
    }
}

impl<T, E> Encoder for Result<T, E>
where
    T: Encoder,
    E: Encoder,
{
    fn encode<'c>(&self, env: Env<'c>) -> Term<'c> {
        match *self {
            Ok(ref value) => (atom::ok().encode(env), value.encode(env)).encode(env),
            Err(ref err) => (atom::error().encode(env), err.encode(env)).encode(env),
        }
    }
}

impl<'a, T, E> Decoder<'a> for Result<T, E>
where
    T: Decoder<'a>,
    E: Decoder<'a>,
{
    fn decode(term: Term<'a>) -> NifResult<Self> {
        let (decoded_atom, inner_term): (atom::Atom, Term) = term.decode()?;
        if decoded_atom == atom::ok() {
            let ok_value: T = inner_term.decode()?;
            Ok(Ok(ok_value))
        } else if decoded_atom == atom::error() {
            let err_value: E = inner_term.decode()?;
            Ok(Err(err_value))
        } else {
            Err(Error::BadArg)
        }
    }
}

impl<'a, K, V> Decoder<'a> for std::collections::HashMap<K, V>
where
    K: Decoder<'a> + Eq + std::hash::Hash,
    V: Decoder<'a>,
{
    fn decode(term: Term<'a>) -> NifResult<Self> {
        let size = term.map_size()?;

        let it = MapIterator::new(term).ok_or_else(|| Error::BadArg)?;

        let mut map = std::collections::HashMap::with_capacity(size);

        for (k, v) in it {
            let k = k.decode()?;
            let v = v.decode()?;
            map.insert(k, v);
        }

        Ok(map)
    }
}

impl<K, V> Encoder for std::collections::HashMap<K, V>
where
    K: Encoder + Eq + std::hash::Hash,
    V: Encoder,
{
    fn encode<'c>(&self, env: Env<'c>) -> Term<'c> {
        let (keys, values): (Vec<_>, Vec<_>) = self
            .iter()
            .map(|(k, v)| (k.encode(env), v.encode(env)))
            .unzip();
        Term::map_from_arrays(env, &keys, &values).unwrap()
    }
}
<commit_msg>Fix clippy lint<commit_after>use crate::{Env, Error, NifResult, Term};

#[macro_use]
pub mod atom;
pub use crate::types::atom::Atom;

#[doc(hidden)]
pub mod binary;
pub use crate::types::binary::{Binary, OwnedBinary};

#[doc(hidden)]
pub mod list;
pub use crate::types::list::ListIterator;

#[doc(hidden)]
pub mod map;
pub use self::map::MapIterator;

#[doc(hidden)]
pub mod primitive;
#[doc(hidden)]
pub mod string;
pub mod tuple;

#[doc(hidden)]
pub mod local_pid;
pub use self::local_pid::LocalPid;

#[deprecated(since = "0.22.0", note = "Please use local_pid instead")]
pub mod pid {
    #[deprecated(since = "0.22.0", note = "Please use LocalPid instead")]
    pub use super::LocalPid as Pid;
}
#[deprecated(since = "0.22.0", note = "Please use LocalPid instead")]
pub use self::LocalPid as Pid;

pub mod truthy;

pub mod elixir_struct;

pub trait Encoder {
    fn encode<'a>(&self, env: Env<'a>) -> Term<'a>;
}
pub trait Decoder<'a>: Sized + 'a {
    fn decode(term: Term<'a>) -> NifResult<Self>;
}

impl<'a> Encoder for Term<'a> {
    fn encode<'b>(&self, env: Env<'b>) -> Term<'b> {
        self.in_env(env)
    }
}
impl<'a> Decoder<'a> for Term<'a> {
    fn decode(term: Term<'a>) -> NifResult<Self> {
        Ok(term)
    }
}

impl<'a, T> Encoder for &'a T
where
    T: Encoder,
{
    fn encode<'c>(&self, env: Env<'c>) -> Term<'c> {
        <T as Encoder>::encode(self, env)
    }
}

impl<T> Encoder for Option<T>
where
    T: Encoder,
{
    fn encode<'c>(&self, env: Env<'c>) -> Term<'c> {
        match *self {
            Some(ref value) => value.encode(env),
            None => atom::nil().encode(env),
        }
    }
}

impl<'a, T> Decoder<'a> for Option<T>
where
    T: Decoder<'a>,
{
    fn decode(term: Term<'a>) -> NifResult<Self> {
        if let Ok(term) = term.decode::<T>() {
            Ok(Some(term))
        } else {
            let decoded_atom: atom::Atom = term.decode()?;
            if decoded_atom == atom::nil() {
                Ok(None)
            } else {
                Err(Error::BadArg)
            }
        }
    }
}

impl<T, E> Encoder for Result<T, E>
where
    T: Encoder,
    E: Encoder,
{
    fn encode<'c>(&self, env: Env<'c>) -> Term<'c> {
        match *self {
            Ok(ref value) => (atom::ok().encode(env), value.encode(env)).encode(env),
            Err(ref err) => (atom::error().encode(env), err.encode(env)).encode(env),
        }
    }
}

impl<'a, T, E> Decoder<'a> for Result<T, E>
where
    T: Decoder<'a>,
    E: Decoder<'a>,
{
    fn decode(term: Term<'a>) -> NifResult<Self> {
        let (decoded_atom, inner_term): (atom::Atom, Term) = term.decode()?;
        if decoded_atom == atom::ok() {
            let ok_value: T = inner_term.decode()?;
            Ok(Ok(ok_value))
        } else if decoded_atom == atom::error() {
            let err_value: E = inner_term.decode()?;
            Ok(Err(err_value))
        } else {
            Err(Error::BadArg)
        }
    }
}

impl<'a, K, V> Decoder<'a> for std::collections::HashMap<K, V>
where
    K: Decoder<'a> + Eq + std::hash::Hash,
    V: Decoder<'a>,
{
    fn decode(term: Term<'a>) -> NifResult<Self> {
        let size = term.map_size()?;

        let it = MapIterator::new(term).ok_or(Error::BadArg)?;

        let mut map = std::collections::HashMap::with_capacity(size);

        for (k, v) in it {
            let k = k.decode()?;
            let v = v.decode()?;
            map.insert(k, v);
        }

        Ok(map)
    }
}

impl<K, V> Encoder for std::collections::HashMap<K, V>
where
    K: Encoder + Eq + std::hash::Hash,
    V: Encoder,
{
    fn encode<'c>(&self, env: Env<'c>) -> Term<'c> {
        let (keys, values): (Vec<_>, Vec<_>) = self
            .iter()
            .map(|(k, v)| (k.encode(env), v.encode(env)))
            .unzip();
        Term::map_from_arrays(env, &keys, &values).unwrap()
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! TCP Transport
use crate::transport::Transport;
use futures::{
    compat::{Compat01As03, Future01CompatExt},
    future::{self, Future},
    io::{AsyncRead, AsyncWrite},
    ready,
    stream::Stream,
};
use parity_multiaddr::{Multiaddr, Protocol};
use std::{
    io,
    net::{Shutdown, SocketAddr},
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};
use tokio::net::tcp::{ConnectFuture, Incoming, TcpListener, TcpStream};

/// Transport to build TCP connections
#[derive(Debug, Clone, Default)]
pub struct TcpTransport {
    /// Size of the recv buffer size to set for opened sockets, or `None` to keep default.
    recv_buffer_size: Option<usize>,
    /// Size of the send buffer size to set for opened sockets, or `None` to keep default.
    send_buffer_size: Option<usize>,
    /// TTL to set for opened sockets, or `None` to keep default.
    ttl: Option<u32>,
    /// Keep alive duration to set for opened sockets, or `None` to keep default.
    #[allow(clippy::option_option)]
    keepalive: Option<Option<Duration>>,
    /// `TCP_NODELAY` to set for opened sockets, or `None` to keep default.
    nodelay: Option<bool>,
}

impl TcpTransport {
    pub fn set_recv_buffer_size(mut self, size: usize) -> Self {
        self.recv_buffer_size = Some(size);
        self
    }

    pub fn set_send_buffer_size(mut self, size: usize) -> Self {
        self.send_buffer_size = Some(size);
        self
    }

    pub fn set_ttl(mut self, ttl: u32) -> Self {
        self.ttl = Some(ttl);
        self
    }

    pub fn set_keepalive(mut self, keepalive: Option<Duration>) -> Self {
        self.keepalive = Some(keepalive);
        self
    }

    pub fn set_nodelay(mut self, nodelay: bool) -> Self {
        self.nodelay = Some(nodelay);
        self
    }

    fn apply_config(&self, stream: &TcpStream) -> ::std::io::Result<()> {
        if let Some(size) = self.recv_buffer_size {
            stream.set_recv_buffer_size(size)?;
        }

        if let Some(size) = self.send_buffer_size {
            stream.set_send_buffer_size(size)?;
        }

        if let Some(ttl) = self.ttl {
            stream.set_ttl(ttl)?;
        }

        if let Some(keepalive) = self.keepalive {
            stream.set_keepalive(keepalive)?;
        }

        if let Some(nodelay) = self.nodelay {
            stream.set_nodelay(nodelay)?;
        }

        Ok(())
    }
}

impl Transport for TcpTransport {
    type Output = TcpSocket;
    type Error = ::std::io::Error;
    type Listener = TcpListenerStream;
    type Inbound = future::Ready<io::Result<TcpSocket>>;
    type Outbound = TcpOutbound;

    fn listen_on(&self, addr: Multiaddr) -> Result<(Self::Listener, Multiaddr), Self::Error> {
        let socket_addr = multiaddr_to_socketaddr(&addr)?;
        let config = self.clone();
        let listener = TcpListener::bind(&socket_addr)?;
        let local_addr = socketaddr_to_multiaddr(listener.local_addr()?);
        Ok((
            TcpListenerStream {
                inner: Compat01As03::new(listener.incoming()),
                config,
            },
            local_addr,
        ))
    }

    fn dial(&self, addr: Multiaddr) -> Result<Self::Outbound, Self::Error> {
        let socket_addr = multiaddr_to_socketaddr(&addr)?;
        let config = self.clone();
        let f = TcpStream::connect(&socket_addr).compat();
        Ok(TcpOutbound { inner: f, config })
    }
}

#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct TcpListenerStream {
    inner: Compat01As03<Incoming>,
    config: TcpTransport,
}

impl Stream for TcpListenerStream {
    type Item = io::Result<(future::Ready<io::Result<TcpSocket>>, Multiaddr)>;

    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context) -> Poll<Option<Self::Item>> {
        match Pin::new(&mut self.inner).poll_next(context) {
            Poll::Ready(Some(Ok(socket))) => {
                if let Err(e) = self.config.apply_config(&socket) {
                    return Poll::Ready(Some(Err(e)));
                }
                let dialer_addr = match socket.peer_addr() {
                    Ok(addr) => socketaddr_to_multiaddr(addr),
                    Err(e) => return Poll::Ready(Some(Err(e))),
                };
                Poll::Ready(Some(Ok((
                    future::ready(Ok(TcpSocket::new(socket))),
                    dialer_addr,
                ))))
            }
            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => Poll::Pending,
        }
    }
}

#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct TcpOutbound {
    inner: Compat01As03<ConnectFuture>,
    config: TcpTransport,
}

impl Future for TcpOutbound {
    type Output = Result<TcpSocket, ::std::io::Error>;

    fn poll(mut self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
        let socket = ready!(Pin::new(&mut self.inner).poll(context))?;
        self.config.apply_config(&socket)?;
        Poll::Ready(Ok(TcpSocket::new(socket)))
    }
}

/// A wrapper around a tokio TcpStream
///
/// In order to properly implement the AsyncRead/AsyncWrite traits we need to wrap a TcpStream to
/// ensure that the "close" method actually closes the write half of the TcpStream.  This is
/// because the "close" method on a TcpStream just performs a no-op instead of actually shutting
/// down the write side of the TcpStream.
//TODO Probably should add some tests for this
#[derive(Debug)]
pub struct TcpSocket {
    inner: Compat01As03<TcpStream>,
}

impl TcpSocket {
    fn new(socket: TcpStream) -> Self {
        Self {
            inner: Compat01As03::new(socket),
        }
    }
}

impl AsyncRead for TcpSocket {
    fn poll_read(
        mut self: Pin<&mut Self>,
        context: &mut Context,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.inner).poll_read(context, buf)
    }
}

impl AsyncWrite for TcpSocket {
    fn poll_write(
        mut self: Pin<&mut Self>,
        context: &mut Context,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.inner).poll_write(context, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_flush(context)
    }

    fn poll_close(self: Pin<&mut Self>, _context: &mut Context) -> Poll<io::Result<()>> {
        Poll::Ready(self.inner.get_ref().shutdown(Shutdown::Write))
    }
}

fn socketaddr_to_multiaddr(socketaddr: SocketAddr) -> Multiaddr {
    let ipaddr: Multiaddr = socketaddr.ip().into();
    ipaddr.with(Protocol::Tcp(socketaddr.port()))
}

fn multiaddr_to_socketaddr(addr: &Multiaddr) -> ::std::io::Result<SocketAddr> {
    let mut iter = addr.iter();
    let proto1 = iter.next().ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("Invalid Multiaddr '{:?}'", addr),
        )
    })?;
    let proto2 = iter.next().ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("Invalid Multiaddr '{:?}'", addr),
        )
    })?;

    if iter.next().is_some() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("Invalid Multiaddr '{:?}'", addr),
        ));
    }

    match (proto1, proto2) {
        (Protocol::Ip4(ip), Protocol::Tcp(port)) => Ok(SocketAddr::new(ip.into(), port)),
        (Protocol::Ip6(ip), Protocol::Tcp(port)) => Ok(SocketAddr::new(ip.into(), port)),
        _ => Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("Invalid Multiaddr '{:?}'", addr),
        )),
    }
}

#[cfg(test)]
mod test {
    use crate::transport::{tcp::TcpTransport, ConnectionOrigin, Transport, TransportExt};
    use futures::{
        executor::block_on,
        future::{join, FutureExt},
        io::{AsyncReadExt, AsyncWriteExt},
        stream::StreamExt,
    };

    #[test]
    fn simple_listen_and_dial() -> Result<(), ::std::io::Error> {
        let t = TcpTransport::default().and_then(|mut out, connection| {
            async move {
                match connection {
                    ConnectionOrigin::Inbound => {
                        out.write_all(b"Earth").await?;
                        let mut buf = [0; 3];
                        out.read_exact(&mut buf).await?;
                        assert_eq!(&buf, b"Air");
                    }
                    ConnectionOrigin::Outbound => {
                        let mut buf = [0; 5];
                        out.read_exact(&mut buf).await?;
                        assert_eq!(&buf, b"Earth");
                        out.write_all(b"Air").await?;
                    }
                }
                Ok(())
            }
        });

        let (listener, addr) = t.listen_on("/ip4/127.0.0.1/tcp/0".parse().unwrap())?;

        let dial = t.dial(addr)?;
        let listener = listener.into_future().then(|(maybe_result, _stream)| {
            let (incoming, _addr) = maybe_result.unwrap().unwrap();
            incoming.map(Result::unwrap)
        });

        let (outgoing, _incoming) = block_on(join(dial, listener));
        assert!(outgoing.is_ok());
        Ok(())
    }

    #[test]
    fn unsupported_multiaddrs() {
        let t = TcpTransport::default();

        let result = t.listen_on("/memory/0".parse().unwrap());
        assert!(result.is_err());

        let result = t.dial("/memory/22".parse().unwrap());
        assert!(result.is_err());
    }
}
<commit_msg>Removes unused setters from netcore::tcp::transport<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! TCP Transport
use crate::transport::Transport;
use futures::{
    compat::{Compat01As03, Future01CompatExt},
    future::{self, Future},
    io::{AsyncRead, AsyncWrite},
    ready,
    stream::Stream,
};
use parity_multiaddr::{Multiaddr, Protocol};
use std::{
    io,
    net::{Shutdown, SocketAddr},
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};
use tokio::net::tcp::{ConnectFuture, Incoming, TcpListener, TcpStream};

/// Transport to build TCP connections
#[derive(Debug, Clone, Default)]
pub struct TcpTransport {
    /// Size of the recv buffer size to set for opened sockets, or `None` to keep default.
    recv_buffer_size: Option<usize>,
    /// Size of the send buffer size to set for opened sockets, or `None` to keep default.
    send_buffer_size: Option<usize>,
    /// TTL to set for opened sockets, or `None` to keep default.
    ttl: Option<u32>,
    /// Keep alive duration to set for opened sockets, or `None` to keep default.
    #[allow(clippy::option_option)]
    keepalive: Option<Option<Duration>>,
    /// `TCP_NODELAY` to set for opened sockets, or `None` to keep default.
    nodelay: Option<bool>,
}

impl TcpTransport {
    fn apply_config(&self, stream: &TcpStream) -> ::std::io::Result<()> {
        if let Some(size) = self.recv_buffer_size {
            stream.set_recv_buffer_size(size)?;
        }

        if let Some(size) = self.send_buffer_size {
            stream.set_send_buffer_size(size)?;
        }

        if let Some(ttl) = self.ttl {
            stream.set_ttl(ttl)?;
        }

        if let Some(keepalive) = self.keepalive {
            stream.set_keepalive(keepalive)?;
        }

        if let Some(nodelay) = self.nodelay {
            stream.set_nodelay(nodelay)?;
        }

        Ok(())
    }
}

impl Transport for TcpTransport {
    type Output = TcpSocket;
    type Error = ::std::io::Error;
    type Listener = TcpListenerStream;
    type Inbound = future::Ready<io::Result<TcpSocket>>;
    type Outbound = TcpOutbound;

    fn listen_on(&self, addr: Multiaddr) -> Result<(Self::Listener, Multiaddr), Self::Error> {
        let socket_addr = multiaddr_to_socketaddr(&addr)?;
        let config = self.clone();
        let listener = TcpListener::bind(&socket_addr)?;
        let local_addr = socketaddr_to_multiaddr(listener.local_addr()?);
        Ok((
            TcpListenerStream {
                inner: Compat01As03::new(listener.incoming()),
                config,
            },
            local_addr,
        ))
    }

    fn dial(&self, addr: Multiaddr) -> Result<Self::Outbound, Self::Error> {
        let socket_addr = multiaddr_to_socketaddr(&addr)?;
        let config = self.clone();
        let f = TcpStream::connect(&socket_addr).compat();
        Ok(TcpOutbound { inner: f, config })
    }
}

#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct TcpListenerStream {
    inner: Compat01As03<Incoming>,
    config: TcpTransport,
}

impl Stream for TcpListenerStream {
    type Item = io::Result<(future::Ready<io::Result<TcpSocket>>, Multiaddr)>;

    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context) -> Poll<Option<Self::Item>> {
        match Pin::new(&mut self.inner).poll_next(context) {
            Poll::Ready(Some(Ok(socket))) => {
                if let Err(e) = self.config.apply_config(&socket) {
                    return Poll::Ready(Some(Err(e)));
                }
                let dialer_addr = match socket.peer_addr() {
                    Ok(addr) => socketaddr_to_multiaddr(addr),
                    Err(e) => return Poll::Ready(Some(Err(e))),
                };
                Poll::Ready(Some(Ok((
                    future::ready(Ok(TcpSocket::new(socket))),
                    dialer_addr,
                ))))
            }
            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => Poll::Pending,
        }
    }
}

#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct TcpOutbound {
    inner: Compat01As03<ConnectFuture>,
    config: TcpTransport,
}

impl Future for TcpOutbound {
    type Output = Result<TcpSocket, ::std::io::Error>;

    fn poll(mut self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
        let socket = ready!(Pin::new(&mut self.inner).poll(context))?;
        self.config.apply_config(&socket)?;
        Poll::Ready(Ok(TcpSocket::new(socket)))
    }
}

/// A wrapper around a tokio TcpStream
///
/// In order to properly implement the AsyncRead/AsyncWrite traits we need to wrap a TcpStream to
/// ensure that the "close" method actually closes the write half of the TcpStream.  This is
/// because the "close" method on a TcpStream just performs a no-op instead of actually shutting
/// down the write side of the TcpStream.
//TODO Probably should add some tests for this
#[derive(Debug)]
pub struct TcpSocket {
    inner: Compat01As03<TcpStream>,
}

impl TcpSocket {
    fn new(socket: TcpStream) -> Self {
        Self {
            inner: Compat01As03::new(socket),
        }
    }
}

impl AsyncRead for TcpSocket {
    fn poll_read(
        mut self: Pin<&mut Self>,
        context: &mut Context,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.inner).poll_read(context, buf)
    }
}

impl AsyncWrite for TcpSocket {
    fn poll_write(
        mut self: Pin<&mut Self>,
        context: &mut Context,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.inner).poll_write(context, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_flush(context)
    }

    fn poll_close(self: Pin<&mut Self>, _context: &mut Context) -> Poll<io::Result<()>> {
        Poll::Ready(self.inner.get_ref().shutdown(Shutdown::Write))
    }
}

fn socketaddr_to_multiaddr(socketaddr: SocketAddr) -> Multiaddr {
    let ipaddr: Multiaddr = socketaddr.ip().into();
    ipaddr.with(Protocol::Tcp(socketaddr.port()))
}

fn multiaddr_to_socketaddr(addr: &Multiaddr) -> ::std::io::Result<SocketAddr> {
    let mut iter = addr.iter();
    let proto1 = iter.next().ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("Invalid Multiaddr '{:?}'", addr),
        )
    })?;
    let proto2 = iter.next().ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("Invalid Multiaddr '{:?}'", addr),
        )
    })?;

    if iter.next().is_some() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("Invalid Multiaddr '{:?}'", addr),
        ));
    }

    match (proto1, proto2) {
        (Protocol::Ip4(ip), Protocol::Tcp(port)) => Ok(SocketAddr::new(ip.into(), port)),
        (Protocol::Ip6(ip), Protocol::Tcp(port)) => Ok(SocketAddr::new(ip.into(), port)),
        _ => Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("Invalid Multiaddr '{:?}'", addr),
        )),
    }
}

#[cfg(test)]
mod test {
    use crate::transport::{tcp::TcpTransport, ConnectionOrigin, Transport, TransportExt};
    use futures::{
        executor::block_on,
        future::{join, FutureExt},
        io::{AsyncReadExt, AsyncWriteExt},
        stream::StreamExt,
    };

    #[test]
    fn simple_listen_and_dial() -> Result<(), ::std::io::Error> {
        let t = TcpTransport::default().and_then(|mut out, connection| {
            async move {
                match connection {
                    ConnectionOrigin::Inbound => {
                        out.write_all(b"Earth").await?;
                        let mut buf = [0; 3];
                        out.read_exact(&mut buf).await?;
                        assert_eq!(&buf, b"Air");
                    }
                    ConnectionOrigin::Outbound => {
                        let mut buf = [0; 5];
                        out.read_exact(&mut buf).await?;
                        assert_eq!(&buf, b"Earth");
                        out.write_all(b"Air").await?;
                    }
                }
                Ok(())
            }
        });

        let (listener, addr) = t.listen_on("/ip4/127.0.0.1/tcp/0".parse().unwrap())?;

        let dial = t.dial(addr)?;
        let listener = listener.into_future().then(|(maybe_result, _stream)| {
            let (incoming, _addr) = maybe_result.unwrap().unwrap();
            incoming.map(Result::unwrap)
        });

        let (outgoing, _incoming) = block_on(join(dial, listener));
        assert!(outgoing.is_ok());
        Ok(())
    }

    #[test]
    fn unsupported_multiaddrs() {
        let t = TcpTransport::default();

        let result = t.listen_on("/memory/0".parse().unwrap());
        assert!(result.is_err());

        let result = t.dial("/memory/22".parse().unwrap());
        assert!(result.is_err());
    }
}
<|endoftext|>