client.rs

#![feature(phase)]

#[phase(syntax, link)] extern crate log;
extern crate collections;
extern crate crypto = "rust-crypto";

use collections::Bitv;
use common::{KiB, BlockId, SequenceId, Packet};
use std::slice::MutableCloneableVector;

mod common;

// 12KiB should be large enough to support Jumbo ethernet frames if desired
static RECEIVE_BUFFER_SIZE: uint = 12*KiB;

struct BlockDownload {
    data: Vec<u8>,
    /// Sequence id of the next packet that is expected.
    next_packet: SequenceId,
    /// SHA-1 hash of the block. Present after the last packet is received.
    hash: Option<[u8, ..32]>,
    /// List of pending packets that were skipped over and not received yet.
    missed_packets: Vec<SequenceId>,
}

impl BlockDownload {
    /// Creates a new block of the specified size.
    pub fn new(block_size: uint) -> BlockDownload {
        BlockDownload {
            data: Vec::from_elem(block_size, 0u8),
            next_packet: 0,
            hash: None,
            missed_packets: Vec::new(),
        }
    }

    /// Processes the given packet, merging it into the block if necessary.
    pub fn add_packet(&mut self, packet: &Packet) {
        if self.update_sequence(packet.sequence_id) {
            if packet.offset != !0 {
                self.data.mut_slice_from(packet.offset).copy_from(packet.data);
            } else {
                // Terminating packet. `data` contains SHA-1 hash of block contents.
                if packet.data.len() != 32 {
                    warn!("Malformed terminating packet with length {}", packet.data.len());
                    return;
                }
                self.hash = Some([0u8, ..32]);
                self.hash.unwrap().copy_from(packet.data);
            }
        }
    }

    /// Queries if the download of the block has completed successfully.
    pub fn is_done(&self) -> bool {
        self.hash.is_some() && self.missed_packets.is_empty()
    }

    /// Verifies if the block's content's hash matches with the one supplied by the server.
    pub fn is_valid(&self) -> bool {
        use crypto::sha2::Sha256;
        use crypto::digest::Digest;

        if !self.is_done() {
            return false;
        }

        let expected_hash = self.hash.unwrap();
        let mut actual_hash = [0, ..32];

        let mut hasher = Sha256::new();
        hasher.input(self.data.as_slice());
        hasher.result(actual_hash);

        expected_hash == actual_hash
    }

    /// Returns the block data or None if the download has not yet finished.
    pub fn get_data(self) -> Option<Vec<u8>> {
        if self.is_done() {
            Some(self.data)
        } else {
            None
        }
    }

    /// Updates the sequence number fields and returns if the packet with the specified id should be
    /// processed or not.
    fn update_sequence(&mut self, new_id: SequenceId) -> bool {
        if new_id == self.next_packet {
            // In-order packet. Increment and move on.
            self.next_packet += 1;
            return true;
        } else if new_id < self.next_packet {
            // Late packet. check if we haven't processed it yet.
            // Search if packet was in the pending packets list
            match self.missed_packets.iter().position(|x| *x == new_id) {
                Some(i) => {
                    self.missed_packets.swap_remove(i);
                    return true;
                },
                None => return false // Packet had already been received, ignore it
            }
        } else {
            // Future packet. We missed some, add them to the pending list.
            self.missed_packets.extend(range(self.next_packet, new_id));
            self.next_packet = new_id + 1;
            return true;
        }
    }
}

struct FileInfo {
    path: ~str,
    size: u64,
    first_block: BlockId,
}

impl FileInfo {
    fn num_blocks(&self, block_size: uint) -> BlockId {
        // Round up
        ((self.size + block_size as u64 - 1) / block_size as u64) as BlockId
    }

    fn end_block(&self, block_size: uint) -> BlockId {
        self.first_block + self.num_blocks(block_size)
    }
}

struct FileSet {
    block_size: uint,
    files: Vec<FileInfo>,

    /// Blocks which this client has already downloaded.
    downloaded_blocks: Bitv,
}

impl FileSet {
    fn new(mut files: Vec<FileInfo>, block_size: uint) -> FileSet {
        files.sort_by(|a, b| a.first_block.cmp(&b.first_block));
        let num_blocks = match files.last() {
            Some(file) => file.end_block(block_size),
            None => 0
        } as uint;

        FileSet {
            block_size: block_size,
            files: files,
            downloaded_blocks: Bitv::new(num_blocks, false),
        }
    }

    /// Returns the file that contains the specified block.
    fn file_for_block<'a>(&'a self, block: BlockId) -> Option<&'a FileInfo> {
        self.files.as_slice().bsearch(|file| {
            if file.first_block > block {
                Greater
            } else if file.end_block(self.block_size) <= block {
                Less
            } else {
                Equal
            }
        }).map(|i| self.files.get(i))
    }
}

fn block_writer_task(rx: Receiver<BlockDownload>) {
}

fn main() {
}

common.rs

use std::slice::MutableCloneableVector;

pub static KiB: uint = 1024;
pub static MiB: uint = 1024*KiB;

pub type BlockId = u16;
pub type SequenceId = u16;

pub struct Packet<'a> {
    pub block: BlockId,
    pub sequence_id: SequenceId,
    pub offset: uint,
    pub data: &'a [u8],
}

fn read_u16(buf: &[u8]) -> u16 {
    buf[0] as u16 << 8 |
    buf[1] as u16
}

fn read_u32(buf: &[u8]) -> u32 {
    buf[0] as u32 << 24 |
    buf[1] as u32 << 16 |
    buf[2] as u32 <<  8 |
    buf[3] as u32
}

fn write_u16(val: u16, buf: &mut [u8]) {
    buf[0] = (val >> 8) as u8;
    buf[1] = (val >> 0) as u8;
}

fn write_u32(val: u32, buf: &mut [u8]) {
    buf[0] = (val >> 24) as u8;
    buf[1] = (val >> 16) as u8;
    buf[2] = (val >>  8) as u8;
    buf[3] = (val >>  0) as u8;
}

impl<'a> Packet<'a> {
    fn from_buffer(buffer: &'a [u8]) -> Packet<'a> {
        Packet {
            block: read_u16(buffer.slice(0, 2)),
            sequence_id: read_u16(buffer.slice(2, 4)),
            offset: read_u32(buffer.slice(4, 8)) as uint,
            data: buffer.slice_from(8),
        }
    }

    fn to_buffer(&self, buffer: &mut[u8]) -> uint {
        write_u16(self.block, buffer.mut_slice(0, 2));
        write_u16(self.sequence_id, buffer.mut_slice(2, 4));
        write_u32(self.offset as u32, buffer.mut_slice(4, 8));
        buffer.mut_slice_from(8).copy_from(self.data);

        2 + 2 + 4 + self.data.len()
    }
}