<commit_before>// "Tifflin" Kernel - ISO9660 Filesystem Driver
// - By John Hodge (thePowersGang)
//
// Modules/fs_iso9660/lib.rs
#![feature(no_std,core,linkage)]
#![no_std]
#[macro_use] extern crate core;
#[macro_use] extern crate kernel;
use kernel::prelude::*;

use kernel::vfs::{self, mount, node};
use kernel::metadevs::storage::{self,VolumeHandle,SizePrinter};
use kernel::lib::mem::aref::{ArefInner,ArefBorrow};
use kernel::lib::byteorder::{ByteOrder,LittleEndian};
use kernel::lib::byte_str::{ByteStr,ByteString};

module_define!{FS_ISO9660, [VFS], init}

//mod ondisk;

struct Driver;
static S_DRIVER: Driver = Driver;

struct Instance(ArefInner<InstanceInner>);
impl ::core::ops::Deref for Instance {
	type Target = InstanceInner;
	fn deref(&self) -> &InstanceInner { &self.0 }
}

struct InstanceInner
{
	vh: VolumeHandle,
	lb_size: usize,
	root_lba: u32,
	root_size: u32,
}

fn init()
{
	let h = mount::DriverRegistration::new("iso9660", &S_DRIVER);
	// TODO: Remember the registration for unloading
	::core::mem::forget(h);
}

impl mount::Driver for Driver
{
	fn detect(&self, vol: &VolumeHandle) -> vfs::Result<usize> {
		let bs = vol.block_size() as u64;
		let blk = {
			let mut block: Vec<_> = (0 .. bs).map(|_|0).collect();
			try!(vol.read_blocks(32*1024 / bs, &mut block));
			block
			};
		if &blk[1..6] == b"CD001" {
			Ok(1)
		}
		else {
			Ok(0)
		}
	}
	fn mount(&self, vol: VolumeHandle) -> vfs::Result<Box<mount::Filesystem>> {
		if 2048 % vol.block_size() != 0 {
			return Err( vfs::Error::Unknown("Can't mount ISO9660 with sector size not a factor of 2048"/*, vol.block_size()*/) );
		}
		let scale = 2048 / vol.block_size();

		let mut block: Vec<u8> = (0 .. 2048).map(|_|0).collect();
		for sector in (16 .. )
		{
			try!(vol.read_blocks((sector*scale) as u64, &mut block));
			if &block[1..6] != b"CD001" {
				return Err( vfs::Error::Unknown("Invalid volume descriptor present") );
			}
			else if block[0] == 255 {
				return Err( vfs::Error::Unknown("Can't find ISO9660 primary volume descriptor") );
			}
			else if block[0] == 0x01 {
				// Found it!
				break ;
			}
			else {
				// Try the next one
			}
		}
		::kernel::logging::hex_dump("ISO966 PVD", &block);

		// Obtain the logical block size (different from medium sector size)
		let lb_size = LittleEndian::read_u16(&block[128..]);
		// Extract the root directory entry
		// - We want the LBA and byte length
		let root_lba  = LittleEndian::read_u32(&block[156+ 2..]);
		let root_size = LittleEndian::read_u32(&block[156+10..]);

		log_debug!("lb_size = {}, root = {:#x} + {:#x} bytes", lb_size, root_lba, root_size);

		Ok( Box::new( Instance(
			unsafe { ArefInner::new( InstanceInner {
				vh: vol,
				lb_size: lb_size as usize,
				root_lba: root_lba,
				root_size: root_size,
				} ) }
			) ) )
	}
}

impl mount::Filesystem for Instance
{
	fn root_inode(&self) -> node::InodeId {
		self.root_lba as node::InodeId
	}
	fn get_node_by_inode(&self, id: node::InodeId) -> Option<node::Node> {
		if id == self.root_lba as node::InodeId {
			Some(Dir::new_node(self.0.borrow(), self.root_lba, self.root_size) )
		}
		else {
			// Look up (or read) parent directory to obtain the
			todo!("Instance::get_node_by_inode({:#x})", id);
		}
	}
}
impl InstanceInner
{
	fn read_sector(&self, sector: u32) -> Result<Vec<u8>, storage::IoError> {
		// TODO: Cache sector data
		// - I would like to keep Rc<[u8]> in a VecMap or similar, keep four or so around as a LRU cache.
		let mut ret: Vec<u8> = ::core::iter::repeat(0).take(self.lb_size).collect();
		try!( self.vh.read_blocks( (sector as usize * self.lb_size / self.vh.block_size()) as u64, &mut ret) );
		Ok( ret )
	}
}

struct Dir
{
	fs: ArefBorrow<InstanceInner>,
	first_lba: u32,
	size: u32,
}
impl Dir
{
	fn new_node(fs: ArefBorrow<InstanceInner>, first_lba: u32, size: u32) -> node::Node {
		node::Node::Dir( Box::new( Dir {
			fs: fs,
			first_lba: first_lba,
			size: size,
			} ) )
	}
}
impl node::NodeBase for Dir
{
	fn get_id(&self) -> node::InodeId { todo!("Dir::get_id") }
}
impl node::Dir for Dir
{
	fn lookup(&self, name: &ByteStr) -> node::Result<node::InodeId> {
		todo!("Dir::lookup({:?})", name)
	}
	fn read(&self, ofs: usize, items: &mut [(node::InodeId,ByteString)]) -> node::Result<(usize,usize)> {
		let (end_sect,end_ofs) = (self.size as usize / self.fs.lb_size, self.size as usize % self.fs.lb_size);
		let mut sector = ofs / self.fs.lb_size;
		let mut ofs = ofs % self.fs.lb_size;

		let mut count = 0;

		let mut data = try!(self.fs.read_sector(self.first_lba + sector as u32));
		while !(sector == end_sect && ofs >= end_ofs)
		{
			let len = data[ofs] as usize;
			if len == 0 {
				ofs += 1;
			}
			else {
				let ent = &data[ofs .. ofs + len];
				ofs += len;

				let namelen = ent[32] as usize;
				let name = &ent[33 .. 33 + namelen];
				let start = LittleEndian::read_u32(&ent[2..]);

				items[count] = (start as node::InodeId, ByteString::from(name));
				count += 1;
				if count == items.len() {
					break;
				}
			}
			if ofs >= self.fs.lb_size {
				sector += 1;
				ofs = 0;
				data = try!(self.fs.read_sector(self.first_lba + sector as u32));
			}
		}

		Ok( (sector*self.fs.lb_size + ofs, count) )
	}

	fn create(&self, name: &ByteStr, nodetype: node::NodeType) -> node::Result<node::InodeId> {
		Err( node::IoError::ReadOnly )
	}
	fn link(&self, name: &ByteStr, inode: node::InodeId) -> node::Result<()> {
		Err( node::IoError::ReadOnly )
	}
	fn unlink(&self, name: &ByteStr) -> node::Result<()> {
		Err( node::IoError::ReadOnly )
	}
}

<commit_msg>Modules/ISO9660 - Cache metadata blocks to reduce IO<commit_after>// "Tifflin" Kernel - ISO9660 Filesystem Driver
// - By John Hodge (thePowersGang)
//
// Modules/fs_iso9660/lib.rs
#![feature(no_std,core,linkage)]
#![no_std]
#[macro_use] extern crate core;
#[macro_use] extern crate kernel;
use kernel::prelude::*;

use kernel::vfs::{self, mount, node};
use kernel::metadevs::storage::{self,VolumeHandle,SizePrinter};
use kernel::lib::mem::aref::{ArefInner,ArefBorrow};
use kernel::lib::byteorder::{ByteOrder,LittleEndian};
use kernel::lib::byte_str::{ByteStr,ByteString};
use kernel::lib::mem::Arc;

module_define!{FS_ISO9660, [VFS], init}

//mod ondisk;

struct Driver;
static S_DRIVER: Driver = Driver;

struct Instance(ArefInner<InstanceInner>);
impl ::core::ops::Deref for Instance {
	type Target = InstanceInner;
	fn deref(&self) -> &InstanceInner { &self.0 }
}

struct InstanceInner
{
	vh: VolumeHandle,
	lb_size: usize,
	root_lba: u32,
	root_size: u32,

	lru_blocks: ::kernel::sync::Mutex< [Option<CachedBlock>; 6] >,
}

struct CachedBlock {
	lba: u32,
	time: ::kernel::time::TickCount,
	data: Arc<[u8]>,
}

fn init()
{
	let h = mount::DriverRegistration::new("iso9660", &S_DRIVER);
	// TODO: Remember the registration for unloading
	::core::mem::forget(h);
}

impl mount::Driver for Driver
{
	fn detect(&self, vol: &VolumeHandle) -> vfs::Result<usize> {
		let bs = vol.block_size() as u64;
		let blk = {
			let mut block: Vec<_> = (0 .. bs).map(|_|0).collect();
			try!(vol.read_blocks(32*1024 / bs, &mut block));
			block
			};
		if &blk[1..6] == b"CD001" {
			Ok(1)
		}
		else {
			Ok(0)
		}
	}
	fn mount(&self, vol: VolumeHandle) -> vfs::Result<Box<mount::Filesystem>> {
		if 2048 % vol.block_size() != 0 {
			return Err( vfs::Error::Unknown("Can't mount ISO9660 with sector size not a factor of 2048"/*, vol.block_size()*/) );
		}
		let scale = 2048 / vol.block_size();

		let mut block: Vec<u8> = (0 .. 2048).map(|_|0).collect();
		for sector in (16 .. )
		{
			try!(vol.read_blocks((sector*scale) as u64, &mut block));
			if &block[1..6] != b"CD001" {
				return Err( vfs::Error::Unknown("Invalid volume descriptor present") );
			}
			else if block[0] == 255 {
				return Err( vfs::Error::Unknown("Can't find ISO9660 primary volume descriptor") );
			}
			else if block[0] == 0x01 {
				// Found it!
				break ;
			}
			else {
				// Try the next one
			}
		}
		::kernel::logging::hex_dump("ISO966 PVD", &block);

		// Obtain the logical block size (different from medium sector size)
		let lb_size = LittleEndian::read_u16(&block[128..]);
		// Extract the root directory entry
		// - We want the LBA and byte length
		let root_lba  = LittleEndian::read_u32(&block[156+ 2..]);
		let root_size = LittleEndian::read_u32(&block[156+10..]);

		log_debug!("lb_size = {}, root = {:#x} + {:#x} bytes", lb_size, root_lba, root_size);

		Ok( Box::new( Instance(
			unsafe { ArefInner::new( InstanceInner {
				vh: vol,
				lb_size: lb_size as usize,
				root_lba: root_lba,
				root_size: root_size,
				lru_blocks: ::kernel::sync::Mutex::new([
					None, None,
					None, None,
					None, None,
					]),
				} ) }
			) ) )
	}
}

impl mount::Filesystem for Instance
{
	fn root_inode(&self) -> node::InodeId {
		self.root_lba as node::InodeId
	}
	fn get_node_by_inode(&self, id: node::InodeId) -> Option<node::Node> {
		if id == self.root_lba as node::InodeId {
			Some(Dir::new_node(self.0.borrow(), self.root_lba, self.root_size) )
		}
		else {
			// Look up (or read) parent directory to obtain the
			todo!("Instance::get_node_by_inode({:#x})", id);
		}
	}
}
impl InstanceInner
{
	fn read_sector(&self, sector: u32, buf: &mut [u8]) -> Result<(), storage::IoError> {
		try!( self.vh.read_blocks( (sector as usize * self.lb_size / self.vh.block_size()) as u64, buf) );
		Ok( () )

	}
	fn get_sector(&self, sector: u32) -> Result<Arc<[u8]>, storage::IoError> {
		assert!(sector > 0);
		// TODO: Cache sector data
		// - I would like to keep Rc<[u8]> in a VecMap or similar, keep four or so around as a LRU cache.
		let mut lh = self.lru_blocks.lock();
		let (mut oldest_i, mut oldest_ts) = (0,!0);
		for (i,e) in lh.iter().enumerate()
		{
			match *e
			{
			Some(ref e) => {
				if e.time < oldest_ts {
					oldest_i = i;
					oldest_ts = e.time;
				}
				if e.lba == sector {
					return Ok(e.data.clone());
				}
				},
			None => {
				oldest_i = i;
				oldest_ts = 0;
				},
			}
		}

		let mut data: Arc<[u8]> = Arc::from_iter( (0 .. self.lb_size).map(|_| 0u8) );
		try!(self.read_sector(sector, ::kernel::lib::mem::arc::get_mut(&mut data).unwrap()));

		lh[oldest_i] = Some(CachedBlock {
			time: ::kernel::time::ticks(),
			lba: sector,
			data: data.clone(),
			});
		Ok( data )
	}
}

struct Dir
{
	fs: ArefBorrow<InstanceInner>,
	first_lba: u32,
	size: u32,
}
impl Dir
{
	fn new_node(fs: ArefBorrow<InstanceInner>, first_lba: u32, size: u32) -> node::Node {
		node::Node::Dir( Box::new( Dir {
			fs: fs,
			first_lba: first_lba,
			size: size,
			} ) )
	}
}
impl node::NodeBase for Dir
{
	fn get_id(&self) -> node::InodeId { todo!("Dir::get_id") }
}
impl node::Dir for Dir
{
	fn lookup(&self, name: &ByteStr) -> node::Result<node::InodeId> {
		todo!("Dir::lookup({:?})", name)
	}
	fn read(&self, ofs: usize, items: &mut [(node::InodeId,ByteString)]) -> node::Result<(usize,usize)> {
		let (end_sect,end_ofs) = (self.size as usize / self.fs.lb_size, self.size as usize % self.fs.lb_size);
		let mut sector = ofs / self.fs.lb_size;
		let mut ofs = ofs % self.fs.lb_size;

		let mut count = 0;

		let mut data = try!(self.fs.get_sector(self.first_lba + sector as u32));
		while !(sector == end_sect && ofs >= end_ofs)
		{
			let len = data[ofs] as usize;
			if len == 0 {
				ofs += 1;
			}
			else {
				let ent = &data[ofs .. ofs + len];
				ofs += len;

				let namelen = ent[32] as usize;
				let name = &ent[33 .. 33 + namelen];
				let start = LittleEndian::read_u32(&ent[2..]);

				items[count] = (start as node::InodeId, ByteString::from(name));
				count += 1;
				if count == items.len() {
					break;
				}
			}
			if ofs >= self.fs.lb_size {
				sector += 1;
				ofs = 0;
				data = try!(self.fs.get_sector(self.first_lba + sector as u32));
			}
		}

		Ok( (sector*self.fs.lb_size + ofs, count) )
	}

	fn create(&self, _name: &ByteStr, _nodetype: node::NodeType) -> node::Result<node::InodeId> {
		Err( node::IoError::ReadOnly )
	}
	fn link(&self, _name: &ByteStr, _inode: node::InodeId) -> node::Result<()> {
		Err( node::IoError::ReadOnly )
	}
	fn unlink(&self, _name: &ByteStr) -> node::Result<()> {
		Err( node::IoError::ReadOnly )
	}
}

<|endoftext|>


<commit_before>use super::*;
use std::cell::Cell;

#[test]
fn allocator_param() {
    use crate::alloc::AllocErr;

    // Writing a test of integration between third-party
    // allocators and `RawVec` is a little tricky because the `RawVec`
    // API does not expose fallible allocation methods, so we
    // cannot check what happens when allocator is exhausted
    // (beyond detecting a panic).
    //
    // Instead, this just checks that the `RawVec` methods do at
    // least go through the Allocator API when it reserves
    // storage.

    // A dumb allocator that consumes a fixed amount of fuel
    // before allocation attempts start failing.
    struct BoundedAlloc {
        fuel: Cell<usize>,
    }
    unsafe impl AllocRef for BoundedAlloc {
        fn alloc(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocErr> {
            let size = layout.size();
            if size > self.fuel.get() {
                return Err(AllocErr);
            }
            match Global.alloc(layout) {
                ok @ Ok(_) => {
                    self.fuel.update(|old| old - size);
                    ok
                }
                err @ Err(_) => err,
            }
        }
        unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
            unsafe { Global.dealloc(ptr, layout) }
        }
    }

    let a = BoundedAlloc { fuel: Cell::new(500) };
    let mut v: RawVec<u8, _> = RawVec::with_capacity_in(50, a);
    assert_eq!(v.alloc.fuel.get(), 450);
    v.reserve(50, 150); // (causes a realloc, thus using 50 + 150 = 200 units of fuel)
    assert_eq!(v.alloc.fuel.get(), 250);
}

#[test]
fn reserve_does_not_overallocate() {
    {
        let mut v: RawVec<u32> = RawVec::new();
        // First, `reserve` allocates like `reserve_exact`.
        v.reserve(0, 9);
        assert_eq!(9, v.capacity());
    }

    {
        let mut v: RawVec<u32> = RawVec::new();
        v.reserve(0, 7);
        assert_eq!(7, v.capacity());
        // 97 is more than double of 7, so `reserve` should work
        // like `reserve_exact`.
        v.reserve(7, 90);
        assert_eq!(97, v.capacity());
    }

    {
        let mut v: RawVec<u32> = RawVec::new();
        v.reserve(0, 12);
        assert_eq!(12, v.capacity());
        v.reserve(12, 3);
        // 3 is less than half of 12, so `reserve` must grow
        // exponentially. At the time of writing this test grow
        // factor is 2, so new capacity is 24, however, grow factor
        // of 1.5 is OK too. Hence `>= 18` in assert.
        assert!(v.capacity() >= 12 + 12 / 2);
    }
}
<commit_msg>Added feature flag to use cell_update<commit_after>#![feature(cell_update)]

use super::*;
use std::cell::Cell;

#[test]
fn allocator_param() {
    use crate::alloc::AllocErr;

    // Writing a test of integration between third-party
    // allocators and `RawVec` is a little tricky because the `RawVec`
    // API does not expose fallible allocation methods, so we
    // cannot check what happens when allocator is exhausted
    // (beyond detecting a panic).
    //
    // Instead, this just checks that the `RawVec` methods do at
    // least go through the Allocator API when it reserves
    // storage.

    // A dumb allocator that consumes a fixed amount of fuel
    // before allocation attempts start failing.
    struct BoundedAlloc {
        fuel: Cell<usize>,
    }
    unsafe impl AllocRef for BoundedAlloc {
        fn alloc(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocErr> {
            let size = layout.size();
            if size > self.fuel.get() {
                return Err(AllocErr);
            }
            match Global.alloc(layout) {
                ok @ Ok(_) => {
                    self.fuel.update(|old| old - size);
                    ok
                }
                err @ Err(_) => err,
            }
        }
        unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
            unsafe { Global.dealloc(ptr, layout) }
        }
    }

    let a = BoundedAlloc { fuel: Cell::new(500) };
    let mut v: RawVec<u8, _> = RawVec::with_capacity_in(50, a);
    assert_eq!(v.alloc.fuel.get(), 450);
    v.reserve(50, 150); // (causes a realloc, thus using 50 + 150 = 200 units of fuel)
    assert_eq!(v.alloc.fuel.get(), 250);
}

#[test]
fn reserve_does_not_overallocate() {
    {
        let mut v: RawVec<u32> = RawVec::new();
        // First, `reserve` allocates like `reserve_exact`.
        v.reserve(0, 9);
        assert_eq!(9, v.capacity());
    }

    {
        let mut v: RawVec<u32> = RawVec::new();
        v.reserve(0, 7);
        assert_eq!(7, v.capacity());
        // 97 is more than double of 7, so `reserve` should work
        // like `reserve_exact`.
        v.reserve(7, 90);
        assert_eq!(97, v.capacity());
    }

    {
        let mut v: RawVec<u32> = RawVec::new();
        v.reserve(0, 12);
        assert_eq!(12, v.capacity());
        v.reserve(12, 3);
        // 3 is less than half of 12, so `reserve` must grow
        // exponentially. At the time of writing this test grow
        // factor is 2, so new capacity is 24, however, grow factor
        // of 1.5 is OK too. Hence `>= 18` in assert.
        assert!(v.capacity() >= 12 + 12 / 2);
    }
}
<|endoftext|>


<commit_before>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use anyhow::{bail, Result};
use mio::event::Event;
use mio::net::TcpListener;
use mio::net::TcpStream;
use mio::{Events, Interest, Poll, Token};
use mio_signals::{Signal, SignalSet, Signals};
use serde::de::DeserializeOwned;
use serde::Serialize;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::collections::HashMap;
use std::io::{ErrorKind, Read, Write};
use std::marker::PhantomData;
use std::sync::atomic::{AtomicUsize, Ordering};

use super::CommandHandler;

const BUFFER_SIZE: usize = 8192;
const EOL_CODE: u8 = b'\n';

fn get_next_token() -> Token {
    static TOCKEN_COUNTER: AtomicUsize = AtomicUsize::new(0);
    Token(TOCKEN_COUNTER.fetch_add(1, Ordering::Relaxed))
}

/// This struct listens on a TCP socket, and maintains a number of concurrent connections,
/// receiving serialized JSON representations of `Msg`, passing them to the given
/// `CommandHandler` to obtain responses to be sent as socket flow contol permits.  Note that
/// this implementaion is not specific to (and does not refer to) any particular protocol.
pub struct JsonSocketServer<Msg: DeserializeOwned + Serialize, T: CommandHandler<Msg>> {
    command_handler: T,
    poll: Poll,
    socket: TcpListener,
    socket_token: Token,
    signals: Signals,
    signal_token: Token,
    connection_map: HashMap<Token, Connection>,
    exit_requested: bool,
    phantom: PhantomData<Msg>,
}

impl<Msg: DeserializeOwned + Serialize, T: CommandHandler<Msg>> JsonSocketServer<Msg, T> {
    pub fn new(command_handler: T, mut socket: TcpListener) -> Result<Self> {
        let poll = Poll::new()?;
        let socket_token = get_next_token();
        poll.registry()
            .register(&mut socket, socket_token, Interest::READABLE)?;
        // Create a `Signals` instance that will catch given set of signals for us.
        let signals: SignalSet = Signal::Terminate | Signal::Interrupt;
        let mut signals = Signals::new(signals)?;
        // And register it with our `Poll` instance.
        let signal_token = get_next_token();
        poll.registry()
            .register(&mut signals, signal_token, Interest::READABLE)?;
        Ok(Self {
            command_handler,
            poll,
            socket,
            socket_token,
            signals,
            signal_token,
            connection_map: HashMap::new(),
            exit_requested: false,
            phantom: PhantomData,
        })
    }

    pub fn run_loop(&mut self) -> Result<()> {
        let mut events = Events::with_capacity(1024);
        while !self.exit_requested {
            self.poll.poll(&mut events, None)?;
            for event in events.iter() {
                if event.token() == self.socket_token {
                    self.process_new_connection()?;
                } else if event.token() == self.signal_token {
                    self.process_signals()?;
                } else {
                    match self.process_connection(event) {
                        Ok(shutdown) => {
                            if shutdown {
                                self.shutdown_connection(event)?;
                            }
                        }
                        Err(e) => {
                            log::warn!("Connection {:#X} error: {}", event.token().0, e,);
                            self.shutdown_connection(event)?;
                        }
                    }
                }
            }
        }
        Ok(())
    }

    /// Accept new socket connections, creating new Connection objects.
    fn process_new_connection(&mut self) -> Result<()> {
        loop {
            match self.socket.accept() {
                Ok((mut conn_socket, _addres)) => {
                    let token = get_next_token();
                    log::info!("New connection id:{:#X}", token.0);
                    match self.connection_map.entry(token) {
                        Vacant(entry) => {
                            self.poll.registry().register(
                                &mut conn_socket,
                                token,
                                Interest::READABLE | Interest::WRITABLE,
                            )?;
                            entry.insert(Connection::new(conn_socket));
                        }
                        Occupied(_) => {
                            panic!("JsonSocketServer error: token colision");
                        }
                    };
                }
                Err(err) if err.kind() == ErrorKind::WouldBlock => {
                    // No more connections ready to accept (or spurious poll event).
                    return Ok(());
                }
                Err(err) => bail!("Error accepting TCP connection: {}", err),
            }
        }
    }

    fn process_signals(&mut self) -> Result<()> {
        loop {
            match self.signals.receive()? {
                Some(Signal::Interrupt) => {
                    log::info!("Got interrupt signal");
                    self.exit_requested = true;
                }
                Some(Signal::Terminate) => {
                    log::info!("Got terminate signal");
                    self.exit_requested = true;
                }
                Some(signal) => {
                    log::info!("Got unexpected signal: {:?}", signal);
                }
                None => return Ok(()),
            }
        }
    }

    /// Read and write as much as possible from one particular socket connection.
    fn process_connection(&mut self, event: &Event) -> Result<bool> {
        match self.connection_map.get_mut(&event.token()) {
            Some(conn) => {
                if event.is_writable() {
                    conn.write()?;
                }
                if event.is_readable() {
                    conn.read()?;
                    Self::process_any_requests(conn, &self.command_handler)?;
                }
                // Return whether this connection object should be dropped.
                return Ok((conn.rx_eof && (conn.tx_buf.len() == 0)) || conn.broken);
            }
            None => bail!("Connection don't exist token:{:#X}", event.token().0),
        }
    }

    /// Close a socket connection and remove it from the poll list.
    fn shutdown_connection(&mut self, event: &Event) -> Result<()> {
        log::info!("Closing connection id:{:#X}", event.token().0);
        let mut conn = self
            .connection_map
            .remove(&event.token())
            .expect("Missing connection this should never happend!!!");
        self.poll.registry().deregister(&mut conn.socket)?;
        // As `conn` runs out of scope here, its `drop()` method will close the OS handle, which
        // in turn causes TCP/IP connection shutdown to be signalled to the remote end.
        Ok(())
    }

    /// Check if the buffer contains at least one full JSON request.  If so, remove it from the
    /// buffer, decode and return it.
    fn get_complete_request(conn: &mut Connection) -> Result<Option<Msg>> {
        if let Some(n) = conn.rx_buf.iter().position(|c| *c == EOL_CODE) {
            let res = serde_json::from_slice::<Msg>(&conn.rx_buf[..n])?;
            if n + 1 < conn.rx_buf.len() {
                // Shuffling bytes around in a Vec is expensive, but realistically, as the
                // clients would be waiting for response to each request before sending the next
                // request, this code will rarely if ever execute.
                conn.rx_buf.rotate_left(n + 1);
            }
            conn.rx_buf.resize(conn.rx_buf.len() - n - 1, 0);
            return Ok(Some(res));
        }
        Ok(None)
    }

    // Look for any completely received requests in the rx_buf, and handle them one by one.
    fn process_any_requests(conn: &mut Connection, command_handler: &T) -> Result<()> {
        while let Some(request) = Self::get_complete_request(conn)? {
            // One complete request received, execute it.
            let resp = command_handler.execute_cmd(&request)?;
            // Encode response into tx_buf.
            serde_json::to_writer(&mut conn.tx_buf, &resp)?;
            conn.tx_buf.push(EOL_CODE);
            // Transmit as much as possible without blocking, leaving any remnant in
            // tx_buf.  poll() will tell us when more can be written.
            conn.write()?;
        }
        Ok(())
    }
}

/// Represents one connection with a remote OpenTitan tool invocation.
struct Connection {
    socket: TcpStream,
    /// Outgoing data waiting to be written when the socket permits.
    tx_buf: Vec<u8>,
    /// Data received from the remote end, but not yet decoded into `Msg`.
    rx_buf: Vec<u8>,
    /// The remote end indicated end-of-stream.  After processing any remaning data in `rx_buf`,
    /// this Connection should be gracefully shut down and dropped.
    rx_eof: bool,
    /// Some error happened during writing or reading from the socket, we cannot meaningfully
    /// continue processing, and the connection should be dropped as soon as possible.
    broken: bool,
}

impl Connection {
    fn new(soc: TcpStream) -> Self {
        Self {
            socket: soc,
            tx_buf: Vec::new(),
            rx_buf: Vec::new(),
            rx_eof: false,
            broken: false,
        }
    }

    // Fill rx_buf with as much data as is available on the socket.
    fn read(&mut self) -> Result<()> {
        let mut rx_buf_len: usize = self.rx_buf.len();
        loop {
            self.rx_buf.resize(rx_buf_len + BUFFER_SIZE, 0);
            match self.socket.read(&mut self.rx_buf[rx_buf_len..]) {
                Ok(0) => {
                    self.rx_eof = true;
                    break;
                }
                Ok(n) => {
                    rx_buf_len += n;
                }
                Err(err) => {
                    if err.kind() != ErrorKind::WouldBlock {
                        self.broken = true;
                    }
                    break; // Break out of loop, also on expected WouldBlock
                }
            }
        }
        self.rx_buf.resize(rx_buf_len, 0);
        Ok(())
    }

    // Transmit as much data out of tx_buf as socket will allow.
    fn write(&mut self) -> Result<()> {
        while self.tx_buf.len() > 0 {
            match self.socket.write(&self.tx_buf) {
                Ok(n) => {
                    if n < self.tx_buf.len() {
                        // Shuffling bytes around in a Vec is expensive, but realistically, as
                        // the clients would be waiting for response to each request before
                        // sending the next request, it is unlikely that the OS transmit buffer
                        // would ever fill up and cause partial writes.
                        self.tx_buf.rotate_left(n);
                    }
                    self.tx_buf.resize(self.tx_buf.len() - n, 0);
                }
                Err(err) => {
                    if err.kind() != ErrorKind::WouldBlock {
                        self.broken = true;
                    }
                    break;
                }
            }
        }
        Ok(())
    }
}
<commit_msg>[opentitantool] Tolerate being interrupted in epoll()<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use anyhow::{bail, Result};
use mio::event::Event;
use mio::net::TcpListener;
use mio::net::TcpStream;
use mio::{Events, Interest, Poll, Token};
use mio_signals::{Signal, SignalSet, Signals};
use serde::de::DeserializeOwned;
use serde::Serialize;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::collections::HashMap;
use std::io::{ErrorKind, Read, Write};
use std::marker::PhantomData;
use std::sync::atomic::{AtomicUsize, Ordering};

use super::CommandHandler;

const BUFFER_SIZE: usize = 8192;
const EOL_CODE: u8 = b'\n';

fn get_next_token() -> Token {
    static TOCKEN_COUNTER: AtomicUsize = AtomicUsize::new(0);
    Token(TOCKEN_COUNTER.fetch_add(1, Ordering::Relaxed))
}

/// This struct listens on a TCP socket, and maintains a number of concurrent connections,
/// receiving serialized JSON representations of `Msg`, passing them to the given
/// `CommandHandler` to obtain responses to be sent as socket flow contol permits.  Note that
/// this implementaion is not specific to (and does not refer to) any particular protocol.
pub struct JsonSocketServer<Msg: DeserializeOwned + Serialize, T: CommandHandler<Msg>> {
    command_handler: T,
    poll: Poll,
    socket: TcpListener,
    socket_token: Token,
    signals: Signals,
    signal_token: Token,
    connection_map: HashMap<Token, Connection>,
    exit_requested: bool,
    phantom: PhantomData<Msg>,
}

impl<Msg: DeserializeOwned + Serialize, T: CommandHandler<Msg>> JsonSocketServer<Msg, T> {
    pub fn new(command_handler: T, mut socket: TcpListener) -> Result<Self> {
        let poll = Poll::new()?;
        let socket_token = get_next_token();
        poll.registry()
            .register(&mut socket, socket_token, Interest::READABLE)?;
        // Create a `Signals` instance that will catch given set of signals for us.
        let signals: SignalSet = Signal::Terminate | Signal::Interrupt;
        let mut signals = Signals::new(signals)?;
        // And register it with our `Poll` instance.
        let signal_token = get_next_token();
        poll.registry()
            .register(&mut signals, signal_token, Interest::READABLE)?;
        Ok(Self {
            command_handler,
            poll,
            socket,
            socket_token,
            signals,
            signal_token,
            connection_map: HashMap::new(),
            exit_requested: false,
            phantom: PhantomData,
        })
    }

    pub fn run_loop(&mut self) -> Result<()> {
        let mut events = Events::with_capacity(1024);
        while !self.exit_requested {
            match self.poll.poll(&mut events, None) {
                Ok(()) => (),
                Err(err) if err.kind() == ErrorKind::Interrupted => {
                    continue;
                }
                Err(err) => bail!("poll: {}", err),
            }
            for event in events.iter() {
                if event.token() == self.socket_token {
                    self.process_new_connection()?;
                } else if event.token() == self.signal_token {
                    self.process_signals()?;
                } else {
                    match self.process_connection(event) {
                        Ok(shutdown) => {
                            if shutdown {
                                self.shutdown_connection(event)?;
                            }
                        }
                        Err(e) => {
                            log::warn!("Connection {:#X} error: {}", event.token().0, e,);
                            self.shutdown_connection(event)?;
                        }
                    }
                }
            }
        }
        Ok(())
    }

    /// Accept new socket connections, creating new Connection objects.
    fn process_new_connection(&mut self) -> Result<()> {
        loop {
            match self.socket.accept() {
                Ok((mut conn_socket, _addres)) => {
                    let token = get_next_token();
                    log::info!("New connection id:{:#X}", token.0);
                    match self.connection_map.entry(token) {
                        Vacant(entry) => {
                            self.poll.registry().register(
                                &mut conn_socket,
                                token,
                                Interest::READABLE | Interest::WRITABLE,
                            )?;
                            entry.insert(Connection::new(conn_socket));
                        }
                        Occupied(_) => {
                            panic!("JsonSocketServer error: token colision");
                        }
                    };
                }
                Err(err) if err.kind() == ErrorKind::WouldBlock => {
                    // No more connections ready to accept (or spurious poll event).
                    return Ok(());
                }
                Err(err) => bail!("Error accepting TCP connection: {}", err),
            }
        }
    }

    fn process_signals(&mut self) -> Result<()> {
        loop {
            match self.signals.receive()? {
                Some(Signal::Interrupt) => {
                    log::info!("Got interrupt signal");
                    self.exit_requested = true;
                }
                Some(Signal::Terminate) => {
                    log::info!("Got terminate signal");
                    self.exit_requested = true;
                }
                Some(signal) => {
                    log::info!("Got unexpected signal: {:?}", signal);
                }
                None => return Ok(()),
            }
        }
    }

    /// Read and write as much as possible from one particular socket connection.
    fn process_connection(&mut self, event: &Event) -> Result<bool> {
        match self.connection_map.get_mut(&event.token()) {
            Some(conn) => {
                if event.is_writable() {
                    conn.write()?;
                }
                if event.is_readable() {
                    conn.read()?;
                    Self::process_any_requests(conn, &self.command_handler)?;
                }
                // Return whether this connection object should be dropped.
                return Ok((conn.rx_eof && (conn.tx_buf.len() == 0)) || conn.broken);
            }
            None => bail!("Connection don't exist token:{:#X}", event.token().0),
        }
    }

    /// Close a socket connection and remove it from the poll list.
    fn shutdown_connection(&mut self, event: &Event) -> Result<()> {
        log::info!("Closing connection id:{:#X}", event.token().0);
        let mut conn = self
            .connection_map
            .remove(&event.token())
            .expect("Missing connection this should never happend!!!");
        self.poll.registry().deregister(&mut conn.socket)?;
        // As `conn` runs out of scope here, its `drop()` method will close the OS handle, which
        // in turn causes TCP/IP connection shutdown to be signalled to the remote end.
        Ok(())
    }

    /// Check if the buffer contains at least one full JSON request.  If so, remove it from the
    /// buffer, decode and return it.
    fn get_complete_request(conn: &mut Connection) -> Result<Option<Msg>> {
        if let Some(n) = conn.rx_buf.iter().position(|c| *c == EOL_CODE) {
            let res = serde_json::from_slice::<Msg>(&conn.rx_buf[..n])?;
            if n + 1 < conn.rx_buf.len() {
                // Shuffling bytes around in a Vec is expensive, but realistically, as the
                // clients would be waiting for response to each request before sending the next
                // request, this code will rarely if ever execute.
                conn.rx_buf.rotate_left(n + 1);
            }
            conn.rx_buf.resize(conn.rx_buf.len() - n - 1, 0);
            return Ok(Some(res));
        }
        Ok(None)
    }

    // Look for any completely received requests in the rx_buf, and handle them one by one.
    fn process_any_requests(conn: &mut Connection, command_handler: &T) -> Result<()> {
        while let Some(request) = Self::get_complete_request(conn)? {
            // One complete request received, execute it.
            let resp = command_handler.execute_cmd(&request)?;
            // Encode response into tx_buf.
            serde_json::to_writer(&mut conn.tx_buf, &resp)?;
            conn.tx_buf.push(EOL_CODE);
            // Transmit as much as possible without blocking, leaving any remnant in
            // tx_buf.  poll() will tell us when more can be written.
            conn.write()?;
        }
        Ok(())
    }
}

/// Represents one connection with a remote OpenTitan tool invocation.
struct Connection {
    socket: TcpStream,
    /// Outgoing data waiting to be written when the socket permits.
    tx_buf: Vec<u8>,
    /// Data received from the remote end, but not yet decoded into `Msg`.
    rx_buf: Vec<u8>,
    /// The remote end indicated end-of-stream.  After processing any remaning data in `rx_buf`,
    /// this Connection should be gracefully shut down and dropped.
    rx_eof: bool,
    /// Some error happened during writing or reading from the socket, we cannot meaningfully
    /// continue processing, and the connection should be dropped as soon as possible.
    broken: bool,
}

impl Connection {
    fn new(soc: TcpStream) -> Self {
        Self {
            socket: soc,
            tx_buf: Vec::new(),
            rx_buf: Vec::new(),
            rx_eof: false,
            broken: false,
        }
    }

    // Fill rx_buf with as much data as is available on the socket.
    fn read(&mut self) -> Result<()> {
        let mut rx_buf_len: usize = self.rx_buf.len();
        loop {
            self.rx_buf.resize(rx_buf_len + BUFFER_SIZE, 0);
            match self.socket.read(&mut self.rx_buf[rx_buf_len..]) {
                Ok(0) => {
                    self.rx_eof = true;
                    break;
                }
                Ok(n) => {
                    rx_buf_len += n;
                }
                Err(err) => {
                    if err.kind() != ErrorKind::WouldBlock {
                        self.broken = true;
                    }
                    break; // Break out of loop, also on expected WouldBlock
                }
            }
        }
        self.rx_buf.resize(rx_buf_len, 0);
        Ok(())
    }

    // Transmit as much data out of tx_buf as socket will allow.
    fn write(&mut self) -> Result<()> {
        while self.tx_buf.len() > 0 {
            match self.socket.write(&self.tx_buf) {
                Ok(n) => {
                    if n < self.tx_buf.len() {
                        // Shuffling bytes around in a Vec is expensive, but realistically, as
                        // the clients would be waiting for response to each request before
                        // sending the next request, it is unlikely that the OS transmit buffer
                        // would ever fill up and cause partial writes.
                        self.tx_buf.rotate_left(n);
                    }
                    self.tx_buf.resize(self.tx_buf.len() - n, 0);
                }
                Err(err) => {
                    if err.kind() != ErrorKind::WouldBlock {
                        self.broken = true;
                    }
                    break;
                }
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! `POST /_matrix/identity/*/store-invite`
//!
//! Endpoint to store pending invitations to a user's 3PID.

pub mod v2 {
    //! `/v2/` ([spec])
    //!
    //! [spec]: https://spec.matrix.org/v1.2/identity-service-api/#post_matrixidentityv2store-invite

    use ruma_api::ruma_api;
    use ruma_common::thirdparty::Medium;
    use ruma_identifiers::{MxcUri, RoomAliasId, RoomId, RoomName, UserId};
    use serde::{ser::SerializeSeq, Deserialize, Serialize};

    ruma_api! {
        metadata: {
            description: "Store pending invitations to a user's 3PID.",
            method: POST,
            name: "store_invitation",
            stable_path: "/_matrix/identity/v2/store-invite",
            authentication: AccessToken,
            rate_limited: false,
            added: 1.0,
        }

        request: {
            /// The type of the third party identifier for the invited user.
            ///
            /// Currently, only `Medium::Email` is supported.
            pub medium: &'a Medium,

            /// The email address of the invited user.
            pub address: &'a str,

            /// The Matrix room ID to which the user is invited.
            pub room_id: &'a RoomId,

            /// The Matrix user ID of the inviting user.
            pub sender: &'a UserId,

            /// The Matrix room alias for the room to which the user is invited.
            ///
            /// This should be retrieved from the `m.room.canonical` state event.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub room_alias: Option<&'a RoomAliasId>,

            /// The Content URI for the room to which the user is invited.
            ///
            /// This should be retrieved from the `m.room.avatar` state event.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub room_avatar_url: Option<&'a MxcUri>,

            /// The `join_rule` for the room to which the user is invited.
            ///
            /// This should be retrieved from the `m.room.join_rules` state event.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub room_join_rules: Option<&'a str>,

            /// The name of the room to which the user is invited.
            ///
            /// This should be retrieved from the `m.room.name` state event.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub room_name: Option<&'a RoomName>,

            /// The display name of the user ID initiating the invite.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub sender_display_name: Option<&'a str>,

            /// The Content URI for the avater of the user ID initiating the invite.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub sender_avatar_url: Option<&'a MxcUri>,
        }

        response: {
            /// The generated token.
            ///
            /// Must be a string consisting of the characters `[0-9a-zA-Z.=_-]`. Its length must not
            /// exceed 255 characters and it must not be empty.
            pub token: String,

            /// A list of [server's long-term public key, generated ephemeral public key].
            pub public_keys: PublicKeys,

            /// The generated (redacted) display_name.
            ///
            /// An example is `f...@b...`.
            pub display_name: String,
        }
    }

    impl<'a> Request<'a> {
        /// Creates a new `Request with the given medium, email address, room ID and sender.
        pub fn new(
            medium: &'a Medium,
            address: &'a str,
            room_id: &'a RoomId,
            sender: &'a UserId,
        ) -> Self {
            Self {
                medium,
                address,
                room_id,
                sender,
                room_alias: None,
                room_avatar_url: None,
                room_join_rules: None,
                room_name: None,
                sender_display_name: None,
                sender_avatar_url: None,
            }
        }

        /// Creates a new `Request` with the given email address, room ID and sender.
        pub fn email(address: &'a str, room_id: &'a RoomId, sender: &'a UserId) -> Self {
            Self::new(&Medium::Email, address, room_id, sender)
        }
    }

    impl Response {
        /// Creates a new `Response` with the given token, public keys and display name.
        pub fn new(token: String, public_keys: PublicKeys, display_name: String) -> Self {
            Self { token, public_keys, display_name }
        }
    }

    /// The server's long-term public key and generated ephemeral public key.
    #[derive(Debug, Clone)]
    #[allow(clippy::exhaustive_structs)]
    pub struct PublicKeys {
        /// The server's long-term public key.
        pub server_key: String,

        /// The generated ephemeral public key.
        pub ephemeral_key: String,
    }

    impl<'de> Deserialize<'de> for PublicKeys {
        fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where
            D: serde::Deserializer<'de>,
        {
            let [server_key, ephemeral_key] = <[String; 2]>::deserialize(deserializer)?;

            Ok(Self { server_key, ephemeral_key })
        }
    }

    impl Serialize for PublicKeys {
        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where
            S: serde::Serializer,
        {
            let mut seq = serializer.serialize_seq(Some(2))?;

            seq.serialize_element(&self.server_key)?;
            seq.serialize_element(&self.ephemeral_key)?;

            seq.end()
        }
    }
}
<commit_msg>identity-api: Add room_type to store_invitation::Request<commit_after>//! `POST /_matrix/identity/*/store-invite`
//!
//! Endpoint to store pending invitations to a user's 3PID.

pub mod v2 {
    //! `/v2/` ([spec])
    //!
    //! [spec]: https://spec.matrix.org/v1.2/identity-service-api/#post_matrixidentityv2store-invite

    use ruma_api::ruma_api;
    use ruma_common::{room::RoomType, thirdparty::Medium};
    use ruma_identifiers::{MxcUri, RoomAliasId, RoomId, RoomName, UserId};
    use serde::{ser::SerializeSeq, Deserialize, Serialize};

    ruma_api! {
        metadata: {
            description: "Store pending invitations to a user's 3PID.",
            method: POST,
            name: "store_invitation",
            stable_path: "/_matrix/identity/v2/store-invite",
            authentication: AccessToken,
            rate_limited: false,
            added: 1.0,
        }

        request: {
            /// The type of the third party identifier for the invited user.
            ///
            /// Currently, only `Medium::Email` is supported.
            pub medium: &'a Medium,

            /// The email address of the invited user.
            pub address: &'a str,

            /// The Matrix room ID to which the user is invited.
            pub room_id: &'a RoomId,

            /// The Matrix user ID of the inviting user.
            pub sender: &'a UserId,

            /// The Matrix room alias for the room to which the user is invited.
            ///
            /// This should be retrieved from the `m.room.canonical` state event.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub room_alias: Option<&'a RoomAliasId>,

            /// The Content URI for the room to which the user is invited.
            ///
            /// This should be retrieved from the `m.room.avatar` state event.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub room_avatar_url: Option<&'a MxcUri>,

            /// The `join_rule` for the room to which the user is invited.
            ///
            /// This should be retrieved from the `m.room.join_rules` state event.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub room_join_rules: Option<&'a str>,

            /// The name of the room to which the user is invited.
            ///
            /// This should be retrieved from the `m.room.name` state event.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub room_name: Option<&'a RoomName>,

            /// The type of the room to which the user is invited.
            ///
            /// This should be retrieved from the `m.room.create` state event.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub room_type: Option<&'a RoomType>,

            /// The display name of the user ID initiating the invite.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub sender_display_name: Option<&'a str>,

            /// The Content URI for the avater of the user ID initiating the invite.
            #[serde(skip_serializing_if = "Option::is_none")]
            pub sender_avatar_url: Option<&'a MxcUri>,
        }

        response: {
            /// The generated token.
            ///
            /// Must be a string consisting of the characters `[0-9a-zA-Z.=_-]`. Its length must not
            /// exceed 255 characters and it must not be empty.
            pub token: String,

            /// A list of [server's long-term public key, generated ephemeral public key].
            pub public_keys: PublicKeys,

            /// The generated (redacted) display_name.
            ///
            /// An example is `f...@b...`.
            pub display_name: String,
        }
    }

    impl<'a> Request<'a> {
        /// Creates a new `Request with the given medium, email address, room ID and sender.
        pub fn new(
            medium: &'a Medium,
            address: &'a str,
            room_id: &'a RoomId,
            sender: &'a UserId,
        ) -> Self {
            Self {
                medium,
                address,
                room_id,
                sender,
                room_alias: None,
                room_avatar_url: None,
                room_join_rules: None,
                room_name: None,
                room_type: None,
                sender_display_name: None,
                sender_avatar_url: None,
            }
        }

        /// Creates a new `Request` with the given email address, room ID and sender.
        pub fn email(address: &'a str, room_id: &'a RoomId, sender: &'a UserId) -> Self {
            Self::new(&Medium::Email, address, room_id, sender)
        }
    }

    impl Response {
        /// Creates a new `Response` with the given token, public keys and display name.
        pub fn new(token: String, public_keys: PublicKeys, display_name: String) -> Self {
            Self { token, public_keys, display_name }
        }
    }

    /// The server's long-term public key and generated ephemeral public key.
    #[derive(Debug, Clone)]
    #[allow(clippy::exhaustive_structs)]
    pub struct PublicKeys {
        /// The server's long-term public key.
        pub server_key: String,

        /// The generated ephemeral public key.
        pub ephemeral_key: String,
    }

    impl<'de> Deserialize<'de> for PublicKeys {
        fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where
            D: serde::Deserializer<'de>,
        {
            let [server_key, ephemeral_key] = <[String; 2]>::deserialize(deserializer)?;

            Ok(Self { server_key, ephemeral_key })
        }
    }

    impl Serialize for PublicKeys {
        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where
            S: serde::Serializer,
        {
            let mut seq = serializer.serialize_seq(Some(2))?;

            seq.serialize_element(&self.server_key)?;
            seq.serialize_element(&self.ephemeral_key)?;

            seq.end()
        }
    }
}
<|endoftext|>


<commit_before>use std::io::BufRead;
use std::io::BufReader;
use std::io::Read;

use common::Solution;

#[derive(Default, Debug)]
pub struct Day12 {
    productions: [bool; 32],
}

fn char_bool(c: char) -> bool {
    match c {
        '#' => true,
        '.' => false,
        _ => panic!("Invalid input {}", c)
    }
}

fn print_state(state: &[(i32, bool)]) {
    for &(_, b) in state {
        if b {
            print!("#");
        } else {
            print!(".");
        }
    }
    println!();
}

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

    fn read_input(&mut self, input: &mut Read) -> Vec<(i32, bool)> {
        let mut state = Vec::new();
        let mut reader = BufReader::new(input);
        {
            let mut line = String::new();
            reader.read_line(&mut line).unwrap();

            for (idx, c) in line.trim().chars().skip("initial state: ".len()).enumerate() {
                state.push((idx as i32, char_bool(c)));
            }
        }

        for line in reader.lines() {
            let line = line.unwrap();

            if line.is_empty() {
                continue;
            }

            let mut index = 0;
            for c in line.chars().take(5) {
                index <<= 1;
                if char_bool(c) {
                    index |= 1;
                }
            }

            self.productions[index] = char_bool(line.chars().last().unwrap());
        }

        state
    }

    fn simulate(&self, state: &[(i32, bool)]) -> Vec<(i32, bool)> {
        let mut new_state = Vec::new();
        let mut index = 0;
        let mut last_idx = None;
        for &(idx, b) in state {
            index = (index << 1) & 0x1f;
            if b {
                index |= 1;
            }

            if self.productions[index] {
                new_state.push((idx - 2, true));
            } else if !new_state.is_empty() {
                new_state.push((idx - 2, false));
            }
            last_idx = Some(idx);
        }
        if let Some(idx) = last_idx {
            for i in 1..=4 {
                index = (index << 1) & 0x1f;

                if self.productions[index] {
                    new_state.push((idx + i - 2, true));
                } else if !new_state.is_empty() {
                    new_state.push((idx + i - 2, false));
                }
            }
        }

        let new_len = new_state.len() - new_state.iter().rev()
            .take_while(|&&(_, x)| !x).count();
        new_state.truncate(new_len);

        new_state
    }
}

impl Solution for Day12 {
    fn part1(&mut self, input: &mut Read) -> String {
        let mut state = self.read_input(input);
        print_state(&state);
        for _ in 1..=20 {
            state = self.simulate(&state);
        }

        let total: i32 = state.iter()
            .filter(|&&(_, x)| x)
            .map(|&(i, _)| i)
            .sum();

        format!("{}", total)
    }

    fn part2(&mut self, input: &mut Read) -> String {
        // Note: this is way too slow
        let mut state = self.read_input(input);
        print_state(&state);
        for _ in 1..=50000000000i64 {
            state = self.simulate(&state);
        }

        let total: i32 = state.iter()
            .filter(|&&(_, x)| x)
            .map(|&(i, _)| i)
            .sum();

        format!("{}", total)
    }
}

#[cfg(test)]
mod tests {
    use common::Solution;
    use day12::Day12;

    const SAMPLE_INPUT: &[u8] = include_bytes!("samples/12.txt");

    #[test]
    fn sample_part1() {
        let mut instance = Day12::new();
        assert_eq!("325", instance.part1(&mut SAMPLE_INPUT));
    }
}
<commit_msg>Implement day 12 part 2.<commit_after>use std::collections::HashMap;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Read;
use std::iter::FromIterator;

use common::Solution;

type State = Vec<(i64, bool)>;

#[derive(Default, Debug)]
pub struct Day12 {
    productions: [bool; 32],
}

fn char_bool(c: char) -> bool {
    match c {
        '#' => true,
        '.' => false,
        _ => panic!("Invalid input {}", c)
    }
}

fn print_state(state: &State) -> String {
    let mut buf = String::with_capacity(state.len());
    for &(_, b) in state {
        if b {
            buf.push('#');
        } else {
            buf.push('.');
        }
    }
    buf
}

fn state_from_string(representation: &str, offset: i64) -> State {
    Vec::from_iter(representation.chars()
        .enumerate()
        .map(|(i, c)| (i as i64 + offset, char_bool(c))))
}

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

    fn read_input(&mut self, input: &mut Read) -> State {
        let mut state = Vec::new();
        let mut reader = BufReader::new(input);
        {
            let mut line = String::new();
            reader.read_line(&mut line).unwrap();

            for (idx, c) in line.trim().chars().skip("initial state: ".len()).enumerate() {
                state.push((idx as i64, char_bool(c)));
            }
        }

        for line in reader.lines() {
            let line = line.unwrap();

            if line.is_empty() {
                continue;
            }

            let mut index = 0;
            for c in line.chars().take(5) {
                index <<= 1;
                if char_bool(c) {
                    index |= 1;
                }
            }

            self.productions[index] = char_bool(line.chars().last().unwrap());
        }

        state
    }

    fn simulate(&self, state: &State) -> State {
        let mut new_state = Vec::new();
        let mut index = 0;
        let mut last_idx = None;
        for &(idx, b) in state {
            index = (index << 1) & 0x1f;
            if b {
                index |= 1;
            }

            if self.productions[index] {
                new_state.push((idx - 2, true));
            } else if !new_state.is_empty() {
                new_state.push((idx - 2, false));
            }
            last_idx = Some(idx);
        }
        if let Some(idx) = last_idx {
            for i in 1..=4 {
                index = (index << 1) & 0x1f;

                if self.productions[index] {
                    new_state.push((idx + i - 2, true));
                } else if !new_state.is_empty() {
                    new_state.push((idx + i - 2, false));
                }
            }
        }

        let new_len = new_state.len() - new_state.iter().rev()
            .take_while(|&&(_, x)| !x).count();
        new_state.truncate(new_len);

        new_state
    }

    fn simulate_n(&self, mut state: State, n: usize) -> State {
        for _ in 0..n {
            state = self.simulate(&state);
        }
        state
    }

    fn sum(&self, state: &State) -> i64 {
        state.iter()
            .filter(|&&(_, x)| x)
            .map(|&(i, _)| i)
            .sum()
    }
}

impl Solution for Day12 {
    fn part1(&mut self, input: &mut Read) -> String {
        let mut state = self.read_input(input);
        state = self.simulate_n(state, 20);

        format!("{}", self.sum(&state))
    }

    fn part2(&mut self, input: &mut Read) -> String {
        // Note: this is way too slow
        let mut state = self.read_input(input);
        let mut seen = HashMap::new();
        let mut time = 1i64;
        const TARGET_TIME: i64 = 50000000000;

        loop {
            state = self.simulate(&state);
            let &(offset, _) = state.first().unwrap();
            let representation = print_state(&state);
            if let Some((o_offset, o_time)) = seen.get(&representation) {
                let remaining_time = TARGET_TIME - o_time;
                let cycle_length = time - o_time;
                let iterations_skipped = remaining_time / cycle_length;
                let remainder = remaining_time % cycle_length;
                let new_offset = o_offset + iterations_skipped * (offset - o_offset);

                state = state_from_string(&representation, new_offset);
                state = self.simulate_n(state, remainder as usize);

                return format!("{}", self.sum(&state));
            }
            seen.insert(representation, (offset, time));
            time += 1;
        }
    }
}

#[cfg(test)]
mod tests {
    use common::Solution;
    use day12::Day12;

    const SAMPLE_INPUT: &[u8] = include_bytes!("samples/12.txt");

    #[test]
    fn sample_part1() {
        let mut instance = Day12::new();
        assert_eq!("325", instance.part1(&mut SAMPLE_INPUT));
    }
}
<|endoftext|>


<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/. */

//! A task that takes a URL and streams back the binary data.

use about_loader;
use data_loader;
use file_loader;
use http_loader;
use cookie_storage::CookieStorage;
use cookie;
use mime_classifier::MIMEClassifier;

use net_traits::{ControlMsg, LoadData, LoadResponse, ResponseSenders, LoadConsumer};
use net_traits::{Metadata, ProgressMsg, ResourceTask, AsyncResponseTarget, ResponseAction};
use net_traits::ProgressMsg::Done;
use util::opts;
use util::task::spawn_named;

use hyper::header::{ContentType, Header, SetCookie, UserAgent};
use hyper::mime::{Mime, TopLevel, SubLevel};

use std::borrow::ToOwned;
use std::boxed;
use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::{BufReader, Read};
use std::str::FromStr;
use std::sync::Arc;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::thunk::Invoke;

static mut HOST_TABLE: Option<*mut HashMap<String, String>> = None;

pub fn global_init() {
<<<<<<< HEAD
    //TODO: handle bad file path
    let path = match env::var("HOST_FILE") {
        Ok(host_file_path) => host_file_path,
        Err(_) => return,
    };

    let mut file = match File::open(&path) {
        Ok(f) => BufReader::new(f),
        Err(_) => return,
    };

    let mut lines = String::new();
    match file.read_to_string(&mut lines) {
        Ok(()) => (),
        Err(_) => return,
    };

    unsafe {
        let host_table = boxed::into_raw(parse_hostsfile(&lines));
        HOST_TABLE = Some(host_table);
    }
}

pub enum ProgressSender {
    Channel(Sender<ProgressMsg>),
    Listener(Box<AsyncResponseTarget>),
}

impl ProgressSender {
    //XXXjdm return actual error
    pub fn send(&self, msg: ProgressMsg) -> Result<(), ()> {
        match *self {
            ProgressSender::Channel(ref c) => c.send(msg).map_err(|_| ()),
            ProgressSender::Listener(ref b) => {
                let action = match msg {
                    ProgressMsg::Payload(buf) => ResponseAction::DataAvailable(buf),
                    ProgressMsg::Done(status) => ResponseAction::ResponseComplete(status),
                };
                b.invoke_with_listener(action);
                Ok(())
            }
        }
    }
}

/// For use by loaders in responding to a Load message.
pub fn start_sending(start_chan: ResponseSenders, metadata: Metadata) -> ProgressSender {
    start_sending_opt(start_chan, metadata).ok().unwrap()
}

/// For use by loaders in responding to a Load message that allows content sniffing.
pub fn start_sending_sniffed(start_chan: ResponseSenders, metadata: Metadata,
                             classifier: Arc<MIMEClassifier>, partial_body: &Vec<u8>)
                             -> ProgressSender {
    start_sending_sniffed_opt(start_chan, metadata, classifier, partial_body).ok().unwrap()
}

/// For use by loaders in responding to a Load message that allows content sniffing.
pub fn start_sending_sniffed_opt(start_chan: ResponseSenders, mut metadata: Metadata,
                                 classifier: Arc<MIMEClassifier>, partial_body: &Vec<u8>)
                                 -> Result<ProgressSender, ()> {
    if opts::get().sniff_mime_types {
        // TODO: should be calculated in the resource loader, from pull requeset #4094
        let nosniff = false;
        let check_for_apache_bug = false;

        let supplied_type = metadata.content_type.map(|ContentType(Mime(toplevel, sublevel, _))| {
            (format!("{}", toplevel), format!("{}", sublevel))
        });
        metadata.content_type = classifier.classify(nosniff, check_for_apache_bug, &supplied_type, &partial_body).map(|(toplevel, sublevel)| {
            let mime_tp: TopLevel = FromStr::from_str(&toplevel).unwrap();
            let mime_sb: SubLevel = FromStr::from_str(&sublevel).unwrap();
            ContentType(Mime(mime_tp, mime_sb, vec!()))
        });

    }

    start_sending_opt(start_chan, metadata)
}

/// For use by loaders in responding to a Load message.
pub fn start_sending_opt(start_chan: ResponseSenders, metadata: Metadata) -> Result<ProgressSender, ()> {
    match start_chan {
        ResponseSenders::Channel(start_chan) => {
            let (progress_chan, progress_port) = channel();
            let result = start_chan.send(LoadResponse {
                metadata:      metadata,
                progress_port: progress_port,
            });
            match result {
                Ok(_) => Ok(ProgressSender::Channel(progress_chan)),
                Err(_) => Err(())
            }
        }
        ResponseSenders::Listener(target) => {
            target.invoke_with_listener(ResponseAction::HeadersAvailable(metadata));
            Ok(ProgressSender::Listener(target))
        }
    }
}

/// Create a ResourceTask
pub fn new_resource_task(user_agent: Option<String>) -> ResourceTask {
    let (setup_chan, setup_port) = channel();
    let setup_chan_clone = setup_chan.clone();
    spawn_named("ResourceManager".to_owned(), move || {
        ResourceManager::new(setup_port, user_agent, setup_chan_clone).start();
    });
    setup_chan
}

pub fn parse_hostsfile(hostsfile_content: &str) -> Box<HashMap<String, String>> {
    let ipv4_regex = regex!(r"^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$");
    let ipv6_regex = regex!(r"^([a-fA-F0-9]{0,4}[:]?){1,8}(/\d{1,3})?$");
    let mut host_table = HashMap::new();
    let lines: Vec<&str> = hostsfile_content.split('\n').collect();

    for line in lines.iter() {
        let ip_host: Vec<&str> = line.trim().split(|c: char| c == ' ' || c == '\t').collect();
        if ip_host.len() > 1 {
            if !ipv4_regex.is_match(ip_host[0]) && !ipv6_regex.is_match(ip_host[0]) { continue; }
            let address = ip_host[0].to_owned();

            for token in ip_host.iter().skip(1) {
                if token.as_bytes()[0] == b'#' {
                    break;
                }
                host_table.insert(token.to_owned().to_string(), address.clone());
            }
        }
    }
    box host_table
}

pub fn replace_hosts(mut load_data: LoadData, host_table: *mut HashMap<String, String>) -> LoadData {
    if let Some(h) = load_data.url.domain_mut() {
        unsafe {
            if let Some(ip) = (*host_table).get(h) {
                *h = ip.clone();
            }
        }
    }
    return load_data;
}

struct ResourceManager {
    from_client: Receiver<ControlMsg>,
    user_agent: Option<String>,
    cookie_storage: CookieStorage,
    resource_task: Sender<ControlMsg>,
    mime_classifier: Arc<MIMEClassifier>,
}

impl ResourceManager {
    fn new(from_client: Receiver<ControlMsg>, user_agent: Option<String>,
           resource_task: Sender<ControlMsg>) -> ResourceManager {
        ResourceManager {
            from_client: from_client,
            user_agent: user_agent,
            cookie_storage: CookieStorage::new(),
            resource_task: resource_task,
            mime_classifier: Arc::new(MIMEClassifier::new()),
        }
    }
}

impl ResourceManager {
    fn start(&mut self) {
        loop {
            match self.from_client.recv().unwrap() {
              ControlMsg::Load(load_data, consumer) => {
                self.load(load_data, consumer)
              }
              ControlMsg::SetCookiesForUrl(request, cookie_list, source) => {
                let header = Header::parse_header(&[cookie_list.into_bytes()]);
                if let Some(SetCookie(cookies)) = header {
                  for bare_cookie in cookies.into_iter() {
                    if let Some(cookie) = cookie::Cookie::new_wrapped(bare_cookie, &request, source) {
                      self.cookie_storage.push(cookie, source);
                    }
                  }
                }
              }
              ControlMsg::GetCookiesForUrl(url, consumer, source) => {
                consumer.send(self.cookie_storage.cookies_for_url(&url, source)).unwrap();
              }
              ControlMsg::Exit => {
                break
              }
            }
        }
    }

    fn load(&mut self, mut load_data: LoadData, consumer: LoadConsumer) {
        unsafe {
            if let Some(host_table) = HOST_TABLE {
                load_data = replace_hosts(load_data, host_table);
            }
        }

        self.user_agent.as_ref().map(|ua| load_data.headers.set(UserAgent(ua.clone())));

        fn from_factory(factory: fn(LoadData, ResponseSenders, Arc<MIMEClassifier>))
                        -> Box<Invoke<(LoadData, ResponseSenders, Arc<MIMEClassifier>)> + Send> {
            box move |(load_data, senders, classifier)| {
                factory(load_data, senders, classifier)
            }
        }

        let senders = ResponseSenders::from_consumer(consumer);
        let loader = match &*load_data.url.scheme {
            "file" => from_factory(file_loader::factory),
            "http" | "https" | "view-source" => http_loader::factory(self.resource_task.clone()),
            "data" => from_factory(data_loader::factory),
            "about" => from_factory(about_loader::factory),
            _ => {
                debug!("resource_task: no loader for scheme {}", load_data.url.scheme);
                start_sending(senders, Metadata::default(load_data.url))
                    .send(ProgressMsg::Done(Err("no loader for scheme".to_string()))).unwrap();
                return
            }
        };
        debug!("resource_task: loading url: {}", load_data.url.serialize());

        loader.invoke((load_data, senders, self.mime_classifier.clone()));
    }
}
<commit_msg>Warning fixes.<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/. */

//! A task that takes a URL and streams back the binary data.

use about_loader;
use data_loader;
use file_loader;
use http_loader;
use cookie_storage::CookieStorage;
use cookie;
use mime_classifier::MIMEClassifier;

use net_traits::{ControlMsg, LoadData, LoadResponse, ResponseSenders, LoadConsumer};
use net_traits::{Metadata, ProgressMsg, ResourceTask, AsyncResponseTarget, ResponseAction};
use net_traits::ProgressMsg::Done;
use util::opts;
use util::task::spawn_named;

use hyper::header::{ContentType, Header, SetCookie, UserAgent};
use hyper::mime::{Mime, TopLevel, SubLevel};

use std::borrow::ToOwned;
use std::boxed;
use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::{BufReader, Read};
use std::str::FromStr;
use std::sync::Arc;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::thunk::Invoke;

static mut HOST_TABLE: Option<*mut HashMap<String, String>> = None;

pub fn global_init() {
    //TODO: handle bad file path
    let path = match env::var("HOST_FILE") {
        Ok(host_file_path) => host_file_path,
        Err(_) => return,
    };

    let mut file = match File::open(&path) {
        Ok(f) => BufReader::new(f),
        Err(_) => return,
    };

    let mut lines = String::new();
    match file.read_to_string(&mut lines) {
        Ok(()) => (),
        Err(_) => return,
    };

    unsafe {
        let host_table = boxed::into_raw(parse_hostsfile(&lines));
        HOST_TABLE = Some(host_table);
    }
}

pub enum ProgressSender {
    Channel(Sender<ProgressMsg>),
    Listener(Box<AsyncResponseTarget>),
}

impl ProgressSender {
    //XXXjdm return actual error
    pub fn send(&self, msg: ProgressMsg) -> Result<(), ()> {
        match *self {
            ProgressSender::Channel(ref c) => c.send(msg).map_err(|_| ()),
            ProgressSender::Listener(ref b) => {
                let action = match msg {
                    ProgressMsg::Payload(buf) => ResponseAction::DataAvailable(buf),
                    ProgressMsg::Done(status) => ResponseAction::ResponseComplete(status),
                };
                b.invoke_with_listener(action);
                Ok(())
            }
        }
    }
}

/// For use by loaders in responding to a Load message.
pub fn start_sending(start_chan: ResponseSenders, metadata: Metadata) -> ProgressSender {
    start_sending_opt(start_chan, metadata).ok().unwrap()
}

/// For use by loaders in responding to a Load message that allows content sniffing.
pub fn start_sending_sniffed(start_chan: ResponseSenders, metadata: Metadata,
                             classifier: Arc<MIMEClassifier>, partial_body: &Vec<u8>)
                             -> ProgressSender {
    start_sending_sniffed_opt(start_chan, metadata, classifier, partial_body).ok().unwrap()
}

/// For use by loaders in responding to a Load message that allows content sniffing.
pub fn start_sending_sniffed_opt(start_chan: ResponseSenders, mut metadata: Metadata,
                                 classifier: Arc<MIMEClassifier>, partial_body: &Vec<u8>)
                                 -> Result<ProgressSender, ()> {
    if opts::get().sniff_mime_types {
        // TODO: should be calculated in the resource loader, from pull requeset #4094
        let nosniff = false;
        let check_for_apache_bug = false;

        let supplied_type = metadata.content_type.map(|ContentType(Mime(toplevel, sublevel, _))| {
            (format!("{}", toplevel), format!("{}", sublevel))
        });
        metadata.content_type = classifier.classify(nosniff, check_for_apache_bug, &supplied_type, &partial_body).map(|(toplevel, sublevel)| {
            let mime_tp: TopLevel = FromStr::from_str(&toplevel).unwrap();
            let mime_sb: SubLevel = FromStr::from_str(&sublevel).unwrap();
            ContentType(Mime(mime_tp, mime_sb, vec!()))
        });

    }

    start_sending_opt(start_chan, metadata)
}

/// For use by loaders in responding to a Load message.
pub fn start_sending_opt(start_chan: ResponseSenders, metadata: Metadata) -> Result<ProgressSender, ()> {
    match start_chan {
        ResponseSenders::Channel(start_chan) => {
            let (progress_chan, progress_port) = channel();
            let result = start_chan.send(LoadResponse {
                metadata:      metadata,
                progress_port: progress_port,
            });
            match result {
                Ok(_) => Ok(ProgressSender::Channel(progress_chan)),
                Err(_) => Err(())
            }
        }
        ResponseSenders::Listener(target) => {
            target.invoke_with_listener(ResponseAction::HeadersAvailable(metadata));
            Ok(ProgressSender::Listener(target))
        }
    }
}

/// Create a ResourceTask
pub fn new_resource_task(user_agent: Option<String>) -> ResourceTask {
    let (setup_chan, setup_port) = channel();
    let setup_chan_clone = setup_chan.clone();
    spawn_named("ResourceManager".to_owned(), move || {
        ResourceManager::new(setup_port, user_agent, setup_chan_clone).start();
    });
    setup_chan
}

pub fn parse_hostsfile(hostsfile_content: &str) -> Box<HashMap<String, String>> {
    let ipv4_regex = regex!(r"^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$");
    let ipv6_regex = regex!(r"^([a-fA-F0-9]{0,4}[:]?){1,8}(/\d{1,3})?$");
    let mut host_table = HashMap::new();
    let lines: Vec<&str> = hostsfile_content.split('\n').collect();

    for line in lines.iter() {
        let ip_host: Vec<&str> = line.trim().split(|c: char| c == ' ' || c == '\t').collect();
        if ip_host.len() > 1 {
            if !ipv4_regex.is_match(ip_host[0]) && !ipv6_regex.is_match(ip_host[0]) { continue; }
            let address = ip_host[0].to_owned();

            for token in ip_host.iter().skip(1) {
                if token.as_bytes()[0] == b'#' {
                    break;
                }
                host_table.insert(token.to_owned().to_string(), address.clone());
            }
        }
    }
    box host_table
}

pub fn replace_hosts(mut load_data: LoadData, host_table: *mut HashMap<String, String>) -> LoadData {
    if let Some(h) = load_data.url.domain_mut() {
        unsafe {
            if let Some(ip) = (*host_table).get(h) {
                *h = ip.clone();
            }
        }
    }
    return load_data;
}

struct ResourceManager {
    from_client: Receiver<ControlMsg>,
    user_agent: Option<String>,
    cookie_storage: CookieStorage,
    resource_task: Sender<ControlMsg>,
    mime_classifier: Arc<MIMEClassifier>,
}

impl ResourceManager {
    fn new(from_client: Receiver<ControlMsg>, user_agent: Option<String>,
           resource_task: Sender<ControlMsg>) -> ResourceManager {
        ResourceManager {
            from_client: from_client,
            user_agent: user_agent,
            cookie_storage: CookieStorage::new(),
            resource_task: resource_task,
            mime_classifier: Arc::new(MIMEClassifier::new()),
        }
    }
}

impl ResourceManager {
    fn start(&mut self) {
        loop {
            match self.from_client.recv().unwrap() {
              ControlMsg::Load(load_data, consumer) => {
                self.load(load_data, consumer)
              }
              ControlMsg::SetCookiesForUrl(request, cookie_list, source) => {
                let header = Header::parse_header(&[cookie_list.into_bytes()]);
                if let Some(SetCookie(cookies)) = header {
                  for bare_cookie in cookies.into_iter() {
                    if let Some(cookie) = cookie::Cookie::new_wrapped(bare_cookie, &request, source) {
                      self.cookie_storage.push(cookie, source);
                    }
                  }
                }
              }
              ControlMsg::GetCookiesForUrl(url, consumer, source) => {
                consumer.send(self.cookie_storage.cookies_for_url(&url, source)).unwrap();
              }
              ControlMsg::Exit => {
                break
              }
            }
        }
    }

    fn load(&mut self, mut load_data: LoadData, consumer: LoadConsumer) {
        unsafe {
            if let Some(host_table) = HOST_TABLE {
                load_data = replace_hosts(load_data, host_table);
            }
        }

        self.user_agent.as_ref().map(|ua| load_data.headers.set(UserAgent(ua.clone())));

        fn from_factory(factory: fn(LoadData, ResponseSenders, Arc<MIMEClassifier>))
                        -> Box<Invoke<(LoadData, ResponseSenders, Arc<MIMEClassifier>)> + Send> {
            box move |(load_data, senders, classifier)| {
                factory(load_data, senders, classifier)
            }
        }

        let senders = ResponseSenders::from_consumer(consumer);
        let loader = match &*load_data.url.scheme {
            "file" => from_factory(file_loader::factory),
            "http" | "https" | "view-source" => http_loader::factory(self.resource_task.clone()),
            "data" => from_factory(data_loader::factory),
            "about" => from_factory(about_loader::factory),
            _ => {
                debug!("resource_task: no loader for scheme {}", load_data.url.scheme);
                start_sending(senders, Metadata::default(load_data.url))
                    .send(ProgressMsg::Done(Err("no loader for scheme".to_string()))).unwrap();
                return
            }
        };
        debug!("resource_task: loading url: {}", load_data.url.serialize());

        loader.invoke((load_data, senders, self.mime_classifier.clone()));
    }
}
<|endoftext|>


<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/. */

use attr::CaseSensitivity;
use bloom::BloomFilter;
use nth_index_cache::NthIndexCache;

/// What kind of selector matching mode we should use.
///
/// There are two modes of selector matching. The difference is only noticeable
/// in presence of pseudo-elements.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum MatchingMode {
    /// Don't ignore any pseudo-element selectors.
    Normal,

    /// Ignores any stateless pseudo-element selectors in the rightmost sequence
    /// of simple selectors.
    ///
    /// This is useful, for example, to match against ::before when you aren't a
    /// pseudo-element yourself.
    ///
    /// For example, in presence of `::before:hover`, it would never match, but
    /// `::before` would be ignored as in "matching".
    ///
    /// It's required for all the selectors you match using this mode to have a
    /// pseudo-element.
    ForStatelessPseudoElement,
}

/// The mode to use when matching unvisited and visited links.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum VisitedHandlingMode {
    /// All links are matched as if they are unvisted.
    AllLinksUnvisited,
    /// All links are matched as if they are visited and unvisited (both :link
    /// and :visited match).
    ///
    /// This is intended to be used from invalidation code, to be conservative
    /// about whether we need to restyle a link.
    AllLinksVisitedAndUnvisited,
    /// A element's "relevant link" is the element being matched if it is a link
    /// or the nearest ancestor link. The relevant link is matched as though it
    /// is visited, and all other links are matched as if they are unvisited.
    RelevantLinkVisited,
}

/// Which quirks mode is this document in.
///
/// See: https://quirks.spec.whatwg.org/
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum QuirksMode {
    /// Quirks mode.
    Quirks,
    /// Limited quirks mode.
    LimitedQuirks,
    /// No quirks mode.
    NoQuirks,
}

impl QuirksMode {
    #[inline]
    pub fn classes_and_ids_case_sensitivity(self) -> CaseSensitivity {
        match self {
            QuirksMode::NoQuirks |
            QuirksMode::LimitedQuirks => CaseSensitivity::CaseSensitive,
            QuirksMode::Quirks => CaseSensitivity::AsciiCaseInsensitive,
        }
    }
}

/// Data associated with the matching process for a element.  This context is
/// used across many selectors for an element, so it's not appropriate for
/// transient data that applies to only a single selector.
pub struct MatchingContext<'a> {
    /// Input with the matching mode we should use when matching selectors.
    pub matching_mode: MatchingMode,
    /// Input with the bloom filter used to fast-reject selectors.
    pub bloom_filter: Option<&'a BloomFilter>,
    /// An optional cache to speed up nth-index-like selectors.
    pub nth_index_cache: Option<&'a mut NthIndexCache>,
    /// Input that controls how matching for links is handled.
    pub visited_handling: VisitedHandlingMode,
    /// Output that records whether we encountered a "relevant link" while
    /// matching _any_ selector for this element. (This differs from
    /// `RelevantLinkStatus` which tracks the status for the _current_ selector
    /// only.)
    pub relevant_link_found: bool,

    quirks_mode: QuirksMode,
    classes_and_ids_case_sensitivity: CaseSensitivity,
}

impl<'a> MatchingContext<'a> {
    /// Constructs a new `MatchingContext`.
    pub fn new(matching_mode: MatchingMode,
               bloom_filter: Option<&'a BloomFilter>,
               nth_index_cache: Option<&'a mut NthIndexCache>,
               quirks_mode: QuirksMode)
               -> Self
    {
        Self {
            matching_mode: matching_mode,
            bloom_filter: bloom_filter,
            nth_index_cache: nth_index_cache,
            visited_handling: VisitedHandlingMode::AllLinksUnvisited,
            relevant_link_found: false,
            quirks_mode: quirks_mode,
            classes_and_ids_case_sensitivity: quirks_mode.classes_and_ids_case_sensitivity(),
        }
    }

    /// Constructs a new `MatchingContext` for use in visited matching.
    pub fn new_for_visited(matching_mode: MatchingMode,
                           bloom_filter: Option<&'a BloomFilter>,
                           nth_index_cache: Option<&'a mut NthIndexCache>,
                           visited_handling: VisitedHandlingMode,
                           quirks_mode: QuirksMode)
                           -> Self
    {
        Self {
            matching_mode: matching_mode,
            bloom_filter: bloom_filter,
            visited_handling: visited_handling,
            relevant_link_found: false,
            nth_index_cache: nth_index_cache,
            quirks_mode: quirks_mode,
            classes_and_ids_case_sensitivity: quirks_mode.classes_and_ids_case_sensitivity(),
        }
    }

    /// The quirks mode of the document.
    #[inline]
    pub fn quirks_mode(&self) -> QuirksMode {
        self.quirks_mode
    }

    /// The case-sensitivity for class and ID selectors
    #[inline]
    pub fn classes_and_ids_case_sensitivity(&self) -> CaseSensitivity {
        self.classes_and_ids_case_sensitivity
    }
}
<commit_msg>Auto merge of #18773 - emilio:mc-cleanup, r=nox<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/. */

use attr::CaseSensitivity;
use bloom::BloomFilter;
use nth_index_cache::NthIndexCache;

/// What kind of selector matching mode we should use.
///
/// There are two modes of selector matching. The difference is only noticeable
/// in presence of pseudo-elements.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum MatchingMode {
    /// Don't ignore any pseudo-element selectors.
    Normal,

    /// Ignores any stateless pseudo-element selectors in the rightmost sequence
    /// of simple selectors.
    ///
    /// This is useful, for example, to match against ::before when you aren't a
    /// pseudo-element yourself.
    ///
    /// For example, in presence of `::before:hover`, it would never match, but
    /// `::before` would be ignored as in "matching".
    ///
    /// It's required for all the selectors you match using this mode to have a
    /// pseudo-element.
    ForStatelessPseudoElement,
}

/// The mode to use when matching unvisited and visited links.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum VisitedHandlingMode {
    /// All links are matched as if they are unvisted.
    AllLinksUnvisited,
    /// All links are matched as if they are visited and unvisited (both :link
    /// and :visited match).
    ///
    /// This is intended to be used from invalidation code, to be conservative
    /// about whether we need to restyle a link.
    AllLinksVisitedAndUnvisited,
    /// A element's "relevant link" is the element being matched if it is a link
    /// or the nearest ancestor link. The relevant link is matched as though it
    /// is visited, and all other links are matched as if they are unvisited.
    RelevantLinkVisited,
}

/// Which quirks mode is this document in.
///
/// See: https://quirks.spec.whatwg.org/
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum QuirksMode {
    /// Quirks mode.
    Quirks,
    /// Limited quirks mode.
    LimitedQuirks,
    /// No quirks mode.
    NoQuirks,
}

impl QuirksMode {
    #[inline]
    pub fn classes_and_ids_case_sensitivity(self) -> CaseSensitivity {
        match self {
            QuirksMode::NoQuirks |
            QuirksMode::LimitedQuirks => CaseSensitivity::CaseSensitive,
            QuirksMode::Quirks => CaseSensitivity::AsciiCaseInsensitive,
        }
    }
}

/// Data associated with the matching process for a element.  This context is
/// used across many selectors for an element, so it's not appropriate for
/// transient data that applies to only a single selector.
pub struct MatchingContext<'a> {
    /// Input with the matching mode we should use when matching selectors.
    pub matching_mode: MatchingMode,
    /// Input with the bloom filter used to fast-reject selectors.
    pub bloom_filter: Option<&'a BloomFilter>,
    /// An optional cache to speed up nth-index-like selectors.
    pub nth_index_cache: Option<&'a mut NthIndexCache>,
    /// Input that controls how matching for links is handled.
    pub visited_handling: VisitedHandlingMode,
    /// Output that records whether we encountered a "relevant link" while
    /// matching _any_ selector for this element. (This differs from
    /// `RelevantLinkStatus` which tracks the status for the _current_ selector
    /// only.)
    pub relevant_link_found: bool,

    quirks_mode: QuirksMode,
    classes_and_ids_case_sensitivity: CaseSensitivity,
}

impl<'a> MatchingContext<'a> {
    /// Constructs a new `MatchingContext`.
    pub fn new(
        matching_mode: MatchingMode,
        bloom_filter: Option<&'a BloomFilter>,
        nth_index_cache: Option<&'a mut NthIndexCache>,
        quirks_mode: QuirksMode,
    ) -> Self {
        Self::new_for_visited(
            matching_mode,
            bloom_filter,
            nth_index_cache,
            VisitedHandlingMode::AllLinksUnvisited,
            quirks_mode
        )
    }

    /// Constructs a new `MatchingContext` for use in visited matching.
    pub fn new_for_visited(
        matching_mode: MatchingMode,
        bloom_filter: Option<&'a BloomFilter>,
        nth_index_cache: Option<&'a mut NthIndexCache>,
        visited_handling: VisitedHandlingMode,
        quirks_mode: QuirksMode,
    ) -> Self {
        Self {
            matching_mode,
            bloom_filter,
            visited_handling,
            nth_index_cache,
            quirks_mode,
            relevant_link_found: false,
            classes_and_ids_case_sensitivity: quirks_mode.classes_and_ids_case_sensitivity(),
        }
    }

    /// The quirks mode of the document.
    #[inline]
    pub fn quirks_mode(&self) -> QuirksMode {
        self.quirks_mode
    }

    /// The case-sensitivity for class and ID selectors
    #[inline]
    pub fn classes_and_ids_case_sensitivity(&self) -> CaseSensitivity {
        self.classes_and_ids_case_sensitivity
    }
}
<|endoftext|>