<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this file,
* You can obtain one at http://mozilla.org/MPL/2.0/. */

#![feature(slicing_syntax)]

use std::slice::bytes::copy_memory;

pub static SBOX: [u8, ..256] = [
  0x29, 0x2E, 0x43, 0xC9, 0xA2, 0xD8, 0x7C, 0x01, 0x3D, 0x36, 0x54, 0xA1, 0xEC, 0xF0, 0x06, 0x13,
  0x62, 0xA7, 0x05, 0xF3, 0xC0, 0xC7, 0x73, 0x8C, 0x98, 0x93, 0x2B, 0xD9, 0xBC, 0x4C, 0x82, 0xCA,
  0x1E, 0x9B, 0x57, 0x3C, 0xFD, 0xD4, 0xE0, 0x16, 0x67, 0x42, 0x6F, 0x18, 0x8A, 0x17, 0xE5, 0x12,
  0xBE, 0x4E, 0xC4, 0xD6, 0xDA, 0x9E, 0xDE, 0x49, 0xA0, 0xFB, 0xF5, 0x8E, 0xBB, 0x2F, 0xEE, 0x7A,
  0xA9, 0x68, 0x79, 0x91, 0x15, 0xB2, 0x07, 0x3F, 0x94, 0xC2, 0x10, 0x89, 0x0B, 0x22, 0x5F, 0x21,
  0x80, 0x7F, 0x5D, 0x9A, 0x5A, 0x90, 0x32, 0x27, 0x35, 0x3E, 0xCC, 0xE7, 0xBF, 0xF7, 0x97, 0x03,
  0xFF, 0x19, 0x30, 0xB3, 0x48, 0xA5, 0xB5, 0xD1, 0xD7, 0x5E, 0x92, 0x2A, 0xAC, 0x56, 0xAA, 0xC6,
  0x4F, 0xB8, 0x38, 0xD2, 0x96, 0xA4, 0x7D, 0xB6, 0x76, 0xFC, 0x6B, 0xE2, 0x9C, 0x74, 0x04, 0xF1,
  0x45, 0x9D, 0x70, 0x59, 0x64, 0x71, 0x87, 0x20, 0x86, 0x5B, 0xCF, 0x65, 0xE6, 0x2D, 0xA8, 0x02,
  0x1B, 0x60, 0x25, 0xAD, 0xAE, 0xB0, 0xB9, 0xF6, 0x1C, 0x46, 0x61, 0x69, 0x34, 0x40, 0x7E, 0x0F,
  0x55, 0x47, 0xA3, 0x23, 0xDD, 0x51, 0xAF, 0x3A, 0xC3, 0x5C, 0xF9, 0xCE, 0xBA, 0xC5, 0xEA, 0x26,
  0x2C, 0x53, 0x0D, 0x6E, 0x85, 0x28, 0x84, 0x09, 0xD3, 0xDF, 0xCD, 0xF4, 0x41, 0x81, 0x4D, 0x52,
  0x6A, 0xDC, 0x37, 0xC8, 0x6C, 0xC1, 0xAB, 0xFA, 0x24, 0xE1, 0x7B, 0x08, 0x0C, 0xBD, 0xB1, 0x4A,
  0x78, 0x88, 0x95, 0x8B, 0xE3, 0x63, 0xE8, 0x6D, 0xE9, 0xCB, 0xD5, 0xFE, 0x3B, 0x00, 0x1D, 0x39,
  0xF2, 0xEF, 0xB7, 0x0E, 0x66, 0x58, 0xD0, 0xE4, 0xA6, 0x77, 0x72, 0xF8, 0xEB, 0x75, 0x4B, 0x0A,
  0x31, 0x44, 0x50, 0xB4, 0x8F, 0xED, 0x1F, 0x1A, 0xDB, 0x99, 0x8D, 0x33, 0x9F, 0x11, 0x83, 0x14
];

pub static SBOXI: [u8, ..256] = [
  0xDD, 0x07, 0x8F, 0x5F, 0x7E, 0x12, 0x0E, 0x46, 0xCB, 0xB7, 0xEF, 0x4C, 0xCC, 0xB2, 0xE3, 0x9F,
  0x4A, 0xFD, 0x2F, 0x0F, 0xFF, 0x44, 0x27, 0x2D, 0x2B, 0x61, 0xF7, 0x90, 0x98, 0xDE, 0x20, 0xF6,
  0x87, 0x4F, 0x4D, 0xA3, 0xC8, 0x92, 0xAF, 0x57, 0xB5, 0x00, 0x6B, 0x1A, 0xB0, 0x8D, 0x01, 0x3D,
  0x62, 0xF0, 0x56, 0xFB, 0x9C, 0x58, 0x09, 0xC2, 0x72, 0xDF, 0xA7, 0xDC, 0x23, 0x08, 0x59, 0x47,
  0x9D, 0xBC, 0x29, 0x02, 0xF1, 0x80, 0x99, 0xA1, 0x64, 0x37, 0xCF, 0xEE, 0x1D, 0xBE, 0x31, 0x70,
  0xF2, 0xA5, 0xBF, 0xB1, 0x0A, 0xA0, 0x6D, 0x22, 0xE5, 0x83, 0x54, 0x89, 0xA9, 0x52, 0x69, 0x4E,
  0x91, 0x9A, 0x10, 0xD5, 0x84, 0x8B, 0xE4, 0x28, 0x41, 0x9B, 0xC0, 0x7A, 0xC4, 0xD7, 0xB3, 0x2A,
  0x82, 0x85, 0xEA, 0x16, 0x7D, 0xED, 0x78, 0xE9, 0xD0, 0x42, 0x3F, 0xCA, 0x06, 0x76, 0x9E, 0x51,
  0x50, 0xBD, 0x1E, 0xFE, 0xB6, 0xB4, 0x88, 0x86, 0xD1, 0x4B, 0x2C, 0xD3, 0x17, 0xFA, 0x3B, 0xF4,
  0x55, 0x43, 0x6A, 0x19, 0x48, 0xD2, 0x74, 0x5E, 0x18, 0xF9, 0x53, 0x21, 0x7C, 0x81, 0x35, 0xFC,
  0x38, 0x0B, 0x04, 0xA2, 0x75, 0x65, 0xE8, 0x11, 0x8E, 0x40, 0x6E, 0xC6, 0x6C, 0x93, 0x94, 0xA6,
  0x95, 0xCE, 0x45, 0x63, 0xF3, 0x66, 0x77, 0xE2, 0x71, 0x96, 0xAC, 0x3C, 0x1C, 0xCD, 0x30, 0x5C,
  0x14, 0xC5, 0x49, 0xA8, 0x32, 0xAD, 0x6F, 0x15, 0xC3, 0x03, 0x1F, 0xD9, 0x5A, 0xBA, 0xAB, 0x8A,
  0xE6, 0x67, 0x73, 0xB8, 0x25, 0xDA, 0x33, 0x68, 0x05, 0x1B, 0x34, 0xF8, 0xC1, 0xA4, 0x36, 0xB9,
  0x26, 0xC9, 0x7B, 0xD4, 0xE7, 0x2E, 0x8C, 0x5B, 0xD6, 0xD8, 0xAE, 0xEC, 0x0C, 0xF5, 0x3E, 0xE1,
  0x0D, 0x7F, 0xE0, 0x13, 0xBB, 0x3A, 0x97, 0x5D, 0xEB, 0xAA, 0xC7, 0x39, 0x79, 0x24, 0xDB, 0x60
];

fn pad(msg: &[u8]) -> Vec<u8> {
  let mut msg = msg.to_vec();
  let pad = 16 - msg.len() % 16;
  msg.grow_fn(pad, |_| pad as u8);
  msg
}

fn checksum(msg: &[u8]) -> [u8, ..16] {
  // Message must be padded.
  assert!(msg.len() % 16 == 0);

  let mut checksum = [0u8, ..16];
  let mut last = 0u8;

  for chunk in msg.chunks(16) {
    for (mbyte, cbyte) in chunk.iter().zip(checksum.iter_mut()) {
      *cbyte ^= SBOX[(*mbyte ^ last) as uint];
      last = *cbyte;
    }
  }

  checksum
}

pub fn compress(state: &[u8], msg: &[u8]) -> [u8, ..16] {
  // Two 128 bit blocks in.
  assert!(state.len() == 16 && msg.len() == 16);

  let mut x = [0u8, ..48];
  let mut result = [0u8, ..16];

  // Copy over the previous state.
  copy_memory(x[mut ..16], state);

  // Copy over the message block.
  copy_memory(x[mut 16..32], msg);

  // XOR the previous state and the message block.
  for (i, byte) in msg.iter().enumerate() {
    x[32 + i] = *byte ^ x[i];
  }

  // Encrypt block (18 rounds).
  let mut t = 0u8;
  for i in range(0, 18) {
    for byte in x.iter_mut() {
      *byte ^= SBOX[t as uint];
      t = *byte;
    }
    t += i;
  }

  // Copy the new state.
  copy_memory(&mut result, x[..16]);
  result
}

pub fn digest(msg: &[u8]) -> [u8, ..16] {
  // Pad the message to be a multiple of 16 bytes long.
  let msg = pad(msg);

  // Compress all message blocks.
  let state = msg.chunks(16).fold([0u8, ..16], |s, m| compress(&s, m));

  // Compute the checksum.
  let checksum = checksum(msg[]);

  // Compress checksum and return.
  compress(&state, &checksum)
}

#[cfg(test)]
mod test {
  use digest;

  fn check(exp: &str, msg: &str) {
    let hash = digest(msg.as_bytes());
    let hex = hash.iter().fold(String::new(), |a, &b| format!("{}{:02x}", a, b));
    assert_eq!(exp.to_string(), hex);
  }

  #[test]
  fn test() {
    check("8350e5a3e24c153df2275c9f80692773", "");
    check("32ec01ec4a6dac72c0ab96fb34c0b5d1", "a");
    check("da853b0d3f88d99b30283a69e6ded6bb", "abc");
    check("ab4f496bfb2a530b219ff33031fe06b0", "message digest");
    check("4e8ddff3650292ab5a4108c3aa47940b", "abcdefghijklmnopqrstuvwxyz");
    check("da33def2a42df13975352846c30338cd", "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789");
    check("d5976f79d83d3a0dc9806c3c66f3efd8", "12345678901234567890123456789012345678901234567890123456789012345678901234567890");
  }
}
<commit_msg>lower case hex consts<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this file,
* You can obtain one at http://mozilla.org/MPL/2.0/. */

#![feature(slicing_syntax)]

use std::slice::bytes::copy_memory;

pub static SBOX: [u8, ..256] = [
  0x29, 0x2e, 0x43, 0xc9, 0xa2, 0xd8, 0x7c, 0x01, 0x3d, 0x36, 0x54, 0xa1, 0xec, 0xf0, 0x06, 0x13,
  0x62, 0xa7, 0x05, 0xf3, 0xc0, 0xc7, 0x73, 0x8c, 0x98, 0x93, 0x2b, 0xd9, 0xbc, 0x4c, 0x82, 0xca,
  0x1e, 0x9b, 0x57, 0x3c, 0xfd, 0xd4, 0xe0, 0x16, 0x67, 0x42, 0x6f, 0x18, 0x8a, 0x17, 0xe5, 0x12,
  0xbe, 0x4e, 0xc4, 0xd6, 0xda, 0x9e, 0xde, 0x49, 0xa0, 0xfb, 0xf5, 0x8e, 0xbb, 0x2f, 0xee, 0x7a,
  0xa9, 0x68, 0x79, 0x91, 0x15, 0xb2, 0x07, 0x3f, 0x94, 0xc2, 0x10, 0x89, 0x0b, 0x22, 0x5f, 0x21,
  0x80, 0x7f, 0x5d, 0x9a, 0x5a, 0x90, 0x32, 0x27, 0x35, 0x3e, 0xcc, 0xe7, 0xbf, 0xf7, 0x97, 0x03,
  0xff, 0x19, 0x30, 0xb3, 0x48, 0xa5, 0xb5, 0xd1, 0xd7, 0x5e, 0x92, 0x2a, 0xac, 0x56, 0xaa, 0xc6,
  0x4f, 0xb8, 0x38, 0xd2, 0x96, 0xa4, 0x7d, 0xb6, 0x76, 0xfc, 0x6b, 0xe2, 0x9c, 0x74, 0x04, 0xf1,
  0x45, 0x9d, 0x70, 0x59, 0x64, 0x71, 0x87, 0x20, 0x86, 0x5b, 0xcf, 0x65, 0xe6, 0x2d, 0xa8, 0x02,
  0x1b, 0x60, 0x25, 0xad, 0xae, 0xb0, 0xb9, 0xf6, 0x1c, 0x46, 0x61, 0x69, 0x34, 0x40, 0x7e, 0x0f,
  0x55, 0x47, 0xa3, 0x23, 0xdd, 0x51, 0xaf, 0x3a, 0xc3, 0x5c, 0xf9, 0xce, 0xba, 0xc5, 0xea, 0x26,
  0x2c, 0x53, 0x0d, 0x6e, 0x85, 0x28, 0x84, 0x09, 0xd3, 0xdf, 0xcd, 0xf4, 0x41, 0x81, 0x4d, 0x52,
  0x6a, 0xdc, 0x37, 0xc8, 0x6c, 0xc1, 0xab, 0xfa, 0x24, 0xe1, 0x7b, 0x08, 0x0c, 0xbd, 0xb1, 0x4a,
  0x78, 0x88, 0x95, 0x8b, 0xe3, 0x63, 0xe8, 0x6d, 0xe9, 0xcb, 0xd5, 0xfe, 0x3b, 0x00, 0x1d, 0x39,
  0xf2, 0xef, 0xb7, 0x0e, 0x66, 0x58, 0xd0, 0xe4, 0xa6, 0x77, 0x72, 0xf8, 0xeb, 0x75, 0x4b, 0x0a,
  0x31, 0x44, 0x50, 0xb4, 0x8f, 0xed, 0x1f, 0x1a, 0xdb, 0x99, 0x8d, 0x33, 0x9f, 0x11, 0x83, 0x14
];

pub static SBOXI: [u8, ..256] = [
  0xdd, 0x07, 0x8f, 0x5f, 0x7e, 0x12, 0x0e, 0x46, 0xcb, 0xb7, 0xef, 0x4c, 0xcc, 0xb2, 0xe3, 0x9f,
  0x4a, 0xfd, 0x2f, 0x0f, 0xff, 0x44, 0x27, 0x2d, 0x2b, 0x61, 0xf7, 0x90, 0x98, 0xde, 0x20, 0xf6,
  0x87, 0x4f, 0x4d, 0xa3, 0xc8, 0x92, 0xaf, 0x57, 0xb5, 0x00, 0x6b, 0x1a, 0xb0, 0x8d, 0x01, 0x3d,
  0x62, 0xf0, 0x56, 0xfb, 0x9c, 0x58, 0x09, 0xc2, 0x72, 0xdf, 0xa7, 0xdc, 0x23, 0x08, 0x59, 0x47,
  0x9d, 0xbc, 0x29, 0x02, 0xf1, 0x80, 0x99, 0xa1, 0x64, 0x37, 0xcf, 0xee, 0x1d, 0xbe, 0x31, 0x70,
  0xf2, 0xa5, 0xbf, 0xb1, 0x0a, 0xa0, 0x6d, 0x22, 0xe5, 0x83, 0x54, 0x89, 0xa9, 0x52, 0x69, 0x4e,
  0x91, 0x9a, 0x10, 0xd5, 0x84, 0x8b, 0xe4, 0x28, 0x41, 0x9b, 0xc0, 0x7a, 0xc4, 0xd7, 0xb3, 0x2a,
  0x82, 0x85, 0xea, 0x16, 0x7d, 0xed, 0x78, 0xe9, 0xd0, 0x42, 0x3f, 0xca, 0x06, 0x76, 0x9e, 0x51,
  0x50, 0xbd, 0x1e, 0xfe, 0xb6, 0xb4, 0x88, 0x86, 0xd1, 0x4b, 0x2c, 0xd3, 0x17, 0xfa, 0x3b, 0xf4,
  0x55, 0x43, 0x6a, 0x19, 0x48, 0xd2, 0x74, 0x5e, 0x18, 0xf9, 0x53, 0x21, 0x7c, 0x81, 0x35, 0xfc,
  0x38, 0x0b, 0x04, 0xa2, 0x75, 0x65, 0xe8, 0x11, 0x8e, 0x40, 0x6e, 0xc6, 0x6c, 0x93, 0x94, 0xa6,
  0x95, 0xce, 0x45, 0x63, 0xf3, 0x66, 0x77, 0xe2, 0x71, 0x96, 0xac, 0x3c, 0x1c, 0xcd, 0x30, 0x5c,
  0x14, 0xc5, 0x49, 0xa8, 0x32, 0xad, 0x6f, 0x15, 0xc3, 0x03, 0x1f, 0xd9, 0x5a, 0xba, 0xab, 0x8a,
  0xe6, 0x67, 0x73, 0xb8, 0x25, 0xda, 0x33, 0x68, 0x05, 0x1b, 0x34, 0xf8, 0xc1, 0xa4, 0x36, 0xb9,
  0x26, 0xc9, 0x7b, 0xd4, 0xe7, 0x2e, 0x8c, 0x5b, 0xd6, 0xd8, 0xae, 0xec, 0x0c, 0xf5, 0x3e, 0xe1,
  0x0d, 0x7f, 0xe0, 0x13, 0xbb, 0x3a, 0x97, 0x5d, 0xeb, 0xaa, 0xc7, 0x39, 0x79, 0x24, 0xdb, 0x60
];

fn pad(msg: &[u8]) -> Vec<u8> {
  let mut msg = msg.to_vec();
  let pad = 16 - msg.len() % 16;
  msg.grow_fn(pad, |_| pad as u8);
  msg
}

fn checksum(msg: &[u8]) -> [u8, ..16] {
  // Message must be padded.
  assert!(msg.len() % 16 == 0);

  let mut checksum = [0u8, ..16];
  let mut last = 0u8;

  for chunk in msg.chunks(16) {
    for (mbyte, cbyte) in chunk.iter().zip(checksum.iter_mut()) {
      *cbyte ^= SBOX[(*mbyte ^ last) as uint];
      last = *cbyte;
    }
  }

  checksum
}

pub fn compress(state: &[u8], msg: &[u8]) -> [u8, ..16] {
  // Two 128 bit blocks in.
  assert!(state.len() == 16 && msg.len() == 16);

  let mut x = [0u8, ..48];
  let mut result = [0u8, ..16];

  // Copy over the previous state.
  copy_memory(x[mut ..16], state);

  // Copy over the message block.
  copy_memory(x[mut 16..32], msg);

  // XOR the previous state and the message block.
  for (i, byte) in msg.iter().enumerate() {
    x[32 + i] = *byte ^ x[i];
  }

  // Encrypt block (18 rounds).
  let mut t = 0u8;
  for i in range(0, 18) {
    for byte in x.iter_mut() {
      *byte ^= SBOX[t as uint];
      t = *byte;
    }
    t += i;
  }

  // Copy the new state.
  copy_memory(&mut result, x[..16]);
  result
}

pub fn digest(msg: &[u8]) -> [u8, ..16] {
  // Pad the message to be a multiple of 16 bytes long.
  let msg = pad(msg);

  // Compress all message blocks.
  let state = msg.chunks(16).fold([0u8, ..16], |s, m| compress(&s, m));

  // Compute the checksum.
  let checksum = checksum(msg[]);

  // Compress checksum and return.
  compress(&state, &checksum)
}

#[cfg(test)]
mod test {
  use digest;

  fn check(exp: &str, msg: &str) {
    let hash = digest(msg.as_bytes());
    let hex = hash.iter().fold(String::new(), |a, &b| format!("{}{:02x}", a, b));
    assert_eq!(exp.to_string(), hex);
  }

  #[test]
  fn test() {
    check("8350e5a3e24c153df2275c9f80692773", "");
    check("32ec01ec4a6dac72c0ab96fb34c0b5d1", "a");
    check("da853b0d3f88d99b30283a69e6ded6bb", "abc");
    check("ab4f496bfb2a530b219ff33031fe06b0", "message digest");
    check("4e8ddff3650292ab5a4108c3aa47940b", "abcdefghijklmnopqrstuvwxyz");
    check("da33def2a42df13975352846c30338cd", "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789");
    check("d5976f79d83d3a0dc9806c3c66f3efd8", "12345678901234567890123456789012345678901234567890123456789012345678901234567890");
  }
}
<|endoftext|>


<commit_before>extern crate libc;

pub mod sys;
pub mod os;
pub mod num;
pub mod geom;
pub mod application;
pub mod renderer;
pub mod creator;
pub mod error;
pub mod frame;
pub mod event;
pub mod event_loop;
pub mod draw;
pub mod raster;
pub mod devel;
pub mod node;

pub use application::Application;
pub use renderer::Renderer;
pub use frame::Frame;
<commit_msg>Including the new lang module<commit_after>extern crate libc;

pub mod sys;
pub mod os;
pub mod num;
pub mod geom;
pub mod application;
pub mod renderer;
pub mod creator;
pub mod error;
pub mod frame;
pub mod event;
pub mod event_loop;
pub mod draw;
pub mod raster;
pub mod devel;
pub mod node;
pub mod lang;

pub use application::Application;
pub use renderer::Renderer;
pub use frame::Frame;
<|endoftext|>


<commit_before>extern crate bufstream;

mod message;
mod user;
mod client;

pub use std::str::FromStr;
pub use message::{IrcMessage, NumericReply};
pub use user::IrcUser;
pub use client::IrcClient;
<commit_msg>Export new trait for sending messages to IRC<commit_after>extern crate bufstream;

mod message;
mod user;
mod client;

pub use std::str::FromStr;
pub use message::{IrcMessage, NumericReply};
pub use user::IrcUser;
pub use client::{IrcClient, IrcWrite};
<|endoftext|>


<commit_before>//! # generator
//!
//! Rust generator library
//!

#![feature(fnbox)]
#![feature(rustc_private)]
#![feature(rt)]
#![feature(box_raw)]
#![cfg_attr(test, deny(warnings))]
#![deny(missing_docs)]

#[macro_use]
extern crate log;
extern crate libc;
extern crate context;

pub mod env;
pub use env::ContextStack;

use context::Stack;
use context::Context as RegContext;
use std::rt::unwind::try;
use std::rt::util::min_stack;
use std::boxed::FnBox;
use std::any::Any;
use std::mem;
use std::ptr;

/// generator context
pub struct Context {
    /// generator regs context
    regs: RegContext,
    /// generator execution stack
    stack: Stack,
    /// passed in para for send
    para: *mut Any,
    /// this is just a buffer for the return value
    ret:  *mut Any,
    /// track generator ref, yield will -1, send will +1
    _ref: u32,
    /// priave flag that control the execution flow
    _flag: bool,
}

impl Context {
    /// return a default generator context
    pub fn new() -> Context {
        Context {
            regs: RegContext::empty(),
            stack: Stack::new(min_stack()),
            para: unsafe { mem::transmute(&0 as &Any) },
            ret: unsafe { mem::transmute(&0 as &Any) },
            _ref: 0,
            _flag: false,
        }
    }

    /// empty context used for the normal thread
    pub fn empty() -> Context {
        Context {
            regs: RegContext::empty(),
            stack: unsafe { Stack::dummy_stack() },
            para: unsafe { mem::transmute(&0 as &Any) },
            ret: unsafe { mem::transmute(&0 as &Any) },
            _ref: 0,
            _flag: false,
        }
    }

    /// save current generator context
    #[inline(always)]
    pub fn save(&mut self) {
        self._ref -= 1;
        self._flag = true;
        RegContext::save(&mut self.regs);
    }

    /// generator involke flag, save it along with context
    #[inline]
    pub fn get_flag(&mut self)-> &'static mut bool {
        unsafe {mem::transmute(&mut self._flag)}
    }

    /// get current generator send para
    #[inline]
    pub fn get_para<T>(&self) -> Option<T> where T: Any {
        let para = unsafe { &mut *self.para };
        let val = para.downcast_mut::<Option<T>>().unwrap();
        // simple consume the value
        mem::replace(val, None)
    }

    /// set current generator return value
    #[inline]
    pub fn set_ret<T>(&mut self, v: T) where T: Any {
        let ret = unsafe { &mut *self.ret };
        // add type check and panic with message
        let val = ret.downcast_mut::<Option<T>>().unwrap();
        mem::replace(val, Some(v));
    }
}

/// generator trait
pub trait Generator<A> {
    /// Output type
    type Output;

    /// send interface
    fn send(&mut self, para: A) -> Self::Output;

    /// is finished
    fn is_done(&self) -> bool;
}

/// FnGenerator
pub struct FnGenerator<A, T> {
    context: Context,
    // save the input
    para: Option<A>,
    // save the output
    ret: Option<T>,
    // boxed functor
    f: Option<Box<FnBox()->T>>
}

impl<A, T> FnGenerator<A, T> {
    fn resume_gen(&mut self) {
        let env = ContextStack::current();
        let cur = &mut env.top().regs;
        let ctx = &mut self.context as *mut Context;
        let to = unsafe {&mut (*ctx).regs};
        // save current generator context on stack
        env.push(ctx);
        // switch context
        RegContext::swap(cur, to);
    }

    #[inline]
    fn is_started(&self) -> bool {
        // when the f is consumed we think it's running
        self.f.is_none()
    }

    fn send_impl(&mut self, para: Option<A>) -> Option<T> {
        if self.is_started() && self.context._ref != 0 {
            return None;
        }

        // every time we call the function, increase the ref count
        // yiled will decrease it and return will not
        self.context._ref += 1;
        self.resume_gen();

        // this is the passed in value of the send primitive
        // the yield part would read out this value in the next round
        // use the replace is would drop the old value
        mem::replace(&mut self.para, para);

        self.ret.take()
    }
}

impl<A: Any, T: Any> Generator<A> for FnGenerator<A, T> {
    type Output = T;

    fn send(&mut self, para: A) -> T {
        let ret = self.send_impl(Some(para));
        ret.unwrap()
    }

    fn is_done(&self) -> bool {
       self.is_started() && self.context._ref != 0
    }
}

impl<A: Any, T: Any> Iterator for FnGenerator<A, T> {
    type Item = T;
    // The 'Iterator' trait only requires the 'next' method to be defined. The
    // return type is 'Option<T>', 'None' is returned when the 'Iterator' is
    // over, otherwise the next value is returned wrapped in 'Some'
    fn next(&mut self) -> Option<T> {
        self.send_impl(None)
    }
}

/// switch back to parent context
pub fn yield_now() {
    let env = ContextStack::current();
    let ctx = env.top();
    let sp = ctx.stack.start();
    // judge if this is root context
    if sp != ptr::null() {
        env.pop();
        let ctx = env.top();
        RegContext::load(&ctx.regs);
    }
}

#[allow(unused_variables)]
extern "C" fn gen_init(arg: usize, f: *mut libc::c_void) {
    let func: Box<Box<FnBox()>> = unsafe {
        Box::from_raw(f as *mut Box<FnBox()>)
    };

    if let Err(cause) = unsafe { try(move|| func()) } {
        error!("Panicked inside: {:?}", cause.downcast::<&str>());
    }

    yield_now();
}

/// create generator
pub fn make_gen<A: Any, T: Any>(p: Option<A>, f: Box<FnBox()->T>) -> Box<FnGenerator<A, T>> {
    let mut g = Box::new(FnGenerator {
       para: p, ret: None, f: Some(f),
       context: Context::new()
    });

    g.context.para = &mut g.para as &mut Any;
    g.context.ret = &mut g.ret as &mut Any;

    unsafe {
        let ptr = Box::into_raw(g);

        //let start = Box::new(||{g.ret = Some((g.f.take().unwrap())())});
        let start = Box::new(move||{
            let mut g = Box::from_raw(ptr);
            g.ret = Some((g.f.take().unwrap())());
            // don't free the box here
            Box::into_raw(g);
        });

        let stk = &mut (*ptr).context.stack;
        let reg = &mut (*ptr).context.regs;
        reg.init_with(gen_init, ptr as usize, start, stk);
        Box::from_raw(ptr)
    }
}

#[macro_export]
macro_rules! generator {
    // `(func, <type>)`
    // func: the expression for unsafe async function which contains yiled
    // para: default send para type to the generator
    ($func:expr, <$para:ty>) => (
        generator::make_gen::<$para, _>(None, Box::new(move|| unsafe{$func}))
    );

    // `(func, para)`
    // func: the expression for unsafe async function which contains yiled
    // para: default send para to the generator
    ($func:expr, $para:expr) => (
        generator::make_gen(Some($para), Box::new(move|| unsafe{$func}))
    );

    ($func:expr) => (generator!($func, ()));
}

#[macro_export]
macro_rules! _yield {
    // `(para)`
    // val: the value that need to be yield
    // and got the send para from context
    ($val:expr) => ({
        _yield_!($val);
        generator::ContextStack::current().top().get_para().unwrap()
    });

    () => (_yield!(()));
}

/// yield without get the passed in para
#[macro_export]
macro_rules! _yield_ {
    // `(para)`
    // val: the value that need to be yield
    // and got the send para from context
    ($val:expr) => ({
        let context = generator::ContextStack::current().top();
        let _no_use = context.get_flag();
        context.save();
        if *_no_use {
            *_no_use = false;
            context.set_ret($val);
            // don't use the return instruction
            generator::yield_now();
            // context.load();
            return $val;
        }
    });

    () => (_yield_!(()));
}

<commit_msg>use option take<commit_after>//! # generator
//!
//! Rust generator library
//!

#![feature(fnbox)]
#![feature(rustc_private)]
#![feature(rt)]
#![feature(box_raw)]
#![cfg_attr(test, deny(warnings))]
#![deny(missing_docs)]

#[macro_use]
extern crate log;
extern crate libc;
extern crate context;

pub mod env;
pub use env::ContextStack;

use context::Stack;
use context::Context as RegContext;
use std::rt::unwind::try;
use std::rt::util::min_stack;
use std::boxed::FnBox;
use std::any::Any;
use std::mem;
use std::ptr;

/// generator context
pub struct Context {
    /// generator regs context
    regs: RegContext,
    /// generator execution stack
    stack: Stack,
    /// passed in para for send
    para: *mut Any,
    /// this is just a buffer for the return value
    ret:  *mut Any,
    /// track generator ref, yield will -1, send will +1
    _ref: u32,
    /// priave flag that control the execution flow
    _flag: bool,
}

impl Context {
    /// return a default generator context
    pub fn new() -> Context {
        Context {
            regs: RegContext::empty(),
            stack: Stack::new(min_stack()),
            para: unsafe { mem::transmute(&0 as &Any) },
            ret: unsafe { mem::transmute(&0 as &Any) },
            _ref: 0,
            _flag: false,
        }
    }

    /// empty context used for the normal thread
    pub fn empty() -> Context {
        Context {
            regs: RegContext::empty(),
            stack: unsafe { Stack::dummy_stack() },
            para: unsafe { mem::transmute(&0 as &Any) },
            ret: unsafe { mem::transmute(&0 as &Any) },
            _ref: 0,
            _flag: false,
        }
    }

    /// save current generator context
    #[inline(always)]
    pub fn save(&mut self) {
        self._ref -= 1;
        self._flag = true;
        RegContext::save(&mut self.regs);
    }

    /// generator involke flag, save it along with context
    #[inline]
    pub fn get_flag(&mut self)-> &'static mut bool {
        unsafe {mem::transmute(&mut self._flag)}
    }

    /// get current generator send para
    #[inline]
    pub fn get_para<T>(&self) -> Option<T> where T: Any {
        let para = unsafe { &mut *self.para };
        let val = para.downcast_mut::<Option<T>>().unwrap();
        val.take()
    }

    /// set current generator return value
    #[inline]
    pub fn set_ret<T>(&mut self, v: T) where T: Any {
        let ret = unsafe { &mut *self.ret };
        // add type check and panic with message
        let val = ret.downcast_mut::<Option<T>>().unwrap();
        mem::replace(val, Some(v));
    }
}

/// generator trait
pub trait Generator<A> {
    /// Output type
    type Output;

    /// send interface
    fn send(&mut self, para: A) -> Self::Output;

    /// is finished
    fn is_done(&self) -> bool;
}

/// FnGenerator
pub struct FnGenerator<A, T> {
    context: Context,
    // save the input
    para: Option<A>,
    // save the output
    ret: Option<T>,
    // boxed functor
    f: Option<Box<FnBox()->T>>
}

impl<A, T> FnGenerator<A, T> {
    fn resume_gen(&mut self) {
        let env = ContextStack::current();
        let cur = &mut env.top().regs;
        let ctx = &mut self.context as *mut Context;
        let to = unsafe {&mut (*ctx).regs};
        // save current generator context on stack
        env.push(ctx);
        // switch context
        RegContext::swap(cur, to);
    }

    #[inline]
    fn is_started(&self) -> bool {
        // when the f is consumed we think it's running
        self.f.is_none()
    }

    fn send_impl(&mut self, para: Option<A>) -> Option<T> {
        if self.is_started() && self.context._ref != 0 {
            return None;
        }

        // every time we call the function, increase the ref count
        // yiled will decrease it and return will not
        self.context._ref += 1;
        self.resume_gen();

        // this is the passed in value of the send primitive
        // the yield part would read out this value in the next round
        // use the replace is would drop the old value
        mem::replace(&mut self.para, para);

        self.ret.take()
    }
}

impl<A: Any, T: Any> Generator<A> for FnGenerator<A, T> {
    type Output = T;

    fn send(&mut self, para: A) -> T {
        let ret = self.send_impl(Some(para));
        ret.unwrap()
    }

    fn is_done(&self) -> bool {
       self.is_started() && self.context._ref != 0
    }
}

impl<A: Any, T: Any> Iterator for FnGenerator<A, T> {
    type Item = T;
    // The 'Iterator' trait only requires the 'next' method to be defined. The
    // return type is 'Option<T>', 'None' is returned when the 'Iterator' is
    // over, otherwise the next value is returned wrapped in 'Some'
    fn next(&mut self) -> Option<T> {
        self.send_impl(None)
    }
}

/// switch back to parent context
pub fn yield_now() {
    let env = ContextStack::current();
    let ctx = env.top();
    let sp = ctx.stack.start();
    // judge if this is root context
    if sp != ptr::null() {
        env.pop();
        let ctx = env.top();
        RegContext::load(&ctx.regs);
    }
}

#[allow(unused_variables)]
extern "C" fn gen_init(arg: usize, f: *mut libc::c_void) {
    let func: Box<Box<FnBox()>> = unsafe {
        Box::from_raw(f as *mut Box<FnBox()>)
    };

    if let Err(cause) = unsafe { try(move|| func()) } {
        error!("Panicked inside: {:?}", cause.downcast::<&str>());
    }

    yield_now();
}

/// create generator
pub fn make_gen<A: Any, T: Any>(p: Option<A>, f: Box<FnBox()->T>) -> Box<FnGenerator<A, T>> {
    let mut g = Box::new(FnGenerator {
       para: p, ret: None, f: Some(f),
       context: Context::new()
    });

    g.context.para = &mut g.para as &mut Any;
    g.context.ret = &mut g.ret as &mut Any;

    unsafe {
        let ptr = Box::into_raw(g);

        //let start = Box::new(||{g.ret = Some((g.f.take().unwrap())())});
        let start = Box::new(move||{
            let mut g = Box::from_raw(ptr);
            g.ret = Some((g.f.take().unwrap())());
            // don't free the box here
            Box::into_raw(g);
        });

        let stk = &mut (*ptr).context.stack;
        let reg = &mut (*ptr).context.regs;
        reg.init_with(gen_init, ptr as usize, start, stk);
        Box::from_raw(ptr)
    }
}

#[macro_export]
macro_rules! generator {
    // `(func, <type>)`
    // func: the expression for unsafe async function which contains yiled
    // para: default send para type to the generator
    ($func:expr, <$para:ty>) => (
        generator::make_gen::<$para, _>(None, Box::new(move|| unsafe{$func}))
    );

    // `(func, para)`
    // func: the expression for unsafe async function which contains yiled
    // para: default send para to the generator
    ($func:expr, $para:expr) => (
        generator::make_gen(Some($para), Box::new(move|| unsafe{$func}))
    );

    ($func:expr) => (generator!($func, ()));
}

#[macro_export]
macro_rules! _yield {
    // `(para)`
    // val: the value that need to be yield
    // and got the send para from context
    ($val:expr) => ({
        _yield_!($val);
        generator::ContextStack::current().top().get_para().unwrap()
    });

    () => (_yield!(()));
}

/// yield without get the passed in para
#[macro_export]
macro_rules! _yield_ {
    // `(para)`
    // val: the value that need to be yield
    // and got the send para from context
    ($val:expr) => ({
        let context = generator::ContextStack::current().top();
        let _no_use = context.get_flag();
        context.save();
        if *_no_use {
            *_no_use = false;
            context.set_ret($val);
            // don't use the return instruction
            generator::yield_now();
            // context.load();
            return $val;
        }
    });

    () => (_yield_!(()));
}

<|endoftext|>


<commit_before>// Copyright 2021, The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! A library for passing arbitrary file descriptors when spawning child processes.
//!
//! # Example
//!
//! ```rust
//! use command_fds::{CommandFdExt, FdMapping};
//! use std::fs::File;
//! use std::os::unix::io::AsRawFd;
//! use std::process::Command;
//!
//! // Open a file.
//! let file = File::open("Cargo.toml").unwrap();
//!
//! // Prepare to run `ls -l /proc/self/fd` with some FDs mapped.
//! let mut command = Command::new("ls");
//! command.arg("-l").arg("/proc/self/fd");
//! command
//!     .fd_mappings(vec![
//!         // Map `file` as FD 3 in the child process.
//!         FdMapping {
//!             parent_fd: file.as_raw_fd(),
//!             child_fd: 3,
//!         },
//!         // Map this process's stdin as FD 5 in the child process.
//!         FdMapping {
//!             parent_fd: 0,
//!             child_fd: 5,
//!         },
//!     ])
//!     .unwrap();
//!
//! // Spawn the child process.
//! let mut child = command.spawn().unwrap();
//! child.wait().unwrap();
//! ```

use nix::fcntl::{fcntl, FcntlArg, FdFlag};
use nix::unistd::dup2;
use std::cmp::max;
use std::io::{self, ErrorKind};
use std::os::unix::io::RawFd;
use std::os::unix::process::CommandExt;
use std::process::Command;
use thiserror::Error;

/// A mapping from a file descriptor in the parent to a file descriptor in the child, to be applied
/// when spawning a child process.
///
/// The parent_fd must be kept open until after the child is spawned.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FdMapping {
    pub parent_fd: RawFd,
    pub child_fd: RawFd,
}

/// Error setting up FD mappings, because there were two or more mappings for the same child FD.
#[derive(Copy, Clone, Debug, Eq, Error, PartialEq)]
#[error("Two or more mappings for the same child FD")]
pub struct FdMappingCollision;

/// Extension to add file descriptor mappings to a [`Command`].
pub trait CommandFdExt {
    /// Adds the given set of file descriptors to the command.
    ///
    /// Calling this more than once on the same command may result in unexpected behaviour.
    fn fd_mappings(&mut self, mappings: Vec<FdMapping>) -> Result<&mut Self, FdMappingCollision>;

    /// Adds the given set of file descriptors to be passed on to the child process when the command
    /// is run.
    fn preserved_fds(&mut self, fds: Vec<RawFd>) -> &mut Self;
}

impl CommandFdExt for Command {
    fn fd_mappings(
        &mut self,
        mut mappings: Vec<FdMapping>,
    ) -> Result<&mut Self, FdMappingCollision> {
        // Validate that there are no conflicting mappings to the same child FD.
        let mut child_fds: Vec<RawFd> = mappings.iter().map(|mapping| mapping.child_fd).collect();
        child_fds.sort_unstable();
        child_fds.dedup();
        if child_fds.len() != mappings.len() {
            return Err(FdMappingCollision);
        }

        // Register the callback to apply the mappings after forking but before execing.
        // Safety: `map_fds` will not allocate, so it is safe to call from this hook.
        unsafe {
            self.pre_exec(move || map_fds(&mut mappings, &child_fds));
        }

        Ok(self)
    }

    fn preserved_fds(&mut self, fds: Vec<RawFd>) -> &mut Self {
        unsafe {
            self.pre_exec(move || preserve_fds(&fds));
        }

        self
    }
}

// This function must not do any allocation, as it is called from the pre_exec hook.
fn map_fds(mappings: &mut [FdMapping], child_fds: &[RawFd]) -> io::Result<()> {
    if mappings.is_empty() {
        // No need to do anything, and finding first_unused_fd would fail.
        return Ok(());
    }

    // Find the first FD which is higher than any parent or child FD in the mapping, so we can
    // safely use it and higher FDs as temporary FDs. There may be other files open with these FDs,
    // so we still need to ensure we don't conflict with them.
    let first_safe_fd = mappings
        .iter()
        .map(|mapping| max(mapping.parent_fd, mapping.child_fd))
        .max()
        .unwrap()
        + 1;

    // If any parent FDs conflict with child FDs, then first duplicate them to a temporary FD which
    // is clear of either range. Mappings to the same FD are fine though, we can handle them by just
    // removing the FD_CLOEXEC flag from the existing (parent) FD.
    for mapping in mappings.iter_mut() {
        if child_fds.contains(&mapping.parent_fd) && mapping.parent_fd != mapping.child_fd {
            mapping.parent_fd = fcntl(mapping.parent_fd, FcntlArg::F_DUPFD_CLOEXEC(first_safe_fd))
                .map_err(nix_to_io_error)?;
        }
    }

    // Now we can actually duplicate FDs to the desired child FDs.
    for mapping in mappings {
        if mapping.child_fd == mapping.parent_fd {
            // Remove the FD_CLOEXEC flag, so the FD will be kept open when exec is called for the
            // child.
            fcntl(mapping.parent_fd, FcntlArg::F_SETFD(FdFlag::empty()))
                .map_err(nix_to_io_error)?;
        } else {
            // This closes child_fd if it is already open as something else, and clears the
            // FD_CLOEXEC flag on child_fd.
            dup2(mapping.parent_fd, mapping.child_fd).map_err(nix_to_io_error)?;
        }
    }

    Ok(())
}

fn preserve_fds(fds: &[RawFd]) -> io::Result<()> {
    for fd in fds {
        // Remove the FD_CLOEXEC flag, so the FD will be kept open when exec is called for the
        // child.
        fcntl(*fd, FcntlArg::F_SETFD(FdFlag::empty())).map_err(nix_to_io_error)?;
    }

    Ok(())
}

/// Convert a [`nix::Error`] to a [`std::io::Error`].
fn nix_to_io_error(error: nix::Error) -> io::Error {
    if let nix::Error::Sys(errno) = error {
        io::Error::from_raw_os_error(errno as i32)
    } else {
        io::Error::new(ErrorKind::Other, error)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use nix::unistd::close;
    use std::collections::HashSet;
    use std::fs::{read_dir, File};
    use std::os::unix::io::AsRawFd;
    use std::process::Output;
    use std::str;
    use std::sync::Once;

    static SETUP: Once = Once::new();

    #[test]
    fn conflicting_mappings() {
        setup();

        let mut command = Command::new("ls");

        // The same mapping can't be included twice.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
            ])
            .is_err());

        // Mapping two different FDs to the same FD isn't allowed either.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
                FdMapping {
                    child_fd: 4,
                    parent_fd: 6,
                },
            ])
            .is_err());
    }

    #[test]
    fn no_mappings() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        assert!(command.fd_mappings(vec![]).is_ok());

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3], 0);
    }

    #[test]
    fn none_preserved() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        command.preserved_fds(vec![]);

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3], 0);
    }

    #[test]
    fn one_mapping() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file = File::open("testdata/file1.txt").unwrap();
        // Map the file an otherwise unused FD.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: file.as_raw_fd(),
                child_fd: 5,
            },])
            .is_ok());

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3, 5], 0);
    }

    #[test]
    fn one_preserved() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file = File::open("testdata/file1.txt").unwrap();
        let file_fd = file.as_raw_fd();
        command.preserved_fds(vec![file_fd]);

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3, file_fd], 0);
    }

    #[test]
    fn swap_mappings() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file1 = File::open("testdata/file1.txt").unwrap();
        let file2 = File::open("testdata/file2.txt").unwrap();
        let fd1 = file1.as_raw_fd();
        let fd2 = file2.as_raw_fd();
        // Map files to each other's FDs, to ensure that the temporary FD logic works.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    parent_fd: fd1,
                    child_fd: fd2,
                },
                FdMapping {
                    parent_fd: fd2,
                    child_fd: fd1,
                },
            ])
            .is_ok(),);

        let output = command.output().unwrap();
        // Expect one more Fd for the /proc/self/fd directory. We can't predict what number it will
        // be assigned, because 3 might or might not be taken already by fd1 or fd2.
        expect_fds(&output, &[0, 1, 2, fd1, fd2], 1);
    }

    #[test]
    fn one_to_one_mapping() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file1 = File::open("testdata/file1.txt").unwrap();
        let file2 = File::open("testdata/file2.txt").unwrap();
        let fd1 = file1.as_raw_fd();
        // Map file1 to the same FD it currently has, to ensure the special case for that works.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: fd1,
                child_fd: fd1,
            }])
            .is_ok());

        let output = command.output().unwrap();
        // Expect one more Fd for the /proc/self/fd directory. We can't predict what number it will
        // be assigned, because 3 might or might not be taken already by fd1 or fd2.
        expect_fds(&output, &[0, 1, 2, fd1], 1);

        // Keep file2 open until the end, to ensure that it's not passed to the child.
        drop(file2);
    }

    #[test]
    fn map_stdin() {
        setup();

        let mut command = Command::new("cat");

        let file = File::open("testdata/file1.txt").unwrap();
        // Map the file to stdin.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: file.as_raw_fd(),
                child_fd: 0,
            },])
            .is_ok());

        let output = command.output().unwrap();
        assert!(output.status.success());
        assert_eq!(output.stdout, b"test 1");
    }

    /// Parse the output of ls into a set of filenames
    fn parse_ls_output(output: &[u8]) -> HashSet<String> {
        str::from_utf8(output)
            .unwrap()
            .split_terminator("\n")
            .map(str::to_owned)
            .collect()
    }

    /// Check that the output of `ls /proc/self/fd` contains the expected set of FDs, plus exactly
    /// `extra` extra FDs.
    fn expect_fds(output: &Output, expected_fds: &[RawFd], extra: usize) {
        assert!(output.status.success());
        let expected_fds: HashSet<String> = expected_fds.iter().map(RawFd::to_string).collect();
        let fds = parse_ls_output(&output.stdout);
        if extra == 0 {
            assert_eq!(fds, expected_fds);
        } else {
            assert!(expected_fds.is_subset(&fds));
            assert_eq!(fds.len(), expected_fds.len() + extra);
        }
    }

    fn setup() {
        SETUP.call_once(close_excess_fds);
    }

    /// Close all file descriptors apart from stdin, stdout and stderr.
    ///
    /// This is necessary because GitHub Actions opens a bunch of others for some reason.
    fn close_excess_fds() {
        let dir = read_dir("/proc/self/fd").unwrap();
        for entry in dir {
            let entry = entry.unwrap();
            let fd: RawFd = entry.file_name().to_str().unwrap().parse().unwrap();
            if fd > 3 {
                close(fd).unwrap();
            }
        }
    }
}
<commit_msg>Add warning about running command with mappings multiple times.<commit_after>// Copyright 2021, The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! A library for passing arbitrary file descriptors when spawning child processes.
//!
//! # Example
//!
//! ```rust
//! use command_fds::{CommandFdExt, FdMapping};
//! use std::fs::File;
//! use std::os::unix::io::AsRawFd;
//! use std::process::Command;
//!
//! // Open a file.
//! let file = File::open("Cargo.toml").unwrap();
//!
//! // Prepare to run `ls -l /proc/self/fd` with some FDs mapped.
//! let mut command = Command::new("ls");
//! command.arg("-l").arg("/proc/self/fd");
//! command
//!     .fd_mappings(vec![
//!         // Map `file` as FD 3 in the child process.
//!         FdMapping {
//!             parent_fd: file.as_raw_fd(),
//!             child_fd: 3,
//!         },
//!         // Map this process's stdin as FD 5 in the child process.
//!         FdMapping {
//!             parent_fd: 0,
//!             child_fd: 5,
//!         },
//!     ])
//!     .unwrap();
//!
//! // Spawn the child process.
//! let mut child = command.spawn().unwrap();
//! child.wait().unwrap();
//! ```

use nix::fcntl::{fcntl, FcntlArg, FdFlag};
use nix::unistd::dup2;
use std::cmp::max;
use std::io::{self, ErrorKind};
use std::os::unix::io::RawFd;
use std::os::unix::process::CommandExt;
use std::process::Command;
use thiserror::Error;

/// A mapping from a file descriptor in the parent to a file descriptor in the child, to be applied
/// when spawning a child process.
///
/// The parent_fd must be kept open until after the child is spawned.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FdMapping {
    pub parent_fd: RawFd,
    pub child_fd: RawFd,
}

/// Error setting up FD mappings, because there were two or more mappings for the same child FD.
#[derive(Copy, Clone, Debug, Eq, Error, PartialEq)]
#[error("Two or more mappings for the same child FD")]
pub struct FdMappingCollision;

/// Extension to add file descriptor mappings to a [`Command`].
pub trait CommandFdExt {
    /// Adds the given set of file descriptors to the command.
    ///
    /// Warning: Calling this more than once on the same command, or attempting to run the same
    /// command more than once after calling this, may result in unexpected behaviour.
    fn fd_mappings(&mut self, mappings: Vec<FdMapping>) -> Result<&mut Self, FdMappingCollision>;

    /// Adds the given set of file descriptors to be passed on to the child process when the command
    /// is run.
    fn preserved_fds(&mut self, fds: Vec<RawFd>) -> &mut Self;
}

impl CommandFdExt for Command {
    fn fd_mappings(
        &mut self,
        mut mappings: Vec<FdMapping>,
    ) -> Result<&mut Self, FdMappingCollision> {
        // Validate that there are no conflicting mappings to the same child FD.
        let mut child_fds: Vec<RawFd> = mappings.iter().map(|mapping| mapping.child_fd).collect();
        child_fds.sort_unstable();
        child_fds.dedup();
        if child_fds.len() != mappings.len() {
            return Err(FdMappingCollision);
        }

        // Register the callback to apply the mappings after forking but before execing.
        // Safety: `map_fds` will not allocate, so it is safe to call from this hook.
        unsafe {
            // If the command is run more than once, and hence this closure is called multiple
            // times, then `mappings` may be in an incorrect state. It would be good if we could
            // reset it to the initial state somehow, or use something else for saving the temporary
            // mappings.
            self.pre_exec(move || map_fds(&mut mappings, &child_fds));
        }

        Ok(self)
    }

    fn preserved_fds(&mut self, fds: Vec<RawFd>) -> &mut Self {
        unsafe {
            self.pre_exec(move || preserve_fds(&fds));
        }

        self
    }
}

// This function must not do any allocation, as it is called from the pre_exec hook.
fn map_fds(mappings: &mut [FdMapping], child_fds: &[RawFd]) -> io::Result<()> {
    if mappings.is_empty() {
        // No need to do anything, and finding first_unused_fd would fail.
        return Ok(());
    }

    // Find the first FD which is higher than any parent or child FD in the mapping, so we can
    // safely use it and higher FDs as temporary FDs. There may be other files open with these FDs,
    // so we still need to ensure we don't conflict with them.
    let first_safe_fd = mappings
        .iter()
        .map(|mapping| max(mapping.parent_fd, mapping.child_fd))
        .max()
        .unwrap()
        + 1;

    // If any parent FDs conflict with child FDs, then first duplicate them to a temporary FD which
    // is clear of either range. Mappings to the same FD are fine though, we can handle them by just
    // removing the FD_CLOEXEC flag from the existing (parent) FD.
    for mapping in mappings.iter_mut() {
        if child_fds.contains(&mapping.parent_fd) && mapping.parent_fd != mapping.child_fd {
            mapping.parent_fd = fcntl(mapping.parent_fd, FcntlArg::F_DUPFD_CLOEXEC(first_safe_fd))
                .map_err(nix_to_io_error)?;
        }
    }

    // Now we can actually duplicate FDs to the desired child FDs.
    for mapping in mappings {
        if mapping.child_fd == mapping.parent_fd {
            // Remove the FD_CLOEXEC flag, so the FD will be kept open when exec is called for the
            // child.
            fcntl(mapping.parent_fd, FcntlArg::F_SETFD(FdFlag::empty()))
                .map_err(nix_to_io_error)?;
        } else {
            // This closes child_fd if it is already open as something else, and clears the
            // FD_CLOEXEC flag on child_fd.
            dup2(mapping.parent_fd, mapping.child_fd).map_err(nix_to_io_error)?;
        }
    }

    Ok(())
}

fn preserve_fds(fds: &[RawFd]) -> io::Result<()> {
    for fd in fds {
        // Remove the FD_CLOEXEC flag, so the FD will be kept open when exec is called for the
        // child.
        fcntl(*fd, FcntlArg::F_SETFD(FdFlag::empty())).map_err(nix_to_io_error)?;
    }

    Ok(())
}

/// Convert a [`nix::Error`] to a [`std::io::Error`].
fn nix_to_io_error(error: nix::Error) -> io::Error {
    if let nix::Error::Sys(errno) = error {
        io::Error::from_raw_os_error(errno as i32)
    } else {
        io::Error::new(ErrorKind::Other, error)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use nix::unistd::close;
    use std::collections::HashSet;
    use std::fs::{read_dir, File};
    use std::os::unix::io::AsRawFd;
    use std::process::Output;
    use std::str;
    use std::sync::Once;

    static SETUP: Once = Once::new();

    #[test]
    fn conflicting_mappings() {
        setup();

        let mut command = Command::new("ls");

        // The same mapping can't be included twice.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
            ])
            .is_err());

        // Mapping two different FDs to the same FD isn't allowed either.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    child_fd: 4,
                    parent_fd: 5,
                },
                FdMapping {
                    child_fd: 4,
                    parent_fd: 6,
                },
            ])
            .is_err());
    }

    #[test]
    fn no_mappings() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        assert!(command.fd_mappings(vec![]).is_ok());

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3], 0);
    }

    #[test]
    fn none_preserved() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        command.preserved_fds(vec![]);

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3], 0);
    }

    #[test]
    fn one_mapping() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file = File::open("testdata/file1.txt").unwrap();
        // Map the file an otherwise unused FD.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: file.as_raw_fd(),
                child_fd: 5,
            },])
            .is_ok());

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3, 5], 0);
    }

    #[test]
    fn one_preserved() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file = File::open("testdata/file1.txt").unwrap();
        let file_fd = file.as_raw_fd();
        command.preserved_fds(vec![file_fd]);

        let output = command.output().unwrap();
        expect_fds(&output, &[0, 1, 2, 3, file_fd], 0);
    }

    #[test]
    fn swap_mappings() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file1 = File::open("testdata/file1.txt").unwrap();
        let file2 = File::open("testdata/file2.txt").unwrap();
        let fd1 = file1.as_raw_fd();
        let fd2 = file2.as_raw_fd();
        // Map files to each other's FDs, to ensure that the temporary FD logic works.
        assert!(command
            .fd_mappings(vec![
                FdMapping {
                    parent_fd: fd1,
                    child_fd: fd2,
                },
                FdMapping {
                    parent_fd: fd2,
                    child_fd: fd1,
                },
            ])
            .is_ok(),);

        let output = command.output().unwrap();
        // Expect one more Fd for the /proc/self/fd directory. We can't predict what number it will
        // be assigned, because 3 might or might not be taken already by fd1 or fd2.
        expect_fds(&output, &[0, 1, 2, fd1, fd2], 1);
    }

    #[test]
    fn one_to_one_mapping() {
        setup();

        let mut command = Command::new("ls");
        command.arg("/proc/self/fd");

        let file1 = File::open("testdata/file1.txt").unwrap();
        let file2 = File::open("testdata/file2.txt").unwrap();
        let fd1 = file1.as_raw_fd();
        // Map file1 to the same FD it currently has, to ensure the special case for that works.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: fd1,
                child_fd: fd1,
            }])
            .is_ok());

        let output = command.output().unwrap();
        // Expect one more Fd for the /proc/self/fd directory. We can't predict what number it will
        // be assigned, because 3 might or might not be taken already by fd1 or fd2.
        expect_fds(&output, &[0, 1, 2, fd1], 1);

        // Keep file2 open until the end, to ensure that it's not passed to the child.
        drop(file2);
    }

    #[test]
    fn map_stdin() {
        setup();

        let mut command = Command::new("cat");

        let file = File::open("testdata/file1.txt").unwrap();
        // Map the file to stdin.
        assert!(command
            .fd_mappings(vec![FdMapping {
                parent_fd: file.as_raw_fd(),
                child_fd: 0,
            },])
            .is_ok());

        let output = command.output().unwrap();
        assert!(output.status.success());
        assert_eq!(output.stdout, b"test 1");
    }

    /// Parse the output of ls into a set of filenames
    fn parse_ls_output(output: &[u8]) -> HashSet<String> {
        str::from_utf8(output)
            .unwrap()
            .split_terminator("\n")
            .map(str::to_owned)
            .collect()
    }

    /// Check that the output of `ls /proc/self/fd` contains the expected set of FDs, plus exactly
    /// `extra` extra FDs.
    fn expect_fds(output: &Output, expected_fds: &[RawFd], extra: usize) {
        assert!(output.status.success());
        let expected_fds: HashSet<String> = expected_fds.iter().map(RawFd::to_string).collect();
        let fds = parse_ls_output(&output.stdout);
        if extra == 0 {
            assert_eq!(fds, expected_fds);
        } else {
            assert!(expected_fds.is_subset(&fds));
            assert_eq!(fds.len(), expected_fds.len() + extra);
        }
    }

    fn setup() {
        SETUP.call_once(close_excess_fds);
    }

    /// Close all file descriptors apart from stdin, stdout and stderr.
    ///
    /// This is necessary because GitHub Actions opens a bunch of others for some reason.
    fn close_excess_fds() {
        let dir = read_dir("/proc/self/fd").unwrap();
        for entry in dir {
            let entry = entry.unwrap();
            let fd: RawFd = entry.file_name().to_str().unwrap().parse().unwrap();
            if fd > 3 {
                close(fd).unwrap();
            }
        }
    }
}
<|endoftext|>


<commit_before>//! The `sentinel_list` crate provides a simple sentinel based list implementation.
//! Once you insert element to the list you are given back a handle into this list.
//! This handle allows you to remove element from the list in O(1) time.
//!
//! # Examples
//!
//! ```rust
//! extern crate sentinel_list;
//! use sentinel_list::*;
//!
//! fn main() {
//!	 let l = &mut List::new();
//!	 let h1 = l.push_head(1);
//!	 let h2 = l.push_tail(2);
//!	 let h3 = l.push_tail(3);
//!	 {
//!		assert_eq!(&1, (&h1).as_ref());
//!		assert_eq!(&1, l.peek_head().unwrap());
//!		assert_eq!(&3, l.peek_tail().unwrap());
//!	 }
//!	 h1.unlink();
//!	 assert_eq!(&2, l.peek_head().unwrap());
//!	 h3.unlink();
//!	 assert_eq!(&2, l.peek_tail().unwrap());
//!	 h2.unlink();
//!	 assert_eq!(None, l.peek_tail());
//!	 assert_eq!(None, l.peek_head());
//! }
//! ```

#![feature(conservative_impl_trait)]

use std::ptr;
use std::fmt;
use std::ops::Deref;

pub trait ListHandle<T>
{
	fn unlink(self);
	fn as_ref(&self) -> &T;
}

pub struct List<T>
{
	sentinel: Handle<T>,
}

impl<T> List<T>
{
	pub fn new() -> Self
	{
		List {
			sentinel: new_sentinel(),
		}
	}

	pub fn push_head(&mut self, e: T) -> impl ListHandle<T>
	{
		let mut h = new_handle(e);
		insert_after(&mut self.sentinel, &mut h);
		h
	}

	pub fn push_tail(&mut self, e: T) -> impl ListHandle<T>
	{
		let mut h = new_handle(e);
		insert_after(unsafe {&mut *self.sentinel.prev}, &mut h);
		h
	}

	pub fn peek_head(&self) -> Option<&T>
	{
		let link = unsafe { &*self.sentinel.next };
		link.value.as_ref()
	}

	pub fn peek_tail(&self) -> Option<&T>
	{
		let link = unsafe { &*self.sentinel.prev };
		link.value.as_ref()
	}
}

#[derive(PartialEq)]
struct Link<T>
{
	pub next: *mut Link<T>,
	pub prev: *mut Link<T>,
	pub value: Option<T>,
}

type Handle<T> = Box<Link<T>>;

impl<T> Link<T>
{
	fn new(v: T) -> Link<T>
	{
		Link {
			prev: ptr::null_mut(),
			next: ptr::null_mut(),
			value: Some(v),
		}
	}
}

impl<T: fmt::Debug> fmt::Debug for Link<T>
{
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		write!(f, "Link {{ prev: {:p}, next: {:p}, value: {:?} }}", self.prev, self.next, self.value)
	}
}

impl<T> ListHandle<T> for Handle<T>
{
	fn unlink(self)
	{
		unlink(self);
	}

	fn as_ref(&self) -> &T
	{
		&self
	}
}

impl<T> Deref for Link<T>
{
	type Target = T;
	fn deref(&self) -> &Self::Target
	{
		// if you call it on the sentinel
		// you deserve to get what's coming
		self.value.as_ref().unwrap()
	}
}

fn new_sentinel<T>() -> Handle<T>
{
	let mut h = Box::new(
		Link {
			prev: ptr::null_mut(),
			next: ptr::null_mut(),
			value: None,
		}
	);
	h.prev = &mut *h;
	h.next = &mut *h;
	h
}

fn new_handle<T>(v: T) -> Handle<T>
{
	Box::new(Link::new(v))
}

fn insert_after<T>(after: &mut Link<T>, h: &mut Link<T>)
{
	let n: *mut _ = after.next;

	h.prev = &mut *after;
	h.next = n;

	after.next = &mut *h;

	let n = unsafe { &mut *h.next };
	n.prev = &mut *h;
}

fn unlink<T>(h: Handle<T>)
{
	let prev = unsafe { &mut *h.prev };
	let next = unsafe { &mut *h.next };
	next.prev = prev;
	prev.next = next;
}

#[cfg(test)]
fn debug_print<T: fmt::Debug>(s: &mut Handle<T>)
{
	let mut h = s.next;
	while h != &mut **s {
		let n = unsafe { &*h };
		println!("{:?}", n);
		h = n.next;
	}
}

#[cfg(test)]
#[test]
fn link_unlink()
{
	let mut s = new_sentinel();
	let mut h1 = new_handle(1);
	let mut h2 = new_handle(2);
	let mut h3 = new_handle(3);

	insert_after(&mut s,  &mut h1);
	insert_after(&mut h1, &mut h2);
	insert_after(&mut h2, &mut h3);

	assert_eq!(unsafe {&*s.next}, &*h1);
	assert_eq!(unsafe {&*h1.prev}, &*s);

	assert_eq!(unsafe {&*h1.next}, &*h2);
	assert_eq!(unsafe {&*h2.prev}, &*h1);

	assert_eq!(unsafe {&*h2.next}, &*h3);
	assert_eq!(unsafe {&*h3.prev}, &*h2);

	assert_eq!(unsafe {&*h3.next}, &*s);
	assert_eq!(unsafe {&*s.prev}, &*h3);

	unlink(h2);

	assert_eq!(unsafe {&*s.next}, &*h1);
	assert_eq!(unsafe {&*h1.prev}, &*s);

	assert_eq!(unsafe {&*h1.next}, &*h3);
	assert_eq!(unsafe {&*h3.prev}, &*h1);

	assert_eq!(unsafe {&*h3.next}, &*s);
	assert_eq!(unsafe {&*s.prev}, &*h3);

	let mut h2 = new_handle(2);
	insert_after(&mut h1, &mut h2);

	unlink(h1);
	assert_eq!(unsafe {&*s.next}, &*h2);
	unlink(h2);
	assert_eq!(unsafe {&*s.next}, &*h3);
	assert_eq!(unsafe {&*s.prev}, &*h3);
	assert_eq!(unsafe {&*h3.prev}, &*s);
	assert_eq!(unsafe {&*h3.next}, &*s);
	unlink(h3);
	assert_eq!(unsafe {&*s.prev}, &*s);
	assert_eq!(unsafe {&*s.next}, &*s);

	unlink(s);
}
<commit_msg>Adding return value to unlink<commit_after>//! The `sentinel_list` crate provides a simple sentinel based list implementation.
//! Once you insert element to the list you are given back a handle into this list.
//! This handle allows you to remove element from the list in O(1) time.
//!
//! # Examples
//!
//! ```rust
//! extern crate sentinel_list;
//! use sentinel_list::*;
//!
//! fn main() {
//!	 let l = &mut List::new();
//!	 let h1 = l.push_head(1);
//!	 let h2 = l.push_tail(2);
//!	 let h3 = l.push_tail(3);
//!	 {
//!		assert_eq!(&1, (&h1).as_ref());
//!		assert_eq!(&1, l.peek_head().unwrap());
//!		assert_eq!(&3, l.peek_tail().unwrap());
//!	 }
//!	 let e1 = h1.unlink();
//!	 assert_eq!(&2, l.peek_head().unwrap());
//!	 let e3 = h3.unlink();
//!	 assert_eq!(&2, l.peek_tail().unwrap());
//!	 let e2 = h2.unlink();
//!	 assert_eq!(None, l.peek_tail());
//!	 assert_eq!(None, l.peek_head());
//!	 assert_eq!(1, e1);
//!	 assert_eq!(2, e2);
//!	 assert_eq!(3, e3);
//! }
//! ```

#![feature(conservative_impl_trait)]

use std::ptr;
use std::fmt;
use std::ops::Deref;

pub trait ListHandle<T>
{
	fn unlink(self) -> T;
	fn as_ref(&self) -> &T;
}

pub struct List<T>
{
	sentinel: Handle<T>,
}

impl<T> List<T>
{
	pub fn new() -> Self
	{
		List {
			sentinel: new_sentinel(),
		}
	}

	pub fn push_head(&mut self, e: T) -> impl ListHandle<T>
	{
		let mut h = new_handle(e);
		insert_after(&mut self.sentinel, &mut h);
		h
	}

	pub fn push_tail(&mut self, e: T) -> impl ListHandle<T>
	{
		let mut h = new_handle(e);
		insert_after(unsafe {&mut *self.sentinel.prev}, &mut h);
		h
	}

	pub fn peek_head(&self) -> Option<&T>
	{
		let link = unsafe { &*self.sentinel.next };
		link.value.as_ref()
	}

	pub fn peek_tail(&self) -> Option<&T>
	{
		let link = unsafe { &*self.sentinel.prev };
		link.value.as_ref()
	}
}

#[derive(PartialEq)]
struct Link<T>
{
	pub next: *mut Link<T>,
	pub prev: *mut Link<T>,
	pub value: Option<T>,
}

type Handle<T> = Box<Link<T>>;

impl<T> Link<T>
{
	fn new(v: T) -> Link<T>
	{
		Link {
			prev: ptr::null_mut(),
			next: ptr::null_mut(),
			value: Some(v),
		}
	}
}

impl<T: fmt::Debug> fmt::Debug for Link<T>
{
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		write!(f, "Link {{ prev: {:p}, next: {:p}, value: {:?} }}", self.prev, self.next, self.value)
	}
}

impl<T> ListHandle<T> for Handle<T>
{
	fn unlink(self) -> T
	{
		// if you call it on the sentinel
		// you deserve to get what's coming
		unlink(self).unwrap()
	}

	fn as_ref(&self) -> &T
	{
		&self
	}
}

impl<T> Deref for Link<T>
{
	type Target = T;
	fn deref(&self) -> &Self::Target
	{
		// if you call it on the sentinel
		// you deserve to get what's coming
		self.value.as_ref().unwrap()
	}
}

fn new_sentinel<T>() -> Handle<T>
{
	let mut h = Box::new(
		Link {
			prev: ptr::null_mut(),
			next: ptr::null_mut(),
			value: None,
		}
	);
	h.prev = &mut *h;
	h.next = &mut *h;
	h
}

fn new_handle<T>(v: T) -> Handle<T>
{
	Box::new(Link::new(v))
}

fn insert_after<T>(after: &mut Link<T>, h: &mut Link<T>)
{
	let n: *mut _ = after.next;

	h.prev = &mut *after;
	h.next = n;

	after.next = &mut *h;

	let n = unsafe { &mut *h.next };
	n.prev = &mut *h;
}

fn unlink<T>(h: Handle<T>) -> Option<T>
{
	let prev = unsafe { &mut *h.prev };
	let next = unsafe { &mut *h.next };
	next.prev = prev;
	prev.next = next;
    h.value
}

#[cfg(test)]
fn debug_print<T: fmt::Debug>(s: &mut Handle<T>)
{
	let mut h = s.next;
	while h != &mut **s {
		let n = unsafe { &*h };
		println!("{:?}", n);
		h = n.next;
	}
}

#[cfg(test)]
#[test]
fn link_unlink()
{
	let mut s = new_sentinel();
	let mut h1 = new_handle(1);
	let mut h2 = new_handle(2);
	let mut h3 = new_handle(3);

	insert_after(&mut s,  &mut h1);
	insert_after(&mut h1, &mut h2);
	insert_after(&mut h2, &mut h3);

	assert_eq!(unsafe {&*s.next}, &*h1);
	assert_eq!(unsafe {&*h1.prev}, &*s);

	assert_eq!(unsafe {&*h1.next}, &*h2);
	assert_eq!(unsafe {&*h2.prev}, &*h1);

	assert_eq!(unsafe {&*h2.next}, &*h3);
	assert_eq!(unsafe {&*h3.prev}, &*h2);

	assert_eq!(unsafe {&*h3.next}, &*s);
	assert_eq!(unsafe {&*s.prev}, &*h3);

	unlink(h2);

	assert_eq!(unsafe {&*s.next}, &*h1);
	assert_eq!(unsafe {&*h1.prev}, &*s);

	assert_eq!(unsafe {&*h1.next}, &*h3);
	assert_eq!(unsafe {&*h3.prev}, &*h1);

	assert_eq!(unsafe {&*h3.next}, &*s);
	assert_eq!(unsafe {&*s.prev}, &*h3);

	let mut h2 = new_handle(2);
	insert_after(&mut h1, &mut h2);

	unlink(h1);
	assert_eq!(unsafe {&*s.next}, &*h2);
	unlink(h2);
	assert_eq!(unsafe {&*s.next}, &*h3);
	assert_eq!(unsafe {&*s.prev}, &*h3);
	assert_eq!(unsafe {&*h3.prev}, &*s);
	assert_eq!(unsafe {&*h3.next}, &*s);
	unlink(h3);
	assert_eq!(unsafe {&*s.prev}, &*s);
	assert_eq!(unsafe {&*s.next}, &*s);

	unlink(s);
}
<|endoftext|>


<commit_before>#![cfg_attr(test, deny(warnings))]
#![deny(missing_docs)]

//! A set of middleware for sharing data between requests in the Iron
//! framework.

extern crate iron;
extern crate plugin;

use iron::{Request, Response, BeforeMiddleware, AfterMiddleware, IronResult};
use iron::typemap::Key;
use std::sync::{Arc, RwLock, Mutex};
use std::fmt;
use std::error::Error;
use plugin::Plugin;

/// The type that can be returned by `eval` to indicate error.
#[derive(Clone, Debug)]
pub enum PersistentError {
    /// The value was not found.
    NotFound
}

impl Error for PersistentError {
    fn description(&self) -> &str {
        match *self {
            PersistentError::NotFound => "Value not found in extensions."
        }
    }
}

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

/// Helper trait for overloading the constructors of `Read`/`Write`/`State`.
/// This is an implementation detail, and should not be used for any other
/// purpose.
///
/// For example, this trait lets you construct a `Read<T>` from either a `T` or
/// an `Arc<T>`.
pub trait PersistentInto<T> {
    /// Convert `self` into a value of type `T`.
    fn persistent_into(self) -> T;
}

impl<T> PersistentInto<T> for T {
    fn persistent_into(self) -> T { self }
}

impl<T> PersistentInto<Arc<T>> for T {
    fn persistent_into(self) -> Arc<T> {
        Arc::new(self)
    }
}

impl<T> PersistentInto<Arc<Mutex<T>>> for T {
    fn persistent_into(self) -> Arc<Mutex<T>> {
        Arc::new(Mutex::new(self))
    }
}

impl<T> PersistentInto<Arc<RwLock<T>>> for T {
    fn persistent_into(self) -> Arc<RwLock<T>> {
        Arc::new(RwLock::new(self))
    }
}

/// Middleware for data that persists between requests with read and write capabilities.
///
/// The data is stored behind a `RwLock`, so multiple read locks
/// can be taken out concurrently.
///
/// If most threads need to take out a write lock, you may want to
/// consider `Write`, which stores the data behind a `Mutex`, which
/// has a faster locking speed.
///
/// `State` can be linked as `BeforeMiddleware` to add data to the `Request`
/// extensions and it can be linked as an `AfterMiddleware` to add data to
/// the `Response` extensions.
///
/// `State` also implements `Plugin`, so the data stored within can be
/// accessed through `request.get::<State<P>>()` as an `Arc<RwLock<P::Value>>`.
pub struct State<P: Key> {
    data: Arc<RwLock<P::Value>>
}

/// Middleware for data that persists between Requests with read-only capabilities.
///
/// The data is stored behind an Arc, so multiple threads can have
/// concurrent, non-blocking access.
///
/// `Read` can be linked as `BeforeMiddleware` to add data to the `Request`
/// extensions and it can be linked as an `AfterMiddleware` to add data to
/// the `Response` extensions.
///
/// `Read` also implements `Plugin`, so the data stored within can be
/// accessed through `request.get::<Read<P>>()` as an `Arc<P::Value>`.
pub struct Read<P: Key> {
    data: Arc<P::Value>
}

/// Middleware for data that persists between Requests for data which mostly
/// needs to be written instead of read.
///
/// The data is stored behind a `Mutex`, so only one request at a time can
/// access the data. This is more performant than `State` in the case where
/// most uses of the data require a write lock.
///
/// `Write` can be linked as `BeforeMiddleware` to add data to the `Request`
/// extensions and it can be linked as an `AfterMiddleware` to add data to
/// the `Response` extensions.
///
/// `Write` also implements `Plugin`, so the data stored within can be
/// accessed through `request.get::<Write<P>>()` as an `Arc<Mutex<P::Value>>`.
pub struct Write<P: Key> {
    data: Arc<Mutex<P::Value>>
}

impl<P: Key> Clone for Read<P> where P::Value: Send + Sync {
    fn clone(&self) -> Read<P> {
        Read { data: self.data.clone() }
    }
}

impl<P: Key> Clone for State<P> where P::Value: Send + Sync {
    fn clone(&self) -> State<P> {
        State { data: self.data.clone() }
    }
}

impl<P: Key> Clone for Write<P> where P::Value: Send {
    fn clone(&self) -> Write<P> {
        Write { data: self.data.clone() }
    }
}

impl<P: Key> Key for State<P> where P::Value: 'static {
    type Value = Arc<RwLock<P::Value>>;
}

impl<P: Key> Key for Read<P> where P::Value: 'static {
    type Value = Arc<P::Value>;
}

impl<P: Key> Key for Write<P> where P::Value: 'static {
    type Value = Arc<Mutex<P::Value>>;
}

impl<'a, 'b, P: Key> Plugin<Request<'a, 'b>> for State<P> where P::Value: Send + Sync {
    type Error = PersistentError;
    fn eval(req: &mut Request<'a, 'b>) -> Result<Arc<RwLock<P::Value>>, PersistentError> {
        req.extensions.get::<State<P>>().cloned().ok_or(PersistentError::NotFound)
    }
}

impl<'a, 'b, P: Key> Plugin<Request<'a, 'b>> for Read<P> where P::Value: Send + Sync {
    type Error = PersistentError;
    fn eval(req: &mut Request<'a, 'b>) -> Result<Arc<P::Value>, PersistentError> {
        req.extensions.get::<Read<P>>().cloned().ok_or(PersistentError::NotFound)
    }
}

impl<'a, 'b, P: Key> Plugin<Request<'a, 'b>> for Write<P> where P::Value: Send {
    type Error = PersistentError;
    fn eval(req: &mut Request<'a, 'b>) -> Result<Arc<Mutex<P::Value>>, PersistentError> {
        req.extensions.get::<Write<P>>().cloned().ok_or(PersistentError::NotFound)
    }
}

impl<P: Key> BeforeMiddleware for State<P> where P::Value: Send + Sync {
    fn before(&self, req: &mut Request) -> IronResult<()> {
        req.extensions.insert::<State<P>>(self.data.clone());
        Ok(())
    }
}

impl<P: Key> BeforeMiddleware for Read<P> where P::Value: Send + Sync {
    fn before(&self, req: &mut Request) -> IronResult<()> {
        req.extensions.insert::<Read<P>>(self.data.clone());
        Ok(())
    }
}

impl<P: Key> BeforeMiddleware for Write<P> where P::Value: Send {
    fn before(&self, req: &mut Request) -> IronResult<()> {
        req.extensions.insert::<Write<P>>(self.data.clone());
        Ok(())
    }
}

impl<P: Key> AfterMiddleware for State<P> where P::Value: Send + Sync {
    fn after(&self, _: &mut Request, mut res: Response) -> IronResult<Response> {
        res.extensions.insert::<State<P>>(self.data.clone());
        Ok(res)
    }
}

impl<P: Key> AfterMiddleware for Read<P> where P::Value: Send + Sync {
    fn after(&self, _: &mut Request, mut res: Response) -> IronResult<Response> {
        res.extensions.insert::<Read<P>>(self.data.clone());
        Ok(res)
    }
}

impl<P: Key> AfterMiddleware for Write<P> where P::Value: Send {
    fn after(&self, _: &mut Request, mut res: Response) -> IronResult<Response> {
        res.extensions.insert::<Write<P>>(self.data.clone());
        Ok(res)
    }
}

impl<P: Key> State<P> where P::Value: Send + Sync {
    /// Construct a new pair of `State` that can be passed directly to `Chain::link`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn both<T>(start: T) -> (State<P>, State<P>) where T: PersistentInto<Arc<RwLock<P::Value>>> {
        let x = State { data: start.persistent_into() };
        (x.clone(), x)
    }

    /// Construct a new `State` that can be passed directly to
    /// `Chain::link_before` or `Chain::link_after`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn one<T>(start: P::Value) -> State<P> where T: PersistentInto<Arc<RwLock<P::Value>>> {
        State { data: start.persistent_into() }
    }
}

impl<P: Key> Read<P> where P::Value: Send + Sync {
    /// Construct a new pair of `Read` that can be passed directly to `Chain::link`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn both<T>(start: T) -> (Read<P>, Read<P>) where T: PersistentInto<Arc<P::Value>> {
        let x = Read { data: start.persistent_into() };
        (x.clone(), x)
    }

    /// Construct a new `Read` that can be passed directly to
    /// `Chain::link_before` or `Chain::link_after`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn one<T>(start: T) -> Read<P> where T: PersistentInto<Arc<P::Value>> {
        Read { data: start.persistent_into() }
    }
}

impl<P: Key> Write<P> where P::Value: Send {
    /// Construct a new pair of `Write` that can be passed directly to `Chain::link`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn both<T>(start: T) -> (Write<P>, Write<P>) where T: PersistentInto<Arc<Mutex<P::Value>>> {
        let x = Write { data: start.persistent_into() };
        (x.clone(), x)
    }

    /// Construct a new `Write` that can be passed directly to
    /// `Chain::link_before` or `Chain::link_after`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn one<T>(start: T) -> Write<P> where T: PersistentInto<Arc<Mutex<P::Value>>> {
        Write { data: start.persistent_into() }
    }
}
<commit_msg>fix typo in State::one parameter<commit_after>#![cfg_attr(test, deny(warnings))]
#![deny(missing_docs)]

//! A set of middleware for sharing data between requests in the Iron
//! framework.

extern crate iron;
extern crate plugin;

use iron::{Request, Response, BeforeMiddleware, AfterMiddleware, IronResult};
use iron::typemap::Key;
use std::sync::{Arc, RwLock, Mutex};
use std::fmt;
use std::error::Error;
use plugin::Plugin;

/// The type that can be returned by `eval` to indicate error.
#[derive(Clone, Debug)]
pub enum PersistentError {
    /// The value was not found.
    NotFound
}

impl Error for PersistentError {
    fn description(&self) -> &str {
        match *self {
            PersistentError::NotFound => "Value not found in extensions."
        }
    }
}

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

/// Helper trait for overloading the constructors of `Read`/`Write`/`State`.
/// This is an implementation detail, and should not be used for any other
/// purpose.
///
/// For example, this trait lets you construct a `Read<T>` from either a `T` or
/// an `Arc<T>`.
pub trait PersistentInto<T> {
    /// Convert `self` into a value of type `T`.
    fn persistent_into(self) -> T;
}

impl<T> PersistentInto<T> for T {
    fn persistent_into(self) -> T { self }
}

impl<T> PersistentInto<Arc<T>> for T {
    fn persistent_into(self) -> Arc<T> {
        Arc::new(self)
    }
}

impl<T> PersistentInto<Arc<Mutex<T>>> for T {
    fn persistent_into(self) -> Arc<Mutex<T>> {
        Arc::new(Mutex::new(self))
    }
}

impl<T> PersistentInto<Arc<RwLock<T>>> for T {
    fn persistent_into(self) -> Arc<RwLock<T>> {
        Arc::new(RwLock::new(self))
    }
}

/// Middleware for data that persists between requests with read and write capabilities.
///
/// The data is stored behind a `RwLock`, so multiple read locks
/// can be taken out concurrently.
///
/// If most threads need to take out a write lock, you may want to
/// consider `Write`, which stores the data behind a `Mutex`, which
/// has a faster locking speed.
///
/// `State` can be linked as `BeforeMiddleware` to add data to the `Request`
/// extensions and it can be linked as an `AfterMiddleware` to add data to
/// the `Response` extensions.
///
/// `State` also implements `Plugin`, so the data stored within can be
/// accessed through `request.get::<State<P>>()` as an `Arc<RwLock<P::Value>>`.
pub struct State<P: Key> {
    data: Arc<RwLock<P::Value>>
}

/// Middleware for data that persists between Requests with read-only capabilities.
///
/// The data is stored behind an Arc, so multiple threads can have
/// concurrent, non-blocking access.
///
/// `Read` can be linked as `BeforeMiddleware` to add data to the `Request`
/// extensions and it can be linked as an `AfterMiddleware` to add data to
/// the `Response` extensions.
///
/// `Read` also implements `Plugin`, so the data stored within can be
/// accessed through `request.get::<Read<P>>()` as an `Arc<P::Value>`.
pub struct Read<P: Key> {
    data: Arc<P::Value>
}

/// Middleware for data that persists between Requests for data which mostly
/// needs to be written instead of read.
///
/// The data is stored behind a `Mutex`, so only one request at a time can
/// access the data. This is more performant than `State` in the case where
/// most uses of the data require a write lock.
///
/// `Write` can be linked as `BeforeMiddleware` to add data to the `Request`
/// extensions and it can be linked as an `AfterMiddleware` to add data to
/// the `Response` extensions.
///
/// `Write` also implements `Plugin`, so the data stored within can be
/// accessed through `request.get::<Write<P>>()` as an `Arc<Mutex<P::Value>>`.
pub struct Write<P: Key> {
    data: Arc<Mutex<P::Value>>
}

impl<P: Key> Clone for Read<P> where P::Value: Send + Sync {
    fn clone(&self) -> Read<P> {
        Read { data: self.data.clone() }
    }
}

impl<P: Key> Clone for State<P> where P::Value: Send + Sync {
    fn clone(&self) -> State<P> {
        State { data: self.data.clone() }
    }
}

impl<P: Key> Clone for Write<P> where P::Value: Send {
    fn clone(&self) -> Write<P> {
        Write { data: self.data.clone() }
    }
}

impl<P: Key> Key for State<P> where P::Value: 'static {
    type Value = Arc<RwLock<P::Value>>;
}

impl<P: Key> Key for Read<P> where P::Value: 'static {
    type Value = Arc<P::Value>;
}

impl<P: Key> Key for Write<P> where P::Value: 'static {
    type Value = Arc<Mutex<P::Value>>;
}

impl<'a, 'b, P: Key> Plugin<Request<'a, 'b>> for State<P> where P::Value: Send + Sync {
    type Error = PersistentError;
    fn eval(req: &mut Request<'a, 'b>) -> Result<Arc<RwLock<P::Value>>, PersistentError> {
        req.extensions.get::<State<P>>().cloned().ok_or(PersistentError::NotFound)
    }
}

impl<'a, 'b, P: Key> Plugin<Request<'a, 'b>> for Read<P> where P::Value: Send + Sync {
    type Error = PersistentError;
    fn eval(req: &mut Request<'a, 'b>) -> Result<Arc<P::Value>, PersistentError> {
        req.extensions.get::<Read<P>>().cloned().ok_or(PersistentError::NotFound)
    }
}

impl<'a, 'b, P: Key> Plugin<Request<'a, 'b>> for Write<P> where P::Value: Send {
    type Error = PersistentError;
    fn eval(req: &mut Request<'a, 'b>) -> Result<Arc<Mutex<P::Value>>, PersistentError> {
        req.extensions.get::<Write<P>>().cloned().ok_or(PersistentError::NotFound)
    }
}

impl<P: Key> BeforeMiddleware for State<P> where P::Value: Send + Sync {
    fn before(&self, req: &mut Request) -> IronResult<()> {
        req.extensions.insert::<State<P>>(self.data.clone());
        Ok(())
    }
}

impl<P: Key> BeforeMiddleware for Read<P> where P::Value: Send + Sync {
    fn before(&self, req: &mut Request) -> IronResult<()> {
        req.extensions.insert::<Read<P>>(self.data.clone());
        Ok(())
    }
}

impl<P: Key> BeforeMiddleware for Write<P> where P::Value: Send {
    fn before(&self, req: &mut Request) -> IronResult<()> {
        req.extensions.insert::<Write<P>>(self.data.clone());
        Ok(())
    }
}

impl<P: Key> AfterMiddleware for State<P> where P::Value: Send + Sync {
    fn after(&self, _: &mut Request, mut res: Response) -> IronResult<Response> {
        res.extensions.insert::<State<P>>(self.data.clone());
        Ok(res)
    }
}

impl<P: Key> AfterMiddleware for Read<P> where P::Value: Send + Sync {
    fn after(&self, _: &mut Request, mut res: Response) -> IronResult<Response> {
        res.extensions.insert::<Read<P>>(self.data.clone());
        Ok(res)
    }
}

impl<P: Key> AfterMiddleware for Write<P> where P::Value: Send {
    fn after(&self, _: &mut Request, mut res: Response) -> IronResult<Response> {
        res.extensions.insert::<Write<P>>(self.data.clone());
        Ok(res)
    }
}

impl<P: Key> State<P> where P::Value: Send + Sync {
    /// Construct a new pair of `State` that can be passed directly to `Chain::link`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn both<T>(start: T) -> (State<P>, State<P>) where T: PersistentInto<Arc<RwLock<P::Value>>> {
        let x = State { data: start.persistent_into() };
        (x.clone(), x)
    }

    /// Construct a new `State` that can be passed directly to
    /// `Chain::link_before` or `Chain::link_after`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn one<T>(start: T) -> State<P> where T: PersistentInto<Arc<RwLock<P::Value>>> {
        State { data: start.persistent_into() }
    }
}

impl<P: Key> Read<P> where P::Value: Send + Sync {
    /// Construct a new pair of `Read` that can be passed directly to `Chain::link`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn both<T>(start: T) -> (Read<P>, Read<P>) where T: PersistentInto<Arc<P::Value>> {
        let x = Read { data: start.persistent_into() };
        (x.clone(), x)
    }

    /// Construct a new `Read` that can be passed directly to
    /// `Chain::link_before` or `Chain::link_after`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn one<T>(start: T) -> Read<P> where T: PersistentInto<Arc<P::Value>> {
        Read { data: start.persistent_into() }
    }
}

impl<P: Key> Write<P> where P::Value: Send {
    /// Construct a new pair of `Write` that can be passed directly to `Chain::link`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn both<T>(start: T) -> (Write<P>, Write<P>) where T: PersistentInto<Arc<Mutex<P::Value>>> {
        let x = Write { data: start.persistent_into() };
        (x.clone(), x)
    }

    /// Construct a new `Write` that can be passed directly to
    /// `Chain::link_before` or `Chain::link_after`.
    ///
    /// The data is initialized with the passed-in value.
    pub fn one<T>(start: T) -> Write<P> where T: PersistentInto<Arc<Mutex<P::Value>>> {
        Write { data: start.persistent_into() }
    }
}
<|endoftext|>


<commit_before>use super::LENGTH_BYTES;

pub trait Encodable {
    fn ssz_append(&self, s: &mut SszStream);
}

/// Provides a buffer for appending ssz-encodable values.
///
/// Use the `append()` fn to add a value to a list, then use
/// the `drain()` method to consume the struct and return the
/// ssz encoded bytes.
pub struct SszStream {
    buffer: Vec<u8>
}

impl SszStream {
    /// Create a new, empty steam for writing ssz values.
    pub fn new() -> Self {
        SszStream {
            buffer: Vec::new()
        }
    }

    /// Append some ssz encodable value to the stream.
    pub fn append<E>(&mut self, value: &E) -> &mut Self
        where E: Encodable
    {
        value.ssz_append(self);
        self
    }

    /// Append some ssz encoded bytes to the steam.
    ///
    /// The length of the supplied bytes will be concatenated
    /// to the stream before the supplied bytes.
    pub fn append_encoded_val(&mut self, vec: &Vec<u8>) {
        self.buffer.extend_from_slice(
            &encode_length(vec.len(),
            LENGTH_BYTES));
        self.buffer.extend_from_slice(&vec);
    }

    /// Append some vector (list) of encodable values to the stream.
    ///
    /// The length of the list will be concatenated to the stream, then
    /// each item in the vector will be encoded and concatenated.
    pub fn append_vec<E>(&mut self, vec: &Vec<E>)
        where E: Encodable
    {
        self.buffer.extend_from_slice(&encode_length(vec.len(), LENGTH_BYTES));
        for v in vec {
            v.ssz_append(self);
        }
    }

    /// Consume the stream and return the underlying bytes.
    pub fn drain(self) -> Vec<u8> {
        self.buffer
    }
}

/// Encode some length into a ssz size prefix.
///
/// The ssz size prefix is 4 bytes, which is treated as a continuious
/// 32bit big-endian integer.
pub fn encode_length(len: usize, length_bytes: usize) -> Vec<u8> {
    assert!(length_bytes > 0);  // For sanity
    assert!((len as usize) < 2usize.pow(length_bytes as u32 * 8));
    let mut header: Vec<u8> = vec![0; length_bytes];
    for i in 0..length_bytes {
        let offset = (length_bytes - i - 1) * 8;
        header[i] = ((len >> offset) & 0xff) as u8;
    };
    header
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::ethereum_types::{ H256, U256 };

    #[test]
    #[should_panic]
    fn test_encode_length_0_bytes_panic() {
        encode_length(0, 0);
    }

    #[test]
    fn test_encode_length_4_bytes() {
        assert_eq!(
            encode_length(0, LENGTH_BYTES),
            vec![0; 3]
        );
        assert_eq!(
            encode_length(1, LENGTH_BYTES),
            vec![0, 0, 1]
        );
        assert_eq!(
            encode_length(255, LENGTH_BYTES),
            vec![0, 0, 255]
        );
        assert_eq!(
            encode_length(256, LENGTH_BYTES),
            vec![0, 1, 0]
        );
        assert_eq!(
            encode_length(16777215, LENGTH_BYTES),  // 2^(3*8) - 1
            vec![255, 255, 255]
        );
    }

    #[test]
    #[should_panic]
    fn test_encode_length_4_bytes_panic() {
        encode_length(16777216, LENGTH_BYTES);  // 2^(3*8)
    }

    /*
    #[test]
    fn test_encode_struct() {
        pub struct TestStruct {
            pub one: u32,
            pub two: H256,
            pub three: u64,
            pub four: U256,
        }

        impl Encodable for TestStruct {
            fn ssz_append(&self, s: &mut SszStream) {
                s.append(&self.one);
                s.append(&self.two);
                s.append(&self.three);
                s.append(&self.four);
            }
        }

        let t = TestStruct {
            one: 1,
            two: H256::zero(),
            three: 100,
            four: U256::zero(),
        };

        let mut s = SszStream::new();
        s.append(&t);
        let e = s.drain();

        let expected_len = {
            3 + 4 +
            3 + 32 +
            3 + 8 +
            3 + 32
        };

        assert_eq!(e[0..4], [0, 0, 0, 4]);
        assert_eq!(e[4..8], [0, 0, 0, 1]);
        assert_eq!(e[8..12], [0, 0, 0, 32]);
        assert_eq!(e[12..44], [0; 32]);
        assert_eq!(e[44..48], [0, 0, 0, 8]);
        assert_eq!(e[48..56], [0, 0, 0, 0, 0, 0, 0, 100]);
        assert_eq!(e[56..60], [0, 0, 0, 32]);
        assert_eq!(e[60..92], [0; 32]);
        assert_eq!(e.len(), expected_len);
    }
    */
}
<commit_msg>Fix ssz typos and unused import<commit_after>use super::LENGTH_BYTES;

pub trait Encodable {
    fn ssz_append(&self, s: &mut SszStream);
}

/// Provides a buffer for appending ssz-encodable values.
///
/// Use the `append()` fn to add a value to a list, then use
/// the `drain()` method to consume the struct and return the
/// ssz encoded bytes.
pub struct SszStream {
    buffer: Vec<u8>
}

impl SszStream {
    /// Create a new, empty stream for writing ssz values.
    pub fn new() -> Self {
        SszStream {
            buffer: Vec::new()
        }
    }

    /// Append some ssz encodable value to the stream.
    pub fn append<E>(&mut self, value: &E) -> &mut Self
        where E: Encodable
    {
        value.ssz_append(self);
        self
    }

    /// Append some ssz encoded bytes to the stream.
    ///
    /// The length of the supplied bytes will be concatenated
    /// to the stream before the supplied bytes.
    pub fn append_encoded_val(&mut self, vec: &Vec<u8>) {
        self.buffer.extend_from_slice(
            &encode_length(vec.len(),
            LENGTH_BYTES));
        self.buffer.extend_from_slice(&vec);
    }

    /// Append some vector (list) of encodable values to the stream.
    ///
    /// The length of the list will be concatenated to the stream, then
    /// each item in the vector will be encoded and concatenated.
    pub fn append_vec<E>(&mut self, vec: &Vec<E>)
        where E: Encodable
    {
        self.buffer.extend_from_slice(&encode_length(vec.len(), LENGTH_BYTES));
        for v in vec {
            v.ssz_append(self);
        }
    }

    /// Consume the stream and return the underlying bytes.
    pub fn drain(self) -> Vec<u8> {
        self.buffer
    }
}

/// Encode some length into a ssz size prefix.
///
/// The ssz size prefix is 4 bytes, which is treated as a continuious
/// 32bit big-endian integer.
pub fn encode_length(len: usize, length_bytes: usize) -> Vec<u8> {
    assert!(length_bytes > 0);  // For sanity
    assert!((len as usize) < 2usize.pow(length_bytes as u32 * 8));
    let mut header: Vec<u8> = vec![0; length_bytes];
    for i in 0..length_bytes {
        let offset = (length_bytes - i - 1) * 8;
        header[i] = ((len >> offset) & 0xff) as u8;
    };
    header
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    #[should_panic]
    fn test_encode_length_0_bytes_panic() {
        encode_length(0, 0);
    }

    #[test]
    fn test_encode_length_4_bytes() {
        assert_eq!(
            encode_length(0, LENGTH_BYTES),
            vec![0; 3]
        );
        assert_eq!(
            encode_length(1, LENGTH_BYTES),
            vec![0, 0, 1]
        );
        assert_eq!(
            encode_length(255, LENGTH_BYTES),
            vec![0, 0, 255]
        );
        assert_eq!(
            encode_length(256, LENGTH_BYTES),
            vec![0, 1, 0]
        );
        assert_eq!(
            encode_length(16777215, LENGTH_BYTES),  // 2^(3*8) - 1
            vec![255, 255, 255]
        );
    }

    #[test]
    #[should_panic]
    fn test_encode_length_4_bytes_panic() {
        encode_length(16777216, LENGTH_BYTES);  // 2^(3*8)
    }

    /*
    #[test]
    fn test_encode_struct() {
        pub struct TestStruct {
            pub one: u32,
            pub two: H256,
            pub three: u64,
            pub four: U256,
        }

        impl Encodable for TestStruct {
            fn ssz_append(&self, s: &mut SszStream) {
                s.append(&self.one);
                s.append(&self.two);
                s.append(&self.three);
                s.append(&self.four);
            }
        }

        let t = TestStruct {
            one: 1,
            two: H256::zero(),
            three: 100,
            four: U256::zero(),
        };

        let mut s = SszStream::new();
        s.append(&t);
        let e = s.drain();

        let expected_len = {
            3 + 4 +
            3 + 32 +
            3 + 8 +
            3 + 32
        };

        assert_eq!(e[0..4], [0, 0, 0, 4]);
        assert_eq!(e[4..8], [0, 0, 0, 1]);
        assert_eq!(e[8..12], [0, 0, 0, 32]);
        assert_eq!(e[12..44], [0; 32]);
        assert_eq!(e[44..48], [0, 0, 0, 8]);
        assert_eq!(e[48..56], [0, 0, 0, 0, 0, 0, 0, 100]);
        assert_eq!(e[56..60], [0, 0, 0, 32]);
        assert_eq!(e[60..92], [0; 32]);
        assert_eq!(e.len(), expected_len);
    }
    */
}
<|endoftext|>