<commit_before>use std::collections::HashMap;
use serialize::json;
use serialize::json::{ToJson};

// masks for pulling nibbles out of bytes
const BYTE_WIDTH: uint = 8;
const U4_LOW:     u8   = 0b00001111;
const U4_HIGH:    u8   = 0b11110000;

/// If the CPU enters the `Halt` state then any additional ticks will result in
/// the program exhibiting undefined behavior.
///
#[deriving(PartialEq,Show)]
pub enum CpuState {
	Continue,
	Halt,
}

/// The P150 virtual machine
/// Represents all information the CPU needs to continue executing a
/// program stored in its main memory.
pub struct P150Cpu {
	ip:  u8,
	ir: u16,

	reg: [u8, ..16],
	mem: [u8, ..256],
}

impl P150Cpu {
	/// Initializes the P150 CPU
	///
	/// NOTE: This implementation 0s memory; but this is not guaranteed
	/// by the machine specification.
	///
	/// Programs MUST start at memory address 0x00
	pub fn new() -> P150Cpu {
		P150Cpu {
			ip:  0x00,
			ir:  0x0000,

			reg: [0u8, ..16],
			mem: [0u8, ..256],
		}
	}

	pub fn js_dump(&self) -> json::Json {
		let mut dict = HashMap::new();

		dict.insert(format!("ip"), self.ip.to_json());
		dict.insert(format!("ir"), self.ir.to_json());
		dict.insert(format!("reg"), self.reg.to_json());
		dict.insert(format!("mem"), self.mem.to_json());

		return dict.to_json();
	}

	/// Prints the current machine state to the console window
	/// This includes the IP, IR, and all registers.
	/// (Registers will be formatted as 2s complement numbers.)
	pub fn dump(&self) {
		println!("IP: 0x{:02X}, IR: 0x{:04X}", self.ip, self.ir);

		println!("---\nRegisters\n---")
		for (addr, cell) in self.reg.iter().enumerate() {
			println!("{:01X}: {}", addr, cell)
		}

		println!("---")

		debug!("Memory")
		for (addr, cell) in self.mem.iter().enumerate() {
			debug!("{:02X}: {}", addr, cell)
		}
	}

	#[cfg(test)]
	fn get_reg(&self) -> &[u8] { self.reg }

	/// Read an array of instructions into main memory
	/// This reads two bytes at a time from the `memory` array
	/// and loads them into the P150s RAM bank, starting from address 0.
	pub fn init_mem(&mut self, memory: &[u16]) {
		assert!(memory.len() < (256 / 2)); // program cannot be larger than memory
		let mut next_cell = 0x00;

		// zero memory
		self.mem = [0, ..256];
		for op in memory.iter() {
			let byte_1 = (*op >> 8) as u8;
			let byte_2 = *op as u8;

			self.mem[next_cell]     = byte_1;
			self.mem[next_cell + 1] = byte_2;
			next_cell += 2
		}
	}

	/// Runs the entire fetch, execute, decode cycle against the current IP
	pub fn tick(&mut self) -> CpuState {
		self.fetch();

		// decode & execute
		let op = (self.ir >> 12) as u8;
		match op {
			0x0   => { // MLOAD
				let rloc_o0 = lo_nibble((self.ir >> 8) as u8);
				let mloc_i0 = self.ir as u8;

				self.reg[rloc_o0 as uint] = self.mem[mloc_i0 as uint];
				Continue
			},

			0x1   => { // RSET
				let rloc = lo_nibble((self.ir >> 8) as u8);  // lower nibble of first byte
				let rval = self.ir as u8;                    // value is entire second byte

				self.reg[rloc as uint] = rval;               // store value in register
				Continue
			},

			0x2   => { // MSTOR
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let mloc_o0 = self.ir as u8;

				self.mem[mloc_o0 as uint] = self.reg[rloc_i0 as uint];
				Continue
			},

			0x3   => { // RMOV
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let rloc_o0 = hi_nibble(self.ir as u8);

				debug!("moving from {:02X} to {:02X}", rloc_i0, rloc_o0)
				self.reg[rloc_o0 as uint] = self.reg[rloc_i0 as uint];
				Continue
			},

			0x4   => { // ADDB
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);   // first input: lower nibble of first byte
				let rloc_i1 = hi_nibble(self.ir as u8);          // second input: upper nibble of second byte
				let rloc_o0 = lo_nibble(self.ir as u8);          // output: lower nibble of second byte

				self.reg[rloc_o0 as uint] =
					((self.reg[rloc_i0 as uint] as i8) + (self.reg[rloc_i1 as uint]as i8)) as u8;
				Continue
			},

			0x6   => { // OR
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let rloc_i1 = hi_nibble(self.ir as u8);
				let rloc_o0 = lo_nibble(self.ir as u8);

				self.reg[rloc_o0 as uint] = self.reg[rloc_i0 as uint] | self.reg[rloc_i1 as uint];
				Continue
			},

			0x7   => { // AND
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let rloc_i1 = hi_nibble(self.ir as u8);
				let rloc_o0 = lo_nibble(self.ir as u8);

				self.reg[rloc_o0 as uint] = self.reg[rloc_i0 as uint] & self.reg[rloc_i1 as uint];
				Continue
			},

			0x8   => { // XOR
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let rloc_i1 = hi_nibble(self.ir as u8);
				let rloc_o0 = lo_nibble(self.ir as u8);

				self.reg[rloc_o0 as uint] = self.reg[rloc_i0 as uint] ^ self.reg[rloc_i1 as uint];
				Continue
			},

			0x9   => { // ROT
				// LHS shifts <width> bits off the (left) end of the bitstring
				// RHS shifts the bitstring to the right until only the bits which fell off remain.
				//   LHS is the remaining MSB bits; RHS is the remaining LSB bits
				//   ∴ LHS <OR> RHS provides a rotated bitstring
				//
				let rloc_i0   = lo_nibble((self.ir >> 8) as u8) as uint;          // register is first nibble ...
				let swidth    = (hi_nibble(self.ir as u8) & 0b0000_0111) as uint; // last three bytes of second nibble ...
				self.reg[rloc_i0 as uint] = (self.reg[rloc_i0] << swidth) | (self.reg[rloc_i0] >> (BYTE_WIDTH - swidth));

				Continue
			},

			0xA   => { // JMPEQ
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let next_ip = self.ir as u8;

				if self.reg[rloc_i0 as uint] == self.reg[0] { self.ip = next_ip }
				Continue
			},

			0xB   => { Halt },
			_     => { debug!("halt, cpu on fire: {}", op); Halt },
		}
	}

	/// Load the instruction at `IP` and advance the pointer by two bytes.
	/// The instruction is packed into a single `u16` and stored in the instruction register.
	fn fetch(&mut self) {
		// fetch two bytes from PC
		let byte_1 = self.mem[(self.ip+0) as uint];
		let byte_2 = self.mem[(self.ip+1) as uint];

		self.ir  = (byte_1 as u16 << 8) | (byte_2 as u16); // load byets into IR
		self.ip += 2;                                      // increment instruction pointer

		debug!("IR set to 0x{:04X} ({:02X},{:02X})", self.ir, byte_1, byte_2)
	}
}

/// Take the lower nibble of a byte
fn lo_nibble(byte: u8) -> u8 {
	(byte & U4_LOW)
}

/// Take the upper nibble of a byte and shift it
/// towards the least significant bits.
fn hi_nibble(byte: u8) -> u8 {
	(byte & U4_HIGH) >> 4
}

#[test]
fn test_hammer_time() {
	// cpu should stop after 1 tick of this program
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0xB000]);

	assert_eq!(cpu.tick(), Halt)
}

#[test]
fn test_registers() {
	// cpu should set, move registers accordingly
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x1110, 0x3130, 0xB000]);
	boot(&mut cpu);

	assert_eq!(cpu.get_reg()[0x3], 16)
}

#[test]
fn test_memory() {
	// memory sets and memory stores should read back successfully
	// uninitialized registers should not match initialized registers
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x1120, 0x1330, 0x2140, 0x0240, 0xB000]);
	boot(&mut cpu);

	assert!(cpu.get_reg()[0x1] == cpu.get_reg()[0x2]);
	assert!(cpu.get_reg()[0x1] != cpu.get_reg()[0x3]);
}

#[test]
fn test_bin() {
	// tests the various binary operations against their rustc counterparts.
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x1121, 0x1222, 0x6123, 0x7124, 0x8125, 0xB000]);
	boot(&mut cpu);

	assert!(cpu.get_reg()[0x3] == (0x21 | 0x22));
	assert!(cpu.get_reg()[0x4] == (0x21 & 0x22));
	assert!(cpu.get_reg()[0x5] == (0x21 ^ 0x22));
}

#[test]
fn test_math() {
	// tests basic 2s comp. addition
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x1120, 0x120A, 0x4123, 0xB000]);
	boot(&mut cpu);

	assert_eq!(cpu.get_reg()[0x3], cpu.get_reg()[0x1] + cpu.get_reg()[0x2]);
	assert_eq!(cpu.get_reg()[0x3], 0x20 + 0x0A);
}

#[test]
fn test_branch() {
	// checks that the program branches; skipping a halt and setting a status register
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x1021, 0x1121, 0xA108, 0xB000, 0x122A, 0xB000]);
	boot(&mut cpu);

	assert_eq!(cpu.get_reg()[0x2], 0x2A)
}

#[test]
fn test_shift() {
	// checks that the program rotates a single nibble 4 places
	// this should move a single hex digit from the lhs to the rhs.
	//
	// NOTE: shifting 12 bits (12 - 8) and 4 bits should be equivalent.
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x10B0, 0x11B0, 0x90C0, 0x9140, 0xB000]);
	boot(&mut cpu);

	assert_eq!(cpu.get_reg()[0x0], 0x0B);
	assert_eq!(cpu.get_reg()[0x1], 0x0B);
}

#[cfg(test)]
fn boot(cpu: &mut P150Cpu) {
	loop {
		if cpu.tick() == Halt { break; }
	}
}

<commit_msg>Reset machine pointers when loading user memory<commit_after>use std::collections::HashMap;
use serialize::json;
use serialize::json::{ToJson};

// masks for pulling nibbles out of bytes
const BYTE_WIDTH: uint = 8;
const U4_LOW:     u8   = 0b00001111;
const U4_HIGH:    u8   = 0b11110000;

/// If the CPU enters the `Halt` state then any additional ticks will result in
/// the program exhibiting undefined behavior.
///
#[deriving(PartialEq,Show)]
pub enum CpuState {
	Continue,
	Halt,
}

/// The P150 virtual machine
/// Represents all information the CPU needs to continue executing a
/// program stored in its main memory.
pub struct P150Cpu {
	ip:  u8,
	ir: u16,

	reg: [u8, ..16],
	mem: [u8, ..256],
}

impl P150Cpu {
	/// Initializes the P150 CPU
	///
	/// NOTE: This implementation 0s memory; but this is not guaranteed
	/// by the machine specification.
	///
	/// Programs MUST start at memory address 0x00
	pub fn new() -> P150Cpu {
		P150Cpu {
			ip:  0x00,
			ir:  0x0000,

			reg: [0u8, ..16],
			mem: [0u8, ..256],
		}
	}

	pub fn js_dump(&self) -> json::Json {
		let mut dict = HashMap::new();

		dict.insert(format!("ip"), self.ip.to_json());
		dict.insert(format!("ir"), self.ir.to_json());
		dict.insert(format!("reg"), self.reg.to_json());
		dict.insert(format!("mem"), self.mem.to_json());

		return dict.to_json();
	}

	/// Prints the current machine state to the console window
	/// This includes the IP, IR, and all registers.
	/// (Registers will be formatted as 2s complement numbers.)
	pub fn dump(&self) {
		println!("IP: 0x{:02X}, IR: 0x{:04X}", self.ip, self.ir);

		println!("---\nRegisters\n---")
		for (addr, cell) in self.reg.iter().enumerate() {
			println!("{:01X}: {}", addr, cell)
		}

		println!("---")

		debug!("Memory")
		for (addr, cell) in self.mem.iter().enumerate() {
			debug!("{:02X}: {}", addr, cell)
		}
	}

	#[cfg(test)]
	fn get_reg(&self) -> &[u8] { self.reg }

	/// Read an array of instructions into main memory
	/// This reads two bytes at a time from the `memory` array
	/// and loads them into the P150s RAM bank, starting from address 0.
	pub fn init_mem(&mut self, memory: &[u16]) {
		assert!(memory.len() <= (256 / 2)); // program cannot be larger than memory
		let mut next_cell = 0x00;

		// zero memory
		self.ip  = 0;
		self.ir  = 0;
		self.mem = [0, ..256];
		for op in memory.iter() {
			let byte_1 = (*op >> 8) as u8;
			let byte_2 = *op as u8;

			self.mem[next_cell]     = byte_1;
			self.mem[next_cell + 1] = byte_2;
			next_cell += 2
		}
	}

	/// Runs the entire fetch, execute, decode cycle against the current IP
	pub fn tick(&mut self) -> CpuState {
		self.fetch();

		// decode & execute
		let op = (self.ir >> 12) as u8;
		match op {
			0x0   => { // MLOAD
				let rloc_o0 = lo_nibble((self.ir >> 8) as u8);
				let mloc_i0 = self.ir as u8;

				self.reg[rloc_o0 as uint] = self.mem[mloc_i0 as uint];
				Continue
			},

			0x1   => { // RSET
				let rloc = lo_nibble((self.ir >> 8) as u8);  // lower nibble of first byte
				let rval = self.ir as u8;                    // value is entire second byte

				self.reg[rloc as uint] = rval;               // store value in register
				Continue
			},

			0x2   => { // MSTOR
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let mloc_o0 = self.ir as u8;

				self.mem[mloc_o0 as uint] = self.reg[rloc_i0 as uint];
				Continue
			},

			0x3   => { // RMOV
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let rloc_o0 = hi_nibble(self.ir as u8);

				debug!("moving from {:02X} to {:02X}", rloc_i0, rloc_o0)
				self.reg[rloc_o0 as uint] = self.reg[rloc_i0 as uint];
				Continue
			},

			0x4   => { // ADDB
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);   // first input: lower nibble of first byte
				let rloc_i1 = hi_nibble(self.ir as u8);          // second input: upper nibble of second byte
				let rloc_o0 = lo_nibble(self.ir as u8);          // output: lower nibble of second byte

				self.reg[rloc_o0 as uint] =
					((self.reg[rloc_i0 as uint] as i8) + (self.reg[rloc_i1 as uint]as i8)) as u8;
				Continue
			},

			0x6   => { // OR
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let rloc_i1 = hi_nibble(self.ir as u8);
				let rloc_o0 = lo_nibble(self.ir as u8);

				self.reg[rloc_o0 as uint] = self.reg[rloc_i0 as uint] | self.reg[rloc_i1 as uint];
				Continue
			},

			0x7   => { // AND
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let rloc_i1 = hi_nibble(self.ir as u8);
				let rloc_o0 = lo_nibble(self.ir as u8);

				self.reg[rloc_o0 as uint] = self.reg[rloc_i0 as uint] & self.reg[rloc_i1 as uint];
				Continue
			},

			0x8   => { // XOR
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let rloc_i1 = hi_nibble(self.ir as u8);
				let rloc_o0 = lo_nibble(self.ir as u8);

				self.reg[rloc_o0 as uint] = self.reg[rloc_i0 as uint] ^ self.reg[rloc_i1 as uint];
				Continue
			},

			0x9   => { // ROT
				// LHS shifts <width> bits off the (left) end of the bitstring
				// RHS shifts the bitstring to the right until only the bits which fell off remain.
				//   LHS is the remaining MSB bits; RHS is the remaining LSB bits
				//   ∴ LHS <OR> RHS provides a rotated bitstring
				//
				let rloc_i0   = lo_nibble((self.ir >> 8) as u8) as uint;          // register is first nibble ...
				let swidth    = (hi_nibble(self.ir as u8) & 0b0000_0111) as uint; // last three bytes of second nibble ...
				self.reg[rloc_i0 as uint] = (self.reg[rloc_i0] << swidth) | (self.reg[rloc_i0] >> (BYTE_WIDTH - swidth));

				Continue
			},

			0xA   => { // JMPEQ
				let rloc_i0 = lo_nibble((self.ir >> 8) as u8);
				let next_ip = self.ir as u8;

				if self.reg[rloc_i0 as uint] == self.reg[0] { self.ip = next_ip }
				Continue
			},

			0xB   => { Halt },
			_     => { debug!("halt, cpu on fire: {}", op); Halt },
		}
	}

	/// Load the instruction at `IP` and advance the pointer by two bytes.
	/// The instruction is packed into a single `u16` and stored in the instruction register.
	fn fetch(&mut self) {
		// fetch two bytes from PC
		let byte_1 = self.mem[(self.ip+0) as uint];
		let byte_2 = self.mem[(self.ip+1) as uint];

		self.ir  = (byte_1 as u16 << 8) | (byte_2 as u16); // load byets into IR
		self.ip += 2;                                      // increment instruction pointer

		debug!("IR set to 0x{:04X} ({:02X},{:02X})", self.ir, byte_1, byte_2)
	}
}

/// Take the lower nibble of a byte
fn lo_nibble(byte: u8) -> u8 {
	(byte & U4_LOW)
}

/// Take the upper nibble of a byte and shift it
/// towards the least significant bits.
fn hi_nibble(byte: u8) -> u8 {
	(byte & U4_HIGH) >> 4
}

#[test]
fn test_hammer_time() {
	// cpu should stop after 1 tick of this program
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0xB000]);

	assert_eq!(cpu.tick(), Halt)
}

#[test]
fn test_registers() {
	// cpu should set, move registers accordingly
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x1110, 0x3130, 0xB000]);
	boot(&mut cpu);

	assert_eq!(cpu.get_reg()[0x3], 16)
}

#[test]
fn test_memory() {
	// memory sets and memory stores should read back successfully
	// uninitialized registers should not match initialized registers
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x1120, 0x1330, 0x2140, 0x0240, 0xB000]);
	boot(&mut cpu);

	assert!(cpu.get_reg()[0x1] == cpu.get_reg()[0x2]);
	assert!(cpu.get_reg()[0x1] != cpu.get_reg()[0x3]);
}

#[test]
fn test_bin() {
	// tests the various binary operations against their rustc counterparts.
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x1121, 0x1222, 0x6123, 0x7124, 0x8125, 0xB000]);
	boot(&mut cpu);

	assert!(cpu.get_reg()[0x3] == (0x21 | 0x22));
	assert!(cpu.get_reg()[0x4] == (0x21 & 0x22));
	assert!(cpu.get_reg()[0x5] == (0x21 ^ 0x22));
}

#[test]
fn test_math() {
	// tests basic 2s comp. addition
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x1120, 0x120A, 0x4123, 0xB000]);
	boot(&mut cpu);

	assert_eq!(cpu.get_reg()[0x3], cpu.get_reg()[0x1] + cpu.get_reg()[0x2]);
	assert_eq!(cpu.get_reg()[0x3], 0x20 + 0x0A);
}

#[test]
fn test_branch() {
	// checks that the program branches; skipping a halt and setting a status register
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x1021, 0x1121, 0xA108, 0xB000, 0x122A, 0xB000]);
	boot(&mut cpu);

	assert_eq!(cpu.get_reg()[0x2], 0x2A)
}

#[test]
fn test_shift() {
	// checks that the program rotates a single nibble 4 places
	// this should move a single hex digit from the lhs to the rhs.
	//
	// NOTE: shifting 12 bits (12 - 8) and 4 bits should be equivalent.
	let mut cpu = P150Cpu::new();
	cpu.init_mem([0x10B0, 0x11B0, 0x90C0, 0x9140, 0xB000]);
	boot(&mut cpu);

	assert_eq!(cpu.get_reg()[0x0], 0x0B);
	assert_eq!(cpu.get_reg()[0x1], 0x0B);
}

#[cfg(test)]
fn boot(cpu: &mut P150Cpu) {
	loop {
		if cpu.tick() == Halt { break; }
	}
}

<|endoftext|>


<commit_before>use std::collections::HashMap;

use model::Model;

pub struct Classifier {
    models: HashMap<String, Model>
}

impl Classifier {
    pub fn new() -> Classifier {
        Classifier { models: HashMap::new() }
    }

    pub fn train(&mut self, text: &str, name: &str) {
        let model = Model::build_from_text(text);
        self.models.insert(name.to_string(), model);
    }

    pub fn classify(&self, text: &str) -> &String {
        let model = Model::build_from_text(text);
        let (name, matched_model) = find_min_by_key(
            self.models.iter(),
            |&(name, other_model)| model.compare(other_model)
        ).unwrap();
        name
    }
}

fn find_min_by_key<T, I, F>(iterable: I, key: F) -> Option<T>
    where I : IntoIterator<Item=T>,
          F : Fn(&T) -> usize {

    let mut min_value = None;
    let mut min_item = None;
    for item in iterable {
        let value = key(&item);
        if min_value == None || value < min_value.unwrap() {
            min_value = Some(value);
            min_item = Some(item);
        }
    }
    min_item
}
<commit_msg>Use the standard min_by_key method<commit_after>use std::collections::HashMap;

use model::Model;

pub struct Classifier {
    models: HashMap<String, Model>
}

impl Classifier {
    pub fn new() -> Classifier {
        Classifier { models: HashMap::new() }
    }

    pub fn train(&mut self, text: &str, name: &str) {
        let model = Model::build_from_text(text);
        self.models.insert(name.to_string(), model);
    }

    pub fn classify(&self, text: &str) -> &String {
        let model = Model::build_from_text(text);
        let (name, matched_model) = self.models
            .iter()
            .min_by_key(|&(name, other_model)| model.compare(other_model))
            .unwrap();
        name
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use crate::{firmware::Metadata, runtime_settings::RuntimeSettings, update_package::UpdatePackage};

use failure::bail;
use log::debug;
use reqwest::{
    header::{HeaderMap, HeaderName, CONTENT_TYPE, RANGE, USER_AGENT},
    Client, StatusCode,
};
use serde_derive::Serialize;
use std::{path::Path, time::Duration};

#[cfg(test)]
pub(crate) mod tests;

pub(crate) struct Api<'a> {
    server: &'a str,
}

#[derive(Debug)]
pub(crate) enum ProbeResponse {
    NoUpdate,
    Update(UpdatePackage),
    ExtraPoll(i64),
}

impl<'a> Api<'a> {
    pub(crate) fn new(server: &'a str) -> Api<'a> {
        Api { server }
    }

    fn client(&self) -> Result<Client, failure::Error> {
        let mut headers = HeaderMap::new();

        headers.insert(USER_AGENT, "updatehub/next".parse()?);
        headers.insert(CONTENT_TYPE, "application/json".parse()?);
        headers.insert(
            HeaderName::from_static("api-content-type"),
            "application/vnd.updatehub-v1+json".parse()?,
        );

        Ok(Client::builder()
            .timeout(Duration::from_secs(10))
            .default_headers(headers)
            .build()?)
    }

    pub fn probe(
        &self,
        runtime_settings: &RuntimeSettings,
        firmware: &Metadata,
    ) -> Result<ProbeResponse, failure::Error> {
        let mut response = self
            .client()?
            .post(&format!("{}/upgrades", &self.server))
            .header(
                HeaderName::from_static("api-retries"),
                runtime_settings.retries(),
            )
            .json(firmware)
            .send()?;

        match response.status() {
            StatusCode::NOT_FOUND => Ok(ProbeResponse::NoUpdate),
            StatusCode::OK => {
                if let Some(extra_poll) = response
                    .headers()
                    .get("add-extra-poll")
                    .and_then(|extra_poll| extra_poll.to_str().ok())
                    .and_then(|extra_poll| extra_poll.parse().ok())
                {
                    return Ok(ProbeResponse::ExtraPoll(extra_poll));
                }

                Ok(ProbeResponse::Update(UpdatePackage::parse(
                    &response.text()?,
                )?))
            }
            _ => bail!("Invalid response. Status: {}", response.status()),
        }
    }

    pub fn download_object(
        &self,
        product_uid: &str,
        package_uid: &str,
        download_dir: &Path,
        object: &str,
    ) -> Result<(), failure::Error> {
        use std::fs::{create_dir_all, OpenOptions};

        // FIXME: Discuss the need of packages inside the route
        let mut client = self.client()?.get(&format!(
            "{}/products/{}/packages/{}/objects/{}",
            &self.server, product_uid, package_uid, object
        ));

        let path = download_dir;
        if !&path.exists() {
            debug!("Creating directory to store the downloads.");
            create_dir_all(&path)?;
        }

        let file = path.join(object);
        if file.exists() {
            client = client.header(RANGE, format!("bytes={}-", file.metadata()?.len() - 1));
        }

        let mut file = OpenOptions::new().create(true).append(true).open(&file)?;
        let mut response = client.send()?;
        if response.status().is_success() {
            response.copy_to(&mut file)?;
            return Ok(());
        }

        bail!("Couldn't download the object {}", object)
    }

    pub fn report(
        &self,
        state: &str,
        firmware: &'a Metadata,
        package_uid: &str,
        previous_state: Option<&str>,
        error_message: Option<String>,
    ) -> Result<(), failure::Error> {
        #[derive(Serialize)]
        #[serde(rename_all = "kebab-case")]
        struct Payload<'a> {
            status: &'a str,
            #[serde(flatten)]
            firmware: &'a Metadata,
            package_uid: &'a str,
            #[serde(skip_serializing_if = "Option::is_none")]
            previous_state: Option<&'a str>,
            #[serde(skip_serializing_if = "Option::is_none")]
            error_message: Option<String>,
        }

        let payload = Payload {
            status: state,
            firmware,
            package_uid,
            previous_state,
            error_message,
        };

        self.client()?
            .post(&format!("{}/report", &self.server))
            .json(&payload)
            .send()?;
        Ok(())
    }
}
<commit_msg>client: probe: Refactor code to be more obvious<commit_after>// Copyright (C) 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use crate::{firmware::Metadata, runtime_settings::RuntimeSettings, update_package::UpdatePackage};

use failure::bail;
use log::debug;
use reqwest::{
    header::{HeaderMap, HeaderName, CONTENT_TYPE, RANGE, USER_AGENT},
    Client, StatusCode,
};
use serde_derive::Serialize;
use std::{path::Path, time::Duration};

#[cfg(test)]
pub(crate) mod tests;

pub(crate) struct Api<'a> {
    server: &'a str,
}

#[derive(Debug)]
pub(crate) enum ProbeResponse {
    NoUpdate,
    Update(UpdatePackage),
    ExtraPoll(i64),
}

impl<'a> Api<'a> {
    pub(crate) fn new(server: &'a str) -> Api<'a> {
        Api { server }
    }

    fn client(&self) -> Result<Client, failure::Error> {
        let mut headers = HeaderMap::new();

        headers.insert(USER_AGENT, "updatehub/next".parse()?);
        headers.insert(CONTENT_TYPE, "application/json".parse()?);
        headers.insert(
            HeaderName::from_static("api-content-type"),
            "application/vnd.updatehub-v1+json".parse()?,
        );

        Ok(Client::builder()
            .timeout(Duration::from_secs(10))
            .default_headers(headers)
            .build()?)
    }

    pub fn probe(
        &self,
        runtime_settings: &RuntimeSettings,
        firmware: &Metadata,
    ) -> Result<ProbeResponse, failure::Error> {
        let mut response = self
            .client()?
            .post(&format!("{}/upgrades", &self.server))
            .header(
                HeaderName::from_static("api-retries"),
                runtime_settings.retries(),
            )
            .json(firmware)
            .send()?;

        match response.status() {
            StatusCode::NOT_FOUND => Ok(ProbeResponse::NoUpdate),
            StatusCode::OK => {
                match response
                    .headers()
                    .get("add-extra-poll")
                    .and_then(|extra_poll| extra_poll.to_str().ok())
                    .and_then(|extra_poll| extra_poll.parse().ok())
                {
                    Some(extra_poll) => Ok(ProbeResponse::ExtraPoll(extra_poll)),
                    None => Ok(ProbeResponse::Update(UpdatePackage::parse(
                        &response.text()?,
                    )?)),
                }
            }
            _ => bail!("Invalid response. Status: {}", response.status()),
        }
    }

    pub fn download_object(
        &self,
        product_uid: &str,
        package_uid: &str,
        download_dir: &Path,
        object: &str,
    ) -> Result<(), failure::Error> {
        use std::fs::{create_dir_all, OpenOptions};

        // FIXME: Discuss the need of packages inside the route
        let mut client = self.client()?.get(&format!(
            "{}/products/{}/packages/{}/objects/{}",
            &self.server, product_uid, package_uid, object
        ));

        let path = download_dir;
        if !&path.exists() {
            debug!("Creating directory to store the downloads.");
            create_dir_all(&path)?;
        }

        let file = path.join(object);
        if file.exists() {
            client = client.header(RANGE, format!("bytes={}-", file.metadata()?.len() - 1));
        }

        let mut file = OpenOptions::new().create(true).append(true).open(&file)?;
        let mut response = client.send()?;
        if response.status().is_success() {
            response.copy_to(&mut file)?;
            return Ok(());
        }

        bail!("Couldn't download the object {}", object)
    }

    pub fn report(
        &self,
        state: &str,
        firmware: &'a Metadata,
        package_uid: &str,
        previous_state: Option<&str>,
        error_message: Option<String>,
    ) -> Result<(), failure::Error> {
        #[derive(Serialize)]
        #[serde(rename_all = "kebab-case")]
        struct Payload<'a> {
            status: &'a str,
            #[serde(flatten)]
            firmware: &'a Metadata,
            package_uid: &'a str,
            #[serde(skip_serializing_if = "Option::is_none")]
            previous_state: Option<&'a str>,
            #[serde(skip_serializing_if = "Option::is_none")]
            error_message: Option<String>,
        }

        let payload = Payload {
            status: state,
            firmware,
            package_uid,
            previous_state,
            error_message,
        };

        self.client()?
            .post(&format!("{}/report", &self.server))
            .json(&payload)
            .send()?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>#[cfg(feature = "parallel")]
extern crate rayon;
#[cfg(feature = "parallel")]
use rayon::prelude::*;

use std::env;
use std::ffi::OsString;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::process::Stdio;

fn x86_triple(os: &str) -> &'static str {
  match os {
    "darwin" => "-fmacho32",
    "windows" => "-fwin32",
    _ => "-felf32",
  }
}

fn x86_64_triple(os: &str) -> &'static str {
  match os {
    "darwin" => "-fmacho64",
    "windows" => "-fwin64",
    _ => "-felf64",
  }
}

fn parse_triple(trip: &str) -> &'static str {
  let parts = trip.split('-').collect::<Vec<_>>();
  // ARCH-VENDOR-OS-ENVIRONMENT
  // or ARCH-VENDOR-OS
  // we don't care about environ so doesn't matter if triple doesn't have it
  if parts.len() < 3 {
    return "";
  }

  match parts[0] {
    "x86_64" => x86_64_triple(parts[2]),
    "x86" | "i386" | "i586" | "i686" => x86_triple(parts[2]),
    _ => "",
  }
}

/// # Example
///
/// ```no_run
/// nasm_rs::compile_library("libfoo.a", &["foo.s", "bar.s"]);
/// ```
pub fn compile_library(output: &str, files: &[&str]) {
  compile_library_args(output, files, &[]);
}

/// # Example
///
/// ```no_run
/// nasm_rs::compile_library_args("libfoo.a", &["foo.s", "bar.s"], &["-Fdwarf"]);
/// ```
pub fn compile_library_args<P: AsRef<Path>>(
  output: &str, files: &[P], args: &[&str],
) {
  let mut b = Build::new();
  for file in files {
    b.file(file);
  }
  for arg in args {
    b.flag(arg);
  }
  b.compile(output);
}

pub struct Build {
  files: Vec<PathBuf>,
  flags: Vec<String>,
  target: Option<String>,
  out_dir: Option<PathBuf>,
  archiver: Option<PathBuf>,
  nasm: Option<PathBuf>,
  debug: bool,
}

impl Build {
  pub fn new() -> Self {
    Self {
      files: Vec::new(),
      flags: Vec::new(),
      archiver: None,
      out_dir: None,
      nasm: None,
      target: None,
      debug: env::var("DEBUG").ok().map_or(false, |d| d != "false"),
    }
  }

  /// Add a file which will be compiled
  ///
  /// e.g. `"foo.s"`
  pub fn file<P: AsRef<Path>>(&mut self, p: P) -> &mut Self {
    self.files.push(p.as_ref().to_owned());
    self
  }

  /// Set multiple files
  pub fn files<P: AsRef<Path>, I: IntoIterator<Item = P>>(
    &mut self, files: I,
  ) -> &mut Self {
    for file in files {
      self.file(file);
    }
    self
  }

  /// Add a directory to the `-I` include path
  pub fn include<P: AsRef<Path>>(&mut self, dir: P) -> &mut Self {
    let mut flag = format!("-I{}", dir.as_ref().display());
    // nasm requires trailing slash, but `Path` may omit it.
    if !flag.ends_with("/") {
      flag += "/";
    }
    self.flags.push(flag);
    self
  }

  /// Pre-define a macro with an optional value
  pub fn define<'a, V: Into<Option<&'a str>>>(
    &mut self, var: &str, val: V,
  ) -> &mut Self {
    let val = val.into();
    let flag = if let Some(val) = val {
      format!("-D{}={}", var, val)
    } else {
      format!("-D{}", var)
    };
    self.flags.push(flag);
    self
  }

  /// Configures whether the assembler will generate debug information.
  ///
  /// This option is automatically scraped from the `DEBUG` environment
  /// variable by build scripts (only enabled when the profile is "debug"), so
  /// it's not required to call this function.
  pub fn debug(&mut self, enable: bool) -> &mut Self {
    self.debug = enable;
    self
  }

  /// Add an arbitrary flag to the invocation of the assembler
  ///
  /// e.g. `"-Fdwarf"`
  pub fn flag(&mut self, flag: &str) -> &mut Self {
    self.flags.push(flag.to_owned());
    self
  }

  /// Configures the target this configuration will be compiling for.
  ///
  /// This option is automatically scraped from the `TARGET` environment
  /// variable by build scripts, so it's not required to call this function.
  pub fn target(&mut self, target: &str) -> &mut Self {
    self.target = Some(target.to_owned());
    self
  }

  /// Configures the output directory where all object files and static libraries will be located.
  ///
  /// This option is automatically scraped from the OUT_DIR environment variable by build scripts,
  /// so it's not required to call this function.
  pub fn out_dir<P: AsRef<Path>>(&mut self, out_dir: P) -> &mut Self {
    self.out_dir = Some(out_dir.as_ref().to_owned());
    self
  }

  /// Configures the tool used to assemble archives.
  ///
  /// This option is automatically determined from the target platform or a
  /// number of environment variables, so it's not required to call this
  /// function.
  pub fn archiver<P: AsRef<Path>>(&mut self, archiver: P) -> &mut Self {
    self.archiver = Some(archiver.as_ref().to_owned());
    self
  }

  /// Configures path to `nasm` command
  pub fn nasm<P: AsRef<Path>>(&mut self, nasm: P) -> &mut Self {
    self.nasm = Some(nasm.as_ref().to_owned());
    self
  }

  /// Run the compiler, generating the file output
  ///
  /// The name output should be the base name of the library,
  /// without file extension, and without "lib" prefix.
  ///
  /// The output file will have target-specific name,
  /// such as `lib*.a` (non-MSVC) or `*.lib` (MSVC).
  pub fn compile(&mut self, lib_name: &str) {
    // Trim name for backwards comatibility
    let lib_name = if lib_name.starts_with("lib") && lib_name.ends_with(".a") {
      &lib_name[3..lib_name.len() - 2]
    } else {
      lib_name.trim_end_matches(".lib")
    };

    let target = self.get_target();
    let output = if target.ends_with("-msvc") {
      format!("{}.lib", lib_name)
    } else {
      format!("lib{}.a", lib_name)
    };

    let dst = &self.get_out_dir();
    let objects = self.compile_objects();
    self.archive(&dst, &output, &objects[..]);

    println!("cargo:rustc-link-search={}", dst.display());
  }

  /// Run the compiler, generating .o files
  ///
  /// The files can be linked in a separate step, e.g. passed to `cc`
  pub fn compile_objects(&mut self) -> Vec<PathBuf> {
    let target = self.get_target();

    let nasm = self.find_nasm();
    let args = self.get_args(&target);

    let src = &PathBuf::from(
      env::var_os("CARGO_MANIFEST_DIR")
        .expect("CARGO_MANIFEST_DIR must be set"),
    );
    let dst = &self.get_out_dir();

    self
      .make_iter()
      .map(|file| self.compile_file(&nasm, file.as_ref(), &args, src, dst))
      .collect::<Vec<_>>()
  }

  fn get_args(&self, target: &str) -> Vec<&str> {
    let mut args = vec![parse_triple(&target)];

    if self.debug {
      args.push("-g");
    }

    for arg in &self.flags {
      args.push(arg);
    }

    args
  }

  #[cfg(feature = "parallel")]
  fn make_iter(&self) -> rayon::slice::Iter<PathBuf> {
    self.files.par_iter()
  }

  #[cfg(not(feature = "parallel"))]
  fn make_iter(&self) -> std::slice::Iter<PathBuf> {
    self.files.iter()
  }

  fn compile_file(
    &self, nasm: &Path, file: &Path, new_args: &[&str], src: &Path, dst: &Path,
  ) -> PathBuf {
    let obj = dst.join(file).with_extension("o");
    let mut cmd = Command::new(nasm);
    cmd.args(&new_args[..]);
    std::fs::create_dir_all(&obj.parent().unwrap()).unwrap();

    run(cmd.arg(src.join(file)).arg("-o").arg(&obj));
    obj
  }

  fn archive(&self, out_dir: &Path, lib: &str, objs: &[PathBuf]) {
    let ar = if cfg!(target_env = "msvc") {
      self.archiver.clone().unwrap_or_else(|| "lib".into())
    } else {
      self
        .archiver
        .clone()
        .or_else(|| env::var_os("AR").map(|a| a.into()))
        .unwrap_or_else(|| "ar".into())
    };
    if cfg!(target_env = "msvc") {
      let mut out_param = OsString::new();
      out_param.push("/OUT:");
      out_param.push(out_dir.join(lib).as_os_str());
      run(Command::new(ar).arg(out_param).args(objs));
    } else {
      run(Command::new(ar).arg("crus").arg(out_dir.join(lib)).args(objs));
    }
  }

  fn get_out_dir(&self) -> PathBuf {
    self.out_dir.clone().unwrap_or_else(|| {
      PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR must be set"))
    })
  }

  fn get_target(&self) -> String {
    self
      .target
      .clone()
      .unwrap_or_else(|| env::var("TARGET").expect("TARGET must be set"))
  }

  fn find_nasm(&mut self) -> PathBuf {
    match self.nasm.clone() {
      Some(path) => path,
      None => {
        let nasm_path = PathBuf::from("nasm");
        match is_nasm_new_enough(&nasm_path) {
          Ok(_) => nasm_path,
          Err(version) => {
            panic!("This version of NASM is too old: {}", version);
          }
        }
      }
    }
  }
}

fn get_output(cmd: &mut Command) -> Result<String, String> {
  let out = cmd.output().map_err(|e| e.to_string())?;
  if out.status.success() {
    Ok(String::from_utf8_lossy(&out.stdout).to_string())
  } else {
    Err(String::from_utf8_lossy(&out.stderr))?
  }
}

/// Returns version string if nasm is too old,
/// or error message string if it's unusable.
fn is_nasm_new_enough(nasm_path: &Path) -> Result<(), String> {
  match get_output(Command::new(nasm_path).arg("-v")) {
    Ok(version) => {
      if version.contains("NASM version 0.") {
        Err(version)?
      } else {
        Ok(())
      }
    }
    Err(err) => Err(err)?,
  }
}

fn run(cmd: &mut Command) {
  println!("running: {:?}", cmd);

  let status =
    match cmd.stdout(Stdio::inherit()).stderr(Stdio::inherit()).status() {
      Ok(status) => status,

      Err(e) => panic!("failed to spawn process: {}", e),
    };

  if !status.success() {
    panic!("nonzero exit status: {}", status);
  }
}

#[test]
fn test_build() {
  let mut build = Build::new();
  build.file("test");
  build.archiver("ar");
  build.include("./");
  build.include("dir");
  build.define("foo", Some("1"));
  build.define("bar", None);
  build.flag("-test");
  build.target("i686-unknown-linux-musl");
  build.out_dir("/tmp");

  assert_eq!(
    build.get_args("i686-unknown-linux-musl"),
    &["-felf32", "-I./", "-Idir/", "-Dfoo=1", "-Dbar", "-test"]
  );
}
<commit_msg>Fix clippy issues in nasm-rs<commit_after>#[cfg(feature = "parallel")]
extern crate rayon;
#[cfg(feature = "parallel")]
use rayon::prelude::*;

use std::env;
use std::ffi::OsString;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::process::Stdio;

fn x86_triple(os: &str) -> &'static str {
  match os {
    "darwin" => "-fmacho32",
    "windows" => "-fwin32",
    _ => "-felf32",
  }
}

fn x86_64_triple(os: &str) -> &'static str {
  match os {
    "darwin" => "-fmacho64",
    "windows" => "-fwin64",
    _ => "-felf64",
  }
}

fn parse_triple(trip: &str) -> &'static str {
  let parts = trip.split('-').collect::<Vec<_>>();
  // ARCH-VENDOR-OS-ENVIRONMENT
  // or ARCH-VENDOR-OS
  // we don't care about environ so doesn't matter if triple doesn't have it
  if parts.len() < 3 {
    return "";
  }

  match parts[0] {
    "x86_64" => x86_64_triple(parts[2]),
    "x86" | "i386" | "i586" | "i686" => x86_triple(parts[2]),
    _ => "",
  }
}

/// # Example
///
/// ```no_run
/// nasm_rs::compile_library("libfoo.a", &["foo.s", "bar.s"]);
/// ```
pub fn compile_library(output: &str, files: &[&str]) {
  compile_library_args(output, files, &[]);
}

/// # Example
///
/// ```no_run
/// nasm_rs::compile_library_args("libfoo.a", &["foo.s", "bar.s"], &["-Fdwarf"]);
/// ```
pub fn compile_library_args<P: AsRef<Path>>(
  output: &str, files: &[P], args: &[&str],
) {
  let mut b = Build::new();
  for file in files {
    b.file(file);
  }
  for arg in args {
    b.flag(arg);
  }
  b.compile(output);
}

pub struct Build {
  files: Vec<PathBuf>,
  flags: Vec<String>,
  target: Option<String>,
  out_dir: Option<PathBuf>,
  archiver: Option<PathBuf>,
  nasm: Option<PathBuf>,
  debug: bool,
}

impl Default for Build {
  fn default() -> Self {
    Self {
      files: Vec::new(),
      flags: Vec::new(),
      archiver: None,
      out_dir: None,
      nasm: None,
      target: None,
      debug: env::var("DEBUG").ok().map_or(false, |d| d != "false"),
    }
  }
}

impl Build {
  pub fn new() -> Self {
    Self::default()
  }

  /// Add a file which will be compiled
  ///
  /// e.g. `"foo.s"`
  pub fn file<P: AsRef<Path>>(&mut self, p: P) -> &mut Self {
    self.files.push(p.as_ref().to_owned());
    self
  }

  /// Set multiple files
  pub fn files<P: AsRef<Path>, I: IntoIterator<Item = P>>(
    &mut self, files: I,
  ) -> &mut Self {
    for file in files {
      self.file(file);
    }
    self
  }

  /// Add a directory to the `-I` include path
  pub fn include<P: AsRef<Path>>(&mut self, dir: P) -> &mut Self {
    let mut flag = format!("-I{}", dir.as_ref().display());
    // nasm requires trailing slash, but `Path` may omit it.
    if !flag.ends_with('/') {
      flag += "/";
    }
    self.flags.push(flag);
    self
  }

  /// Pre-define a macro with an optional value
  pub fn define<'a, V: Into<Option<&'a str>>>(
    &mut self, var: &str, val: V,
  ) -> &mut Self {
    let val = val.into();
    let flag = if let Some(val) = val {
      format!("-D{}={}", var, val)
    } else {
      format!("-D{}", var)
    };
    self.flags.push(flag);
    self
  }

  /// Configures whether the assembler will generate debug information.
  ///
  /// This option is automatically scraped from the `DEBUG` environment
  /// variable by build scripts (only enabled when the profile is "debug"), so
  /// it's not required to call this function.
  pub fn debug(&mut self, enable: bool) -> &mut Self {
    self.debug = enable;
    self
  }

  /// Add an arbitrary flag to the invocation of the assembler
  ///
  /// e.g. `"-Fdwarf"`
  pub fn flag(&mut self, flag: &str) -> &mut Self {
    self.flags.push(flag.to_owned());
    self
  }

  /// Configures the target this configuration will be compiling for.
  ///
  /// This option is automatically scraped from the `TARGET` environment
  /// variable by build scripts, so it's not required to call this function.
  pub fn target(&mut self, target: &str) -> &mut Self {
    self.target = Some(target.to_owned());
    self
  }

  /// Configures the output directory where all object files and static libraries will be located.
  ///
  /// This option is automatically scraped from the OUT_DIR environment variable by build scripts,
  /// so it's not required to call this function.
  pub fn out_dir<P: AsRef<Path>>(&mut self, out_dir: P) -> &mut Self {
    self.out_dir = Some(out_dir.as_ref().to_owned());
    self
  }

  /// Configures the tool used to assemble archives.
  ///
  /// This option is automatically determined from the target platform or a
  /// number of environment variables, so it's not required to call this
  /// function.
  pub fn archiver<P: AsRef<Path>>(&mut self, archiver: P) -> &mut Self {
    self.archiver = Some(archiver.as_ref().to_owned());
    self
  }

  /// Configures path to `nasm` command
  pub fn nasm<P: AsRef<Path>>(&mut self, nasm: P) -> &mut Self {
    self.nasm = Some(nasm.as_ref().to_owned());
    self
  }

  /// Run the compiler, generating the file output
  ///
  /// The name output should be the base name of the library,
  /// without file extension, and without "lib" prefix.
  ///
  /// The output file will have target-specific name,
  /// such as `lib*.a` (non-MSVC) or `*.lib` (MSVC).
  pub fn compile(&mut self, lib_name: &str) {
    // Trim name for backwards comatibility
    let lib_name = if lib_name.starts_with("lib") && lib_name.ends_with(".a") {
      &lib_name[3..lib_name.len() - 2]
    } else {
      lib_name.trim_end_matches(".lib")
    };

    let target = self.get_target();
    let output = if target.ends_with("-msvc") {
      format!("{}.lib", lib_name)
    } else {
      format!("lib{}.a", lib_name)
    };

    let dst = &self.get_out_dir();
    let objects = self.compile_objects();
    self.archive(&dst, &output, &objects[..]);

    println!("cargo:rustc-link-search={}", dst.display());
  }

  /// Run the compiler, generating .o files
  ///
  /// The files can be linked in a separate step, e.g. passed to `cc`
  pub fn compile_objects(&mut self) -> Vec<PathBuf> {
    let target = self.get_target();

    let nasm = self.find_nasm();
    let args = self.get_args(&target);

    let src = &PathBuf::from(
      env::var_os("CARGO_MANIFEST_DIR")
        .expect("CARGO_MANIFEST_DIR must be set"),
    );
    let dst = &self.get_out_dir();

    self
      .make_iter()
      .map(|file| self.compile_file(&nasm, file.as_ref(), &args, src, dst))
      .collect::<Vec<_>>()
  }

  fn get_args(&self, target: &str) -> Vec<&str> {
    let mut args = vec![parse_triple(&target)];

    if self.debug {
      args.push("-g");
    }

    for arg in &self.flags {
      args.push(arg);
    }

    args
  }

  #[cfg(feature = "parallel")]
  fn make_iter(&self) -> rayon::slice::Iter<PathBuf> {
    self.files.par_iter()
  }

  #[cfg(not(feature = "parallel"))]
  fn make_iter(&self) -> std::slice::Iter<PathBuf> {
    self.files.iter()
  }

  fn compile_file(
    &self, nasm: &Path, file: &Path, new_args: &[&str], src: &Path, dst: &Path,
  ) -> PathBuf {
    let obj = dst.join(file).with_extension("o");
    let mut cmd = Command::new(nasm);
    cmd.args(&new_args[..]);
    std::fs::create_dir_all(&obj.parent().unwrap()).unwrap();

    run(cmd.arg(src.join(file)).arg("-o").arg(&obj));
    obj
  }

  fn archive(&self, out_dir: &Path, lib: &str, objs: &[PathBuf]) {
    let ar = if cfg!(target_env = "msvc") {
      self.archiver.clone().unwrap_or_else(|| "lib".into())
    } else {
      self
        .archiver
        .clone()
        .or_else(|| env::var_os("AR").map(|a| a.into()))
        .unwrap_or_else(|| "ar".into())
    };
    if cfg!(target_env = "msvc") {
      let mut out_param = OsString::new();
      out_param.push("/OUT:");
      out_param.push(out_dir.join(lib).as_os_str());
      run(Command::new(ar).arg(out_param).args(objs));
    } else {
      run(Command::new(ar).arg("crus").arg(out_dir.join(lib)).args(objs));
    }
  }

  fn get_out_dir(&self) -> PathBuf {
    self.out_dir.clone().unwrap_or_else(|| {
      PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR must be set"))
    })
  }

  fn get_target(&self) -> String {
    self
      .target
      .clone()
      .unwrap_or_else(|| env::var("TARGET").expect("TARGET must be set"))
  }

  fn find_nasm(&mut self) -> PathBuf {
    match self.nasm.clone() {
      Some(path) => path,
      None => {
        let nasm_path = PathBuf::from("nasm");
        match is_nasm_new_enough(&nasm_path) {
          Ok(_) => nasm_path,
          Err(version) => {
            panic!("This version of NASM is too old: {}", version);
          }
        }
      }
    }
  }
}

fn get_output(cmd: &mut Command) -> Result<String, String> {
  let out = cmd.output().map_err(|e| e.to_string())?;
  if out.status.success() {
    Ok(String::from_utf8_lossy(&out.stdout).to_string())
  } else {
    Err(String::from_utf8_lossy(&out.stderr))?
  }
}

/// Returns version string if nasm is too old,
/// or error message string if it's unusable.
fn is_nasm_new_enough(nasm_path: &Path) -> Result<(), String> {
  match get_output(Command::new(nasm_path).arg("-v")) {
    Ok(version) => {
      if version.contains("NASM version 0.") {
        Err(version)?
      } else {
        Ok(())
      }
    }
    Err(err) => Err(err)?,
  }
}

fn run(cmd: &mut Command) {
  println!("running: {:?}", cmd);

  let status =
    match cmd.stdout(Stdio::inherit()).stderr(Stdio::inherit()).status() {
      Ok(status) => status,

      Err(e) => panic!("failed to spawn process: {}", e),
    };

  if !status.success() {
    panic!("nonzero exit status: {}", status);
  }
}

#[test]
fn test_build() {
  let mut build = Build::new();
  build.file("test");
  build.archiver("ar");
  build.include("./");
  build.include("dir");
  build.define("foo", Some("1"));
  build.define("bar", None);
  build.flag("-test");
  build.target("i686-unknown-linux-musl");
  build.out_dir("/tmp");

  assert_eq!(
    build.get_args("i686-unknown-linux-musl"),
    &["-felf32", "-I./", "-Idir/", "-Dfoo=1", "-Dbar", "-test"]
  );
}
<|endoftext|>


<commit_before>//! Convenience macros.

use std::{
    fmt, mem, panic,
    sync::atomic::{AtomicUsize, Ordering::SeqCst},
};
#[macro_export]
macro_rules! eprintln {
    ($($tt:tt)*) => {{
        if $crate::is_ci() {
            panic!("Forgot to remove debug-print?")
        }
        std::eprintln!($($tt)*)
    }}
}

/// Appends formatted string to a `String`.
#[macro_export]
macro_rules! format_to {
    ($buf:expr) => ();
    ($buf:expr, $lit:literal $($arg:tt)*) => {
        { use ::std::fmt::Write as _; let _ = ::std::write!($buf, $lit $($arg)*); }
    };
}

/// Generates `From` impls for `Enum E { Foo(Foo), Bar(Bar) }` enums
///
/// # Example
///
/// ```rust
/// impl_from!(Struct, Union, Enum for Adt);
/// ```
#[macro_export]
macro_rules! impl_from {
    ($($variant:ident $(($($sub_variant:ident),*))?),* for $enum:ident) => {
        $(
            impl From<$variant> for $enum {
                fn from(it: $variant) -> $enum {
                    $enum::$variant(it)
                }
            }
            $($(
                impl From<$sub_variant> for $enum {
                    fn from(it: $sub_variant) -> $enum {
                        $enum::$variant($variant::$sub_variant(it))
                    }
                }
            )*)?
        )*
    }
}

/// A version of `assert!` macro which allows to handle an assertion failure.
///
/// In release mode, it returns the condition and logs an error.
///
/// ```
/// if assert_never!(impossible) {
///     // Heh, this shouldn't have happened, but lets try to soldier on...
///     return None;
/// }
/// ```
///
/// Rust analyzer is a long-running process, and crashing really isn't an option.
///
/// Shamelessly stolen from: https://www.sqlite.org/assert.html
#[macro_export]
macro_rules! assert_never {
    ($cond:expr) => { $crate::assert_always!($cond, "") };
    ($cond:expr, $($fmt:tt)*) => {{
        let value = $cond;
        if value {
            $crate::on_assert_failure(
                format_args!($($fmt)*)
            );
        }
        value
    }};
}

type AssertHook = fn(&panic::Location<'_>, fmt::Arguments<'_>);
static HOOK: AtomicUsize = AtomicUsize::new(0);

pub fn set_assert_hook(hook: AssertHook) {
    HOOK.store(hook as usize, SeqCst);
}

#[cold]
#[track_caller]
pub fn on_assert_failure(args: fmt::Arguments) {
    let hook: usize = HOOK.load(SeqCst);
    if hook == 0 {
        panic!("\n  assertion failed: {}\n", args);
    }

    let hook: AssertHook = unsafe { mem::transmute::<usize, AssertHook>(hook) };
    hook(panic::Location::caller(), args)
}
<commit_msg>Fix assert_never invoking assert_always<commit_after>//! Convenience macros.

use std::{
    fmt, mem, panic,
    sync::atomic::{AtomicUsize, Ordering::SeqCst},
};
#[macro_export]
macro_rules! eprintln {
    ($($tt:tt)*) => {{
        if $crate::is_ci() {
            panic!("Forgot to remove debug-print?")
        }
        std::eprintln!($($tt)*)
    }}
}

/// Appends formatted string to a `String`.
#[macro_export]
macro_rules! format_to {
    ($buf:expr) => ();
    ($buf:expr, $lit:literal $($arg:tt)*) => {
        { use ::std::fmt::Write as _; let _ = ::std::write!($buf, $lit $($arg)*); }
    };
}

/// Generates `From` impls for `Enum E { Foo(Foo), Bar(Bar) }` enums
///
/// # Example
///
/// ```rust
/// impl_from!(Struct, Union, Enum for Adt);
/// ```
#[macro_export]
macro_rules! impl_from {
    ($($variant:ident $(($($sub_variant:ident),*))?),* for $enum:ident) => {
        $(
            impl From<$variant> for $enum {
                fn from(it: $variant) -> $enum {
                    $enum::$variant(it)
                }
            }
            $($(
                impl From<$sub_variant> for $enum {
                    fn from(it: $sub_variant) -> $enum {
                        $enum::$variant($variant::$sub_variant(it))
                    }
                }
            )*)?
        )*
    }
}

/// A version of `assert!` macro which allows to handle an assertion failure.
///
/// In release mode, it returns the condition and logs an error.
///
/// ```
/// if assert_never!(impossible) {
///     // Heh, this shouldn't have happened, but lets try to soldier on...
///     return None;
/// }
/// ```
///
/// Rust analyzer is a long-running process, and crashing really isn't an option.
///
/// Shamelessly stolen from: https://www.sqlite.org/assert.html
#[macro_export]
macro_rules! assert_never {
    ($cond:expr) => { $crate::assert_never!($cond, "") };
    ($cond:expr, $($fmt:tt)*) => {{
        let value = $cond;
        if value {
            $crate::on_assert_failure(
                format_args!($($fmt)*)
            );
        }
        value
    }};
}

type AssertHook = fn(&panic::Location<'_>, fmt::Arguments<'_>);
static HOOK: AtomicUsize = AtomicUsize::new(0);

pub fn set_assert_hook(hook: AssertHook) {
    HOOK.store(hook as usize, SeqCst);
}

#[cold]
#[track_caller]
pub fn on_assert_failure(args: fmt::Arguments) {
    let hook: usize = HOOK.load(SeqCst);
    if hook == 0 {
        panic!("\n  assertion failed: {}\n", args);
    }

    let hook: AssertHook = unsafe { mem::transmute::<usize, AssertHook>(hook) };
    hook(panic::Location::caller(), args)
}
<|endoftext|>


<commit_before>use std::net::{SocketAddr, IpAddr, Ipv4Addr, Ipv6Addr, UdpSocket};
use std::os::unix::io::AsRawFd;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering, ATOMIC_BOOL_INIT};
use std::io::{Write, Read};
use mio;
use dns_lookup;
use bincode::SizeLimit;
use bincode::rustc_serialize::{encode, decode};
use tuntap;
use utils;
use snap;

pub static INTERRUPTED: AtomicBool = ATOMIC_BOOL_INIT;

#[derive(RustcEncodable, RustcDecodable, PartialEq, Debug)]
enum Message {
    Request,
    Response { id: u8 },
    Data { id: u8, data: Vec<u8> },
}

const TUN: mio::Token = mio::Token(0);
const SOCK: mio::Token = mio::Token(1);

fn resolve(host: &str) -> Result<IpAddr, String> {
    let mut ip_list = try!(dns_lookup::lookup_host(host).map_err(|_| "dns_lookup::lookup_host"));
    let ip = ip_list.next().unwrap().unwrap();
    Ok(ip)
}

fn initiate(socket: &UdpSocket, addr: &SocketAddr) -> Result<u8, String> {
    let req_msg = Message::Request;
    let encoded_req_msg: Vec<u8> = try!(encode(&req_msg, SizeLimit::Infinite)
        .map_err(|e| e.to_string()));

    let mut remaining_len = encoded_req_msg.len();
    while remaining_len > 0 {
        let sent_bytes = try!(socket.send_to(&encoded_req_msg, addr)
            .map_err(|e| e.to_string()));
        remaining_len -= sent_bytes;
    }
    info!("Request sent to {}.", addr);

    let mut buf = [0u8; 1600];
    let (len, recv_addr) = try!(socket.recv_from(&mut buf).map_err(|e| e.to_string()));
    assert_eq!(&recv_addr, addr);
    info!("Response received from {}.", addr);

    let resp_msg: Message = try!(decode(&buf[0..len]).map_err(|e| e.to_string()));
    match resp_msg {
        Message::Response { id } => Ok(id),
        _ => Err(format!("Invalid message {:?} from {}", resp_msg, addr)),
    }
}

pub fn connect(host: &str, port: u16) {
    info!("Working in client mode.");
    let remote_ip = resolve(host).unwrap();
    let remote_addr = SocketAddr::new(remote_ip, port);
    info!("Remote server: {}", remote_addr);

    let local_addr: SocketAddr = "0.0.0.0:0".parse::<SocketAddr>().unwrap();
    let socket = UdpSocket::bind(&local_addr).unwrap();

    let id = initiate(&socket, &remote_addr).unwrap();
    info!("Session established. Assigned IP address: 10.10.10.{}.", id);

    info!("Bringing up TUN device tun1.");
    let mut tun = tuntap::Tun::create(1);
    let tun_rawfd = tun.as_raw_fd();
    tun.up(id);
    let tunfd = mio::unix::EventedFd(&tun_rawfd);
    info!("TUN device tun1 initialized. Internal IP: 10.10.10.{}/24.",
          id);

    let poll = mio::Poll::new().unwrap();
    info!("Setting up TUN device for polling.");
    poll.register(&tunfd, TUN, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    info!("Setting up socket for polling.");
    let sockfd = mio::udp::UdpSocket::from_socket(socket).unwrap();
    poll.register(&sockfd, SOCK, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    let mut events = mio::Events::with_capacity(1024);
    let mut buf = [0u8; 1600];

    // RAII so ignore unused variable warning
    let _gw = utils::DefaultGateway::create("10.10.10.1", &format!("{}", remote_addr.ip()));
    let mut encoder = snap::Encoder::new();
    let mut decoder = snap::Decoder::new();

    info!("Ready for transmission.");

    loop {
        if INTERRUPTED.load(Ordering::Relaxed) {
            break;
        }

        poll.poll(&mut events, None).unwrap();

        for event in events.iter() {
            match event.token() {
                SOCK => {
                    let (len, addr) = sockfd.recv_from(&mut buf).unwrap().unwrap();
                    let msg: Message = decode(&buf[0..len]).unwrap();
                    match msg {
                        Message::Request |
                        Message::Response { id: _ } => {
                            panic!("Invalid message {:?} from {}", msg, addr);
                        }
                        Message::Data { id: _, data } => {
                            let decompressed_data = decoder.decompress_vec(&data).unwrap();
                            let data_len = decompressed_data.len();
                            let sent_len = tun.write(&decompressed_data).unwrap();
                            assert_eq!(sent_len, data_len);
                        }
                    }
                }
                TUN => {
                    let len: usize = tun.read(&mut buf).unwrap();
                    let data = &buf[0..len];

                    let msg = Message::Data {
                        id: id,
                        data: encoder.compress_vec(data).unwrap(),
                    };
                    let encoded_msg = encode(&msg, SizeLimit::Infinite).unwrap();
                    sockfd.send_to(&encoded_msg, &remote_addr).unwrap().unwrap();
                }
                _ => unreachable!(),
            }
        }
    }
}

pub fn serve(port: u16) {
    if cfg!(not(target_os = "linux")) {
        panic!("Server mode is only available in Linux!");
    }
    info!("Working in server mode.");

    info!("Enabling kernel's IPv4 forwarding.");
    utils::enable_ipv4_forwarding().unwrap();

    info!("Bringing up TUN device tun0.");
    let mut tun = tuntap::Tun::create(0);
    tun.up(1);

    let tun_rawfd = tun.as_raw_fd();
    let tunfd = mio::unix::EventedFd(&tun_rawfd);
    info!("TUN device tun0 initialized. Internal IP: 10.10.10.1/24.");

    let addr = format!("0.0.0.0:{}", port).parse().unwrap();
    let sockfd = mio::udp::UdpSocket::bind(&addr).unwrap();
    info!("Listening on: 0.0.0.0:{}.", port);

    let poll = mio::Poll::new().unwrap();
    poll.register(&sockfd, SOCK, mio::Ready::readable(), mio::PollOpt::level()).unwrap();
    poll.register(&tunfd, TUN, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    let mut events = mio::Events::with_capacity(1024);
    let mut available_ids: Vec<u8> = (2..254).collect();
    let mut client_map: HashMap<u8, SocketAddr> = HashMap::new();

    let mut buf = [0u8; 1600];
    let mut encoder = snap::Encoder::new();
    let mut decoder = snap::Decoder::new();
    info!("Ready for transmission.");

    loop {
        if INTERRUPTED.load(Ordering::Relaxed) {
            break;
        }

        poll.poll(&mut events, None).unwrap();

        for event in events.iter() {
            match event.token() {
                SOCK => {
                    let (len, addr) = sockfd.recv_from(&mut buf).unwrap().unwrap();
                    let msg: Message = decode(&buf[0..len]).unwrap();
                    match msg {
                        Message::Request => {
                            let client_id: u8 = available_ids.pop().unwrap();
                            client_map.insert(client_id, addr);

                            info!("Got request from {}. Assigning IP address: 10.10.10.{}.",
                                  addr,
                                  client_id);

                            let reply = Message::Response { id: client_id };
                            let encoded_reply = encode(&reply, SizeLimit::Infinite).unwrap();
                            let sent_len = sockfd.send_to(&encoded_reply, &addr).unwrap().unwrap();
                            assert_eq!(sent_len, encoded_reply.len());
                        }
                        Message::Response { id: _ } => {
                            warn!("Invalid message {:?} from {}", msg, addr)
                        }
                        Message::Data { id: _, data } => {
                            let decompressed_data = decoder.decompress_vec(&data).unwrap();
                            let data_len = decompressed_data.len();
                            let sent_len = tun.write(&decompressed_data).unwrap();
                            assert_eq!(sent_len, data_len);
                        }
                    }
                }
                TUN => {
                    let len: usize = tun.read(&mut buf).unwrap();
                    let data = &buf[0..len];
                    let client_id: u8 = data[19];

                    match client_map.get(&client_id) {
                        None => warn!("Unknown IP packet from TUN for client {}.", client_id),
                        Some(addr) => {
                            let msg = Message::Data {
                                id: client_id,
                                data: encoder.compress_vec(data).unwrap(),
                            };
                            let encoded_msg = encode(&msg, SizeLimit::Infinite).unwrap();
                            sockfd.send_to(&encoded_msg, addr).unwrap().unwrap();
                        }
                    }
                }
                _ => unreachable!(),
            }
        }
    }
}

#[test]
fn resolve_ip() {
    assert_eq!(resolve("127.0.0.1").unwrap(),
               IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
    assert_eq!(resolve("localhost").unwrap(),
               IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)));
}
<commit_msg>connection: Fix an inconsistent test<commit_after>use std::net::{SocketAddr, IpAddr, Ipv4Addr, Ipv6Addr, UdpSocket};
use std::os::unix::io::AsRawFd;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering, ATOMIC_BOOL_INIT};
use std::io::{Write, Read};
use mio;
use dns_lookup;
use bincode::SizeLimit;
use bincode::rustc_serialize::{encode, decode};
use tuntap;
use utils;
use snap;

pub static INTERRUPTED: AtomicBool = ATOMIC_BOOL_INIT;

#[derive(RustcEncodable, RustcDecodable, PartialEq, Debug)]
enum Message {
    Request,
    Response { id: u8 },
    Data { id: u8, data: Vec<u8> },
}

const TUN: mio::Token = mio::Token(0);
const SOCK: mio::Token = mio::Token(1);

fn resolve(host: &str) -> Result<IpAddr, String> {
    let mut ip_list = try!(dns_lookup::lookup_host(host).map_err(|_| "dns_lookup::lookup_host"));
    let ip = ip_list.next().unwrap().unwrap();
    Ok(ip)
}

fn initiate(socket: &UdpSocket, addr: &SocketAddr) -> Result<u8, String> {
    let req_msg = Message::Request;
    let encoded_req_msg: Vec<u8> = try!(encode(&req_msg, SizeLimit::Infinite)
        .map_err(|e| e.to_string()));

    let mut remaining_len = encoded_req_msg.len();
    while remaining_len > 0 {
        let sent_bytes = try!(socket.send_to(&encoded_req_msg, addr)
            .map_err(|e| e.to_string()));
        remaining_len -= sent_bytes;
    }
    info!("Request sent to {}.", addr);

    let mut buf = [0u8; 1600];
    let (len, recv_addr) = try!(socket.recv_from(&mut buf).map_err(|e| e.to_string()));
    assert_eq!(&recv_addr, addr);
    info!("Response received from {}.", addr);

    let resp_msg: Message = try!(decode(&buf[0..len]).map_err(|e| e.to_string()));
    match resp_msg {
        Message::Response { id } => Ok(id),
        _ => Err(format!("Invalid message {:?} from {}", resp_msg, addr)),
    }
}

pub fn connect(host: &str, port: u16) {
    info!("Working in client mode.");
    let remote_ip = resolve(host).unwrap();
    let remote_addr = SocketAddr::new(remote_ip, port);
    info!("Remote server: {}", remote_addr);

    let local_addr: SocketAddr = "0.0.0.0:0".parse::<SocketAddr>().unwrap();
    let socket = UdpSocket::bind(&local_addr).unwrap();

    let id = initiate(&socket, &remote_addr).unwrap();
    info!("Session established. Assigned IP address: 10.10.10.{}.", id);

    info!("Bringing up TUN device tun1.");
    let mut tun = tuntap::Tun::create(1);
    let tun_rawfd = tun.as_raw_fd();
    tun.up(id);
    let tunfd = mio::unix::EventedFd(&tun_rawfd);
    info!("TUN device tun1 initialized. Internal IP: 10.10.10.{}/24.",
          id);

    let poll = mio::Poll::new().unwrap();
    info!("Setting up TUN device for polling.");
    poll.register(&tunfd, TUN, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    info!("Setting up socket for polling.");
    let sockfd = mio::udp::UdpSocket::from_socket(socket).unwrap();
    poll.register(&sockfd, SOCK, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    let mut events = mio::Events::with_capacity(1024);
    let mut buf = [0u8; 1600];

    // RAII so ignore unused variable warning
    let _gw = utils::DefaultGateway::create("10.10.10.1", &format!("{}", remote_addr.ip()));
    let mut encoder = snap::Encoder::new();
    let mut decoder = snap::Decoder::new();

    info!("Ready for transmission.");

    loop {
        if INTERRUPTED.load(Ordering::Relaxed) {
            break;
        }

        poll.poll(&mut events, None).unwrap();

        for event in events.iter() {
            match event.token() {
                SOCK => {
                    let (len, addr) = sockfd.recv_from(&mut buf).unwrap().unwrap();
                    let msg: Message = decode(&buf[0..len]).unwrap();
                    match msg {
                        Message::Request |
                        Message::Response { id: _ } => {
                            panic!("Invalid message {:?} from {}", msg, addr);
                        }
                        Message::Data { id: _, data } => {
                            let decompressed_data = decoder.decompress_vec(&data).unwrap();
                            let data_len = decompressed_data.len();
                            let sent_len = tun.write(&decompressed_data).unwrap();
                            assert_eq!(sent_len, data_len);
                        }
                    }
                }
                TUN => {
                    let len: usize = tun.read(&mut buf).unwrap();
                    let data = &buf[0..len];

                    let msg = Message::Data {
                        id: id,
                        data: encoder.compress_vec(data).unwrap(),
                    };
                    let encoded_msg = encode(&msg, SizeLimit::Infinite).unwrap();
                    sockfd.send_to(&encoded_msg, &remote_addr).unwrap().unwrap();
                }
                _ => unreachable!(),
            }
        }
    }
}

pub fn serve(port: u16) {
    if cfg!(not(target_os = "linux")) {
        panic!("Server mode is only available in Linux!");
    }
    info!("Working in server mode.");

    info!("Enabling kernel's IPv4 forwarding.");
    utils::enable_ipv4_forwarding().unwrap();

    info!("Bringing up TUN device tun0.");
    let mut tun = tuntap::Tun::create(0);
    tun.up(1);

    let tun_rawfd = tun.as_raw_fd();
    let tunfd = mio::unix::EventedFd(&tun_rawfd);
    info!("TUN device tun0 initialized. Internal IP: 10.10.10.1/24.");

    let addr = format!("0.0.0.0:{}", port).parse().unwrap();
    let sockfd = mio::udp::UdpSocket::bind(&addr).unwrap();
    info!("Listening on: 0.0.0.0:{}.", port);

    let poll = mio::Poll::new().unwrap();
    poll.register(&sockfd, SOCK, mio::Ready::readable(), mio::PollOpt::level()).unwrap();
    poll.register(&tunfd, TUN, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    let mut events = mio::Events::with_capacity(1024);
    let mut available_ids: Vec<u8> = (2..254).collect();
    let mut client_map: HashMap<u8, SocketAddr> = HashMap::new();

    let mut buf = [0u8; 1600];
    let mut encoder = snap::Encoder::new();
    let mut decoder = snap::Decoder::new();
    info!("Ready for transmission.");

    loop {
        if INTERRUPTED.load(Ordering::Relaxed) {
            break;
        }

        poll.poll(&mut events, None).unwrap();

        for event in events.iter() {
            match event.token() {
                SOCK => {
                    let (len, addr) = sockfd.recv_from(&mut buf).unwrap().unwrap();
                    let msg: Message = decode(&buf[0..len]).unwrap();
                    match msg {
                        Message::Request => {
                            let client_id: u8 = available_ids.pop().unwrap();
                            client_map.insert(client_id, addr);

                            info!("Got request from {}. Assigning IP address: 10.10.10.{}.",
                                  addr,
                                  client_id);

                            let reply = Message::Response { id: client_id };
                            let encoded_reply = encode(&reply, SizeLimit::Infinite).unwrap();
                            let sent_len = sockfd.send_to(&encoded_reply, &addr).unwrap().unwrap();
                            assert_eq!(sent_len, encoded_reply.len());
                        }
                        Message::Response { id: _ } => {
                            warn!("Invalid message {:?} from {}", msg, addr)
                        }
                        Message::Data { id: _, data } => {
                            let decompressed_data = decoder.decompress_vec(&data).unwrap();
                            let data_len = decompressed_data.len();
                            let sent_len = tun.write(&decompressed_data).unwrap();
                            assert_eq!(sent_len, data_len);
                        }
                    }
                }
                TUN => {
                    let len: usize = tun.read(&mut buf).unwrap();
                    let data = &buf[0..len];
                    let client_id: u8 = data[19];

                    match client_map.get(&client_id) {
                        None => warn!("Unknown IP packet from TUN for client {}.", client_id),
                        Some(addr) => {
                            let msg = Message::Data {
                                id: client_id,
                                data: encoder.compress_vec(data).unwrap(),
                            };
                            let encoded_msg = encode(&msg, SizeLimit::Infinite).unwrap();
                            sockfd.send_to(&encoded_msg, addr).unwrap().unwrap();
                        }
                    }
                }
                _ => unreachable!(),
            }
        }
    }
}

#[test]
fn resolve_ip() {
    assert_eq!(resolve("127.0.0.1").unwrap(),
               IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
}
<|endoftext|>


<commit_before>use std::net::{SocketAddr};
use std::time::{Duration, Instant};

use tokio_core::net::TcpStream;
use tokio_core::reactor::Core;
use tokio_core::io::{Io, Framed};
use tokio_timer::{Timer};
use futures::Future;
use futures::Sink;
use futures::Stream;
use futures::sync::mpsc;

use mqtt3::{Message, QoS, TopicPath};

use error::*;
use packet::*;
use codec::MqttCodec;
use clientoptions::MqttOptions;

pub struct Connection {
    pub addr: SocketAddr,
    pub state: MqttState,
    pub opts: MqttOptions,
    pub initial_connect: bool,
    pub await_pingresp: bool,
    pub last_flush: Instant,

    pub reactor: Core,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MqttState {
    Handshake,
    Connected,
    Disconnected,
}

#[derive(Debug)]
pub enum NetworkRequest {
    Subscribe(Vec<(TopicPath, QoS)>),
    Publish(Message),
    Retransmit,
    Shutdown,
    Disconnect,
    Ping,
}

#[derive(Debug)]
pub enum NetworkNotification {
    Disconnected,
    Connected,
}

// DESIGN: Initial connect status should be immediately known.
//        Intermediate disconnections should be automatically reconnected
fn _try_reconnect(addr: SocketAddr, reactor: &mut Core) -> Result<Framed<TcpStream, MqttCodec>, Error> {

    let connect = generate_connect_packet("".to_string(), true, None, None);

    let f_response = TcpStream::connect(&addr, &reactor.handle()).and_then(|connection| {
        let framed = connection.framed(MqttCodec);
        let f1 = framed.send(connect);

        f1.and_then(|framed| {
                framed.into_future()
                    .and_then(|(res, stream)| Ok((res, stream)))
                    .map_err(|(err, _stream)| err)
            })
            .boxed()
    });

    let response = reactor.run(f_response);
    let (packet, frame) = response?;
    println!("{:?}", packet);
    Ok(frame)
}

impl Connection {
    pub fn start(addr: SocketAddr, opts: MqttOptions) -> Result<(Self, Framed<TcpStream, MqttCodec>), Error> {
        let mut reactor = Core::new().unwrap();

        let framed = _try_reconnect(addr, &mut reactor)?;

        let connection = Connection {
            addr: addr,
            state: MqttState::Disconnected,
            opts: opts,
            initial_connect: true,
            await_pingresp: false,
            last_flush: Instant::now(),

            reactor: reactor,
        };

        Ok((connection, framed))
    }

    pub fn run(&mut self, framed: Framed<TcpStream, MqttCodec>) -> Result<(), Error> {
        let (mut sender, receiver) = framed.split();
        let rx_future = receiver.for_each(|msg| {
            print!("{:?}", msg);
            Ok(())
        }).map_err(|e| Error::Io(e));

        let (mut sender_tx, sender_rx) = mpsc::channel::<NetworkRequest>(1);

        // Ping timer
        let timer = Timer::default();
        let interval = timer.interval(Duration::new(5, 0));
        let timer_future = interval.for_each(|_| {
                let ref mut sender_tx = sender_tx;
                sender_tx.send(NetworkRequest::Ping).wait().unwrap();
                Ok(())
            }).map_err(|e| Error::Timer(e));

        // Sender which does network writes
        let sender_future = sender_rx.for_each(move |r| {
            let packet = match r {
                NetworkRequest::Ping => {
                    println!("{:?}", r);
                    generate_pingreq_packet()
                }
                _ => panic!("Misc")
            };

            let _ = (&mut sender).send(packet).wait();
            Ok(())
        }).map_err(|_|Error::Sender);

        let mqtt_future =  timer_future.join3(rx_future, sender_future);

        let _ = self.reactor.run(mqtt_future)?;
        Ok(())
    }
}
<commit_msg>take keep_alive from options<commit_after>use std::net::{SocketAddr};
use std::time::{Duration, Instant};

use tokio_core::net::TcpStream;
use tokio_core::reactor::Core;
use tokio_core::io::{Io, Framed};
use tokio_timer::{Timer};
use futures::Future;
use futures::Sink;
use futures::Stream;
use futures::sync::mpsc;

use mqtt3::{Message, QoS, TopicPath};

use error::*;
use packet::*;
use codec::MqttCodec;
use clientoptions::MqttOptions;

pub struct Connection {
    pub addr: SocketAddr,
    pub state: MqttState,
    pub opts: MqttOptions,
    pub initial_connect: bool,
    pub await_pingresp: bool,
    pub last_flush: Instant,

    pub reactor: Core,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MqttState {
    Handshake,
    Connected,
    Disconnected,
}

#[derive(Debug)]
pub enum NetworkRequest {
    Subscribe(Vec<(TopicPath, QoS)>),
    Publish(Message),
    Retransmit,
    Shutdown,
    Disconnect,
    Ping,
}

#[derive(Debug)]
pub enum NetworkNotification {
    Disconnected,
    Connected,
}

// DESIGN: Initial connect status should be immediately known.
//        Intermediate disconnections should be automatically reconnected
fn _try_reconnect(addr: SocketAddr, reactor: &mut Core) -> Result<Framed<TcpStream, MqttCodec>, Error> {

    let connect = generate_connect_packet("".to_string(), true, None, None);

    let f_response = TcpStream::connect(&addr, &reactor.handle()).and_then(|connection| {
        let framed = connection.framed(MqttCodec);
        let f1 = framed.send(connect);

        f1.and_then(|framed| {
                framed.into_future()
                    .and_then(|(res, stream)| Ok((res, stream)))
                    .map_err(|(err, _stream)| err)
            })
            .boxed()
    });

    let response = reactor.run(f_response);
    let (packet, frame) = response?;
    println!("{:?}", packet);
    Ok(frame)
}

impl Connection {
    pub fn start(addr: SocketAddr, opts: MqttOptions) -> Result<(Self, Framed<TcpStream, MqttCodec>), Error> {
        let mut reactor = Core::new().unwrap();

        let framed = _try_reconnect(addr, &mut reactor)?;

        let connection = Connection {
            addr: addr,
            state: MqttState::Disconnected,
            opts: opts,
            initial_connect: true,
            await_pingresp: false,
            last_flush: Instant::now(),

            reactor: reactor,
        };

        Ok((connection, framed))
    }

    pub fn run(&mut self, framed: Framed<TcpStream, MqttCodec>) -> Result<(), Error> {
        let (mut sender, receiver) = framed.split();
        let rx_future = receiver.for_each(|msg| {
            print!("{:?}", msg);
            Ok(())
        }).map_err(|e| Error::Io(e));

        let (mut sender_tx, sender_rx) = mpsc::channel::<NetworkRequest>(1);

        // Ping timer
        let timer = Timer::default();

        let keep_alive = if self.opts.keep_alive.is_some() {
            self.opts.keep_alive.unwrap() as u64
        } else {
            1000
        };

        let interval = timer.interval(Duration::new(keep_alive, 0));
        let timer_future = interval.for_each(|_| {
                let ref mut sender_tx = sender_tx;
                sender_tx.send(NetworkRequest::Ping).wait().unwrap();
                Ok(())
            }).map_err(|e| Error::Timer(e));

        // Sender which does network writes
        let sender_future = sender_rx.for_each(move |r| {
            let packet = match r {
                NetworkRequest::Ping => {
                    println!("{:?}", r);
                    generate_pingreq_packet()
                }
                _ => panic!("Misc")
            };

            let _ = (&mut sender).send(packet).wait();
            Ok(())
        }).map_err(|_|Error::Sender);

        let mqtt_future =  timer_future.join3(rx_future, sender_future);

        let _ = self.reactor.run(mqtt_future)?;
        Ok(())
    }
}
<|endoftext|>