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

// http://smallcultfollowing.com/babysteps/blog/2015/04/06/modeling-graphs-in-rust-using-vector-indices/
pub struct Graph {
    pub nodes: Vec<node::Node>,
    pub edges: Vec<edge::Edge>,
}

impl Graph {
    pub fn add_node(&mut self, identifier: String, attributes: Option<HashMap<String,String>>)
        -> node::Index {
        let index = self.nodes.len();
        self.nodes.push(node::Node { identifier: identifier, first_outgoing_edge: None, attributes: attributes });
        index
    }

    pub fn add_edge(&mut self, source: node::Index, target: node::Index, identifier: String, attributes: Option<HashMap<String,String>>) {
        let edge_index = self.edges.len();
        let node_data = &mut self.nodes[source];
        self.edges.push(edge::Edge {
            identifier: identifier,
            target: target,
            next_outgoing_edge: node_data.first_outgoing_edge,
            attributes: attributes,
        });
        node_data.first_outgoing_edge = Some(edge_index);
    }

    pub fn successors(&self, source: node::Index) -> Successors {
        let first_outgoing_edge = self.nodes[source].first_outgoing_edge;
        Successors { graph: self, current_edge_index: first_outgoing_edge }
    }

    pub fn print(self) {
        for n in 0..self.nodes.len() {
            print!("node::Node {} goes to: ", n);
            let mut suc = self.successors(n);
            loop {
                match suc.next() {
                    Some(s) => { print!("{}, ", s) },
                    None => { print!("\n"); break },
                }
            }
        }
    }
}

pub struct Successors<'graph> {
    graph: &'graph Graph,
    current_edge_index: Option<edge::Index>,
}

impl<'graph> Iterator for Successors<'graph> {
    type Item = node::Index;

    fn next(&mut self) -> Option<node::Index> {
        match self.current_edge_index {
            None => None,
            Some(edge_num) => {
                let edge = &self.graph.edges[edge_num];
                self.current_edge_index = edge.next_outgoing_edge;
                Some(edge.target)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use super::*;
    #[test]
    fn create_small_graph() {
        let mut graph = Graph { nodes: vec![], edges: vec![] };
        let node0 = graph.add_node("node0".to_string(), None);
        let node1 = graph.add_node("node1".to_string(), None);
        graph.add_edge(node0, node1, "edge0".to_string(), None);

        assert_eq!(graph.nodes.len(), 2);
        assert_eq!(graph.edges.len(), 1);
        match graph.nodes[0].first_outgoing_edge {
            Some(index) => assert_eq!(graph.edges[index].target, 1),
            None => assert!(false),
        }
    }

    #[test]
    fn create_list_graph() {
        let mut graph = Graph { nodes: vec![], edges: vec![] };
        let node0 = graph.add_node("node0".to_string(), None);
        let node1 = graph.add_node("node1".to_string(), None);
        let node2 = graph.add_node("node2".to_string(), None);
        let node3 = graph.add_node("node3".to_string(), None);
        graph.add_edge(node0, node1, "edge0".to_string(), None);
        graph.add_edge(node1, node2, "edge1".to_string(), None);
        graph.add_edge(node2, node3, "edge2".to_string(), None);

        let mut targets = vec![];
        for e in graph.edges {
            targets.push(e.target);
        }
        assert_eq!(targets, [1,2,3]);
    }

    #[test]
    fn list_node_edges() {
        let mut graph = Graph { nodes: vec![], edges: vec![] };
        let node0 = graph.add_node("node0".to_string(), None);
        graph.add_edge(node0, node0, "edge0".to_string(), None);
        graph.add_edge(node0, node0, "edge1".to_string(), None);
        graph.add_edge(node0, node0, "edge2".to_string(), None);
        let mut targets = vec![];
        let mut successors = graph.successors(0);
        loop {
            match successors.next() {
                Some(edge) => targets.push(edge),
                None => break,
            }
        }
        assert_eq!(targets, [0,0,0]);
    }

    #[test]
    fn node_attributes() {
        let mut graph = Graph { nodes: vec![], edges: vec![] };

        let mut attributes = HashMap::new();
        attributes.insert("key".to_string(), "value".to_string());

        let node0 = graph.add_node("node0".to_string(), Some(attributes));
        match graph.nodes[node0].attributes {
            Some(ref attrs) => {
                match attrs.get("key") {
                    Some(value) => assert_eq!(value, &"value".to_string()),
                    None => assert!(false),
                }
            },
            None => assert!(false),
        }
    }

    #[test]
    fn edge_attributes() {
        let mut graph = Graph { nodes: vec![], edges: vec![] };

        let mut attributes = HashMap::new();
        attributes.insert("key".to_string(), "value".to_string());

        let node0 = graph.add_node("node0".to_string(), None);
        graph.add_edge(node0, node0, "edge0".to_string(), Some(attributes));
        match graph.edges[0].attributes {
            Some(ref attrs) => {
                match attrs.get("key") {
                    Some(value) => assert_eq!(value, &"value".to_string()),
                    None => assert!(false),
                }
            },
            None => assert!(false),
        }
    }
}
<commit_msg>Add edges for node function<commit_after>use std::collections::HashMap;
use node;
use edge;

// http://smallcultfollowing.com/babysteps/blog/2015/04/06/modeling-graphs-in-rust-using-vector-indices/
pub struct Graph {
    pub nodes: Vec<node::Node>,
    pub edges: Vec<edge::Edge>,
}

impl Graph {
    pub fn add_node(&mut self, identifier: String, attributes: Option<HashMap<String,String>>)
        -> node::Index {
        let index = self.nodes.len();
        self.nodes.push(node::Node { identifier: identifier, first_outgoing_edge: None, attributes: attributes });
        index
    }

    pub fn add_edge(&mut self, source: node::Index, target: node::Index, identifier: String, attributes: Option<HashMap<String,String>>) {
        let edge_index = self.edges.len();
        let node_data = &mut self.nodes[source];
        self.edges.push(edge::Edge {
            identifier: identifier,
            target: target,
            next_outgoing_edge: node_data.first_outgoing_edge,
            attributes: attributes,
        });
        node_data.first_outgoing_edge = Some(edge_index);
    }

    pub fn successors(&self, source: node::Index) -> Successors {
        let first_outgoing_edge = self.nodes[source].first_outgoing_edge;
        Successors { graph: self, current_edge_index: first_outgoing_edge }
    }

    pub fn edges_for_node(&self, node_index: node::Index) -> Vec<edge::Index> {
        let mut edge_indexes: Vec<edge::Index> = vec![];
        match self.nodes[node_index].first_outgoing_edge {
            Some(edge_index) => {
                let mut edge = &self.edges[edge_index];
                edge_indexes.push(edge_index);
                loop {
                    match edge.next_outgoing_edge {
                        Some(edge_index) => {
                            edge = &self.edges[edge_index];
                            edge_indexes.push(edge_index);
                        },
                        None => { break; }
                    }
                }
            },
            None => {}
        }
        return edge_indexes;
    }

    pub fn print(self) {
        for n in 0..self.nodes.len() {
            print!("node::Node {} goes to: ", n);
            let mut suc = self.successors(n);
            loop {
                match suc.next() {
                    Some(s) => { print!("{}, ", s) },
                    None => { print!("\n"); break },
                }
            }
        }
    }
}

pub struct Successors<'graph> {
    graph: &'graph Graph,
    current_edge_index: Option<edge::Index>,
}

impl<'graph> Iterator for Successors<'graph> {
    type Item = node::Index;

    fn next(&mut self) -> Option<node::Index> {
        match self.current_edge_index {
            None => None,
            Some(edge_num) => {
                let edge = &self.graph.edges[edge_num];
                self.current_edge_index = edge.next_outgoing_edge;
                Some(edge.target)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use super::*;
    #[test]
    fn create_small_graph() {
        let mut graph = Graph { nodes: vec![], edges: vec![] };
        let node0 = graph.add_node("node0".to_string(), None);
        let node1 = graph.add_node("node1".to_string(), None);
        graph.add_edge(node0, node1, "edge0".to_string(), None);

        assert_eq!(graph.nodes.len(), 2);
        assert_eq!(graph.edges.len(), 1);
        match graph.nodes[0].first_outgoing_edge {
            Some(index) => assert_eq!(graph.edges[index].target, 1),
            None => assert!(false),
        }
    }

    #[test]
    fn create_list_graph() {
        let mut graph = Graph { nodes: vec![], edges: vec![] };
        let node0 = graph.add_node("node0".to_string(), None);
        let node1 = graph.add_node("node1".to_string(), None);
        let node2 = graph.add_node("node2".to_string(), None);
        let node3 = graph.add_node("node3".to_string(), None);
        graph.add_edge(node0, node1, "edge0".to_string(), None);
        graph.add_edge(node1, node2, "edge1".to_string(), None);
        graph.add_edge(node2, node3, "edge2".to_string(), None);

        let mut targets = vec![];
        for e in graph.edges {
            targets.push(e.target);
        }
        assert_eq!(targets, [1,2,3]);
    }

    #[test]
    fn list_node_edges() {
        let mut graph = Graph { nodes: vec![], edges: vec![] };
        let node0 = graph.add_node("node0".to_string(), None);
        graph.add_edge(node0, node0, "edge0".to_string(), None);
        graph.add_edge(node0, node0, "edge1".to_string(), None);
        graph.add_edge(node0, node0, "edge2".to_string(), None);
        let mut targets = vec![];
        let mut successors = graph.successors(0);
        loop {
            match successors.next() {
                Some(edge) => targets.push(edge),
                None => break,
            }
        }
        assert_eq!(targets, [0,0,0]);
    }

    #[test]
    fn node_attributes() {
        let mut graph = Graph { nodes: vec![], edges: vec![] };

        let mut attributes = HashMap::new();
        attributes.insert("key".to_string(), "value".to_string());

        let node0 = graph.add_node("node0".to_string(), Some(attributes));
        match graph.nodes[node0].attributes {
            Some(ref attrs) => {
                match attrs.get("key") {
                    Some(value) => assert_eq!(value, &"value".to_string()),
                    None => assert!(false),
                }
            },
            None => assert!(false),
        }
    }

    #[test]
    fn edge_attributes() {
        let mut graph = Graph { nodes: vec![], edges: vec![] };

        let mut attributes = HashMap::new();
        attributes.insert("key".to_string(), "value".to_string());

        let node0 = graph.add_node("node0".to_string(), None);
        graph.add_edge(node0, node0, "edge0".to_string(), Some(attributes));
        match graph.edges[0].attributes {
            Some(ref attrs) => {
                match attrs.get("key") {
                    Some(value) => assert_eq!(value, &"value".to_string()),
                    None => assert!(false),
                }
            },
            None => assert!(false),
        }
    }

    #[test]
    fn node_edges() {
        let mut graph = Graph { nodes: vec![], edges: vec![] };
        let node0 = graph.add_node("node0".to_string(), None);
        graph.add_edge(node0, node0, "edge0".to_string(), None);
        graph.add_edge(node0, node0, "edge1".to_string(), None);
        graph.add_edge(node0, node0, "edge1".to_string(), None);

        let edges = graph.edges_for_node(node0);
        assert_eq!(edges, vec![2, 1, 0]);
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use option::{Option, Some, None};
use result::{Ok, Err};
use rt::io::net::ip::IpAddr;
use rt::io::{Reader, Writer};
use rt::io::{io_error, read_error, EndOfFile};
use rt::rtio::{RtioUdpSocketObject, RtioUdpSocket, IoFactory, IoFactoryObject};
use rt::local::Local;

pub struct UdpSocket {
    rtsocket: ~RtioUdpSocketObject
}

impl UdpSocket {
    fn new(s: ~RtioUdpSocketObject) -> UdpSocket {
        UdpSocket { rtsocket: s }
    }

    pub fn bind(addr: IpAddr) -> Option<UdpSocket> {
        let socket = unsafe {
            let io = Local::unsafe_borrow::<IoFactoryObject>();
            (*io).udp_bind(addr)
        };
        match socket {
            Ok(s) => { Some(UdpSocket { rtsocket: s }) }
            Err(ioerr) => {
                io_error::cond.raise(ioerr);
                return None;
            }
        }
    }

    pub fn recvfrom(&self, buf: &mut [u8]) -> Option<(uint, IpAddr)> {
        match (*self.rtsocket).recvfrom(buf) {
            Ok((nread, src)) => Some((nread, src)),
            Err(ioerr) => {
                // EOF is indicated by returning None
                if ioerr.kind != EndOfFile {
                    read_error::cond.raise(ioerr);
                }
                None
            }
        }
    }

    pub fn sendto(&self, buf: &[u8], dst: IpAddr) {
        match (*self.rtsocket).sendto(buf, dst) {
            Ok(_) => (),
            Err(ioerr) => {
                io_error::cond.raise(ioerr);
            }
        }
    }

    // XXX convert ~self to self eventually
    pub fn connect(~self, other: IpAddr) -> UdpStream {
        UdpStream { socket: self, connectedTo: other }
    }
}

pub struct UdpStream {
    socket: ~UdpSocket,
    connectedTo: IpAddr
}

impl UdpStream {
    pub fn as_socket<T>(&self, f: &fn(&UdpSocket) -> T) -> T {
        f(self.socket)
    }

    pub fn disconnect(self) -> ~UdpSocket {
        let UdpStream { socket: s, _ } = self;
        s
    }
}

impl Reader for UdpStream {
    fn read(&mut self, buf: &mut [u8]) -> Option<uint> {
        do self.as_socket |sock| {
            match sock.recvfrom(buf) {
                Some((_nread, src)) if src != self.connectedTo => Some(0),
                Some((nread, _src)) => Some(nread),
                None => None,
            }
        }
    }

    fn eof(&mut self) -> bool { fail!() }
}

impl Writer for UdpStream {
    fn write(&mut self, buf: &[u8]) {
        do self.as_socket |sock| {
            sock.sendto(buf, self.connectedTo);
        }
    }

    fn flush(&mut self) { fail!() }
}

#[cfg(test)]
mod test {
    use super::*;
    use rt::test::*;
    use rt::io::net::ip::Ipv4;
    use rt::io::*;
    use option::{Some, None};

    #[test]  #[ignore]
    fn bind_error() {
        do run_in_newsched_task {
            let mut called = false;
            do io_error::cond.trap(|e| {
                assert!(e.kind == PermissionDenied);
                called = true;
            }).in {
                let addr = Ipv4(0, 0, 0, 0, 1);
                let socket = UdpSocket::bind(addr);
                assert!(socket.is_none());
            }
            assert!(called);
        }
    }

    #[test]
    fn socket_smoke_test() {
        do run_in_newsched_task {
            let server_ip = next_test_ip4();
            let client_ip = next_test_ip4();

            do spawntask_immediately {
                match UdpSocket::bind(server_ip) {
                    Some(server) => {
                        let mut buf = [0];
                        match server.recvfrom(buf) {
                            Some((nread, src)) => {
                                assert_eq!(nread, 1);
                                assert_eq!(buf[0], 99);
                                assert_eq!(src, client_ip);
                            }
                            None => fail!()
                        }
                    }
                    None => fail!()
                }
            }

            do spawntask_immediately {
                match UdpSocket::bind(client_ip) {
                    Some(client) => client.sendto([99], server_ip),
                    None => fail!()
                }
            }
        }
    }

    #[test]
    fn stream_smoke_test() {
        do run_in_newsched_task {
            let server_ip = next_test_ip4();
            let client_ip = next_test_ip4();

            do spawntask_immediately {
                match UdpSocket::bind(server_ip) {
                    Some(server) => {
                        let server = ~server;
                        let mut stream = server.connect(client_ip);
                        let mut buf = [0];
                        match stream.read(buf) {
                            Some(nread) => {
                                assert_eq!(nread, 1);
                                assert_eq!(buf[0], 99);
                            }
                            None => fail!()
                        }
                    }
                    None => fail!()
                }
            }

            do spawntask_immediately {
                match UdpSocket::bind(client_ip) {
                    Some(client) => {
                        let client = ~client;
                        let mut stream = client.connect(server_ip);
                        stream.write([99]);
                    }
                    None => fail!()
                }
            }
        }
    }
}
<commit_msg>Converted UdpSocket into a newtype struct and (dis)connecting uses move semantics rather than ~.<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use option::{Option, Some, None};
use result::{Ok, Err};
use rt::io::net::ip::IpAddr;
use rt::io::{Reader, Writer};
use rt::io::{io_error, read_error, EndOfFile};
use rt::rtio::{RtioUdpSocketObject, RtioUdpSocket, IoFactory, IoFactoryObject};
use rt::local::Local;

pub struct UdpSocket(~RtioUdpSocketObject);

impl UdpSocket {
    fn new(s: ~RtioUdpSocketObject) -> UdpSocket { UdpSocket(s) }

    pub fn bind(addr: IpAddr) -> Option<UdpSocket> {
        let socket = unsafe { (*Local::unsafe_borrow::<IoFactoryObject>()).udp_bind(addr) };
        match socket {
            Ok(s) => Some(UdpSocket::new(s)),
            Err(ioerr) => {
                io_error::cond.raise(ioerr);
                None
            }
        }
    }

    pub fn recvfrom(&self, buf: &mut [u8]) -> Option<(uint, IpAddr)> {
        match (**self).recvfrom(buf) {
            Ok((nread, src)) => Some((nread, src)),
            Err(ioerr) => {
                // EOF is indicated by returning None
                if ioerr.kind != EndOfFile {
                    read_error::cond.raise(ioerr);
                }
                None
            }
        }
    }

    pub fn sendto(&self, buf: &[u8], dst: IpAddr) {
        match (**self).sendto(buf, dst) {
            Ok(_) => (),
            Err(ioerr) => io_error::cond.raise(ioerr),
        }
    }

    pub fn connect(self, other: IpAddr) -> UdpStream {
        UdpStream { socket: self, connectedTo: other }
    }
}

pub struct UdpStream {
    socket: UdpSocket,
    connectedTo: IpAddr
}

impl UdpStream {
    pub fn as_socket<T>(&self, f: &fn(&UdpSocket) -> T) -> T { f(&self.socket) }

    pub fn disconnect(self) -> UdpSocket { self.socket }
}

impl Reader for UdpStream {
    fn read(&mut self, buf: &mut [u8]) -> Option<uint> {
        do self.as_socket |sock| {
            match sock.recvfrom(buf) {
                Some((_nread, src)) if src != self.connectedTo => Some(0),
                Some((nread, _src)) => Some(nread),
                None => None,
            }
        }
    }

    fn eof(&mut self) -> bool { fail!() }
}

impl Writer for UdpStream {
    fn write(&mut self, buf: &[u8]) {
        do self.as_socket |sock| {
            sock.sendto(buf, self.connectedTo);
        }
    }

    fn flush(&mut self) { fail!() }
}

#[cfg(test)]
mod test {
    use super::*;
    use rt::test::*;
    use rt::io::net::ip::Ipv4;
    use rt::io::*;
    use option::{Some, None};

    #[test]  #[ignore]
    fn bind_error() {
        do run_in_newsched_task {
            let mut called = false;
            do io_error::cond.trap(|e| {
                assert!(e.kind == PermissionDenied);
                called = true;
            }).in {
                let addr = Ipv4(0, 0, 0, 0, 1);
                let socket = UdpSocket::bind(addr);
                assert!(socket.is_none());
            }
            assert!(called);
        }
    }

    #[test]
    fn socket_smoke_test() {
        do run_in_newsched_task {
            let server_ip = next_test_ip4();
            let client_ip = next_test_ip4();

            do spawntask_immediately {
                match UdpSocket::bind(server_ip) {
                    Some(server) => {
                        let mut buf = [0];
                        match server.recvfrom(buf) {
                            Some((nread, src)) => {
                                assert_eq!(nread, 1);
                                assert_eq!(buf[0], 99);
                                assert_eq!(src, client_ip);
                            }
                            None => fail!()
                        }
                    }
                    None => fail!()
                }
            }

            do spawntask_immediately {
                match UdpSocket::bind(client_ip) {
                    Some(client) => client.sendto([99], server_ip),
                    None => fail!()
                }
            }
        }
    }

    #[test]
    fn stream_smoke_test() {
        do run_in_newsched_task {
            let server_ip = next_test_ip4();
            let client_ip = next_test_ip4();

            do spawntask_immediately {
                match UdpSocket::bind(server_ip) {
                    Some(server) => {
                        let server = ~server;
                        let mut stream = server.connect(client_ip);
                        let mut buf = [0];
                        match stream.read(buf) {
                            Some(nread) => {
                                assert_eq!(nread, 1);
                                assert_eq!(buf[0], 99);
                            }
                            None => fail!()
                        }
                    }
                    None => fail!()
                }
            }

            do spawntask_immediately {
                match UdpSocket::bind(client_ip) {
                    Some(client) => {
                        let client = ~client;
                        let mut stream = client.connect(server_ip);
                        stream.write([99]);
                    }
                    None => fail!()
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate lwnvulns;
extern crate scraper;
extern crate curl;

use lwnvulns::tokenize::{tokenize, Token, Header, Issue, SourceLink};
use lwnvulns::parse::parse;
use lwnvulns::writer::write;
use lwnvulns::transform::collapse_anemic;
use scraper::{Html, Selector};
use curl::easy::Easy;

use std::io::{Read, Write};
use std::fs::File;

struct LwnRow {
    url: String,
    id: String,
    packages: String,
    description: String,
}

fn main() {
    let mut db: Vec<String> = vec![];
    {
        let mut f = File::open("./db").unwrap();
        let mut s = String::new();
        f.read_to_string(&mut s).unwrap();
        for entry in s.split_whitespace() {
            db.push(entry.to_string());
        }
    }

    let mut tokens: Vec<Token> = tokenize(include_str!("./instructions.md").to_string());

    let mut page = 0;
    let mut pages_with_nothing = 0;

    loop {
        let mut found = 0;
        for token in tokens_from_html(fetch_page(page)) {
            if let Token::Issue(ref issue) = token {
                if db.contains(&issue.source.id) {
                    continue;
                } else {
                    found += 1;
                }
            }
            tokens.push(token);
        }

        writeln!(&mut std::io::stderr(),
                 "Page {page} yielded {found} issues, after {none} pages with nothing",
                 page = page,
                 found = found,
                 none = pages_with_nothing)
            .unwrap();

        page += 1;

        if found == 0 {
            pages_with_nothing += 1;
        } else {
            pages_with_nothing = 0;
        }

        if pages_with_nothing >= 2 {
            break;
        }

        if page > 10 {
            panic!("Not sure we should be hitting 10 pages!");
        }
    }

    print!("{}", write(&collapse_anemic(parse(tokens).unwrap())));

}

fn fetch_page(page: usize) -> scraper::Html {
    let url = format!("https://lwn.net/Vulnerabilities/?n={}&offset={}",
                      100,
                      page * 100);
    let mut data = Vec::new();

    let mut easy = Easy::new();
    easy.url(&url).unwrap();

    {
        let mut transfer = easy.transfer();
        transfer.write_function(|chunk| {
                data.extend_from_slice(chunk);
                Ok(chunk.len())
            })
            .unwrap();
        transfer.perform().unwrap();
    }

    return Html::parse_document(&String::from_utf8(data).unwrap());
}

fn tokens_from_html(html: scraper::Html) -> Vec<Token> {
    let mut tokens: Vec<Token> = vec![];

    let rows = Selector::parse("tr").unwrap();
    let links = Selector::parse("td:nth-child(1) a").unwrap();
    let packages = Selector::parse("td:nth-child(3)").unwrap();
    let descriptions = Selector::parse("td:nth-child(4)").unwrap();
    for row in html.select(&rows) {
        if let Some(link) = row.select(&links).next() {
            let package = row.select(&packages).next().unwrap().text().next().unwrap();
            let description = row.select(&descriptions).next().unwrap().text().next().unwrap();

            let lwnrow = LwnRow {
                url: link.value().attr("href").unwrap().to_string(),
                id: link.inner_html(),
                packages: package.to_string(),
                description: description.to_string(),
            };
            for token in row_to_tokens(lwnrow) {
                tokens.push(token);
            }
        }
    }

    return tokens;
}

fn row_to_tokens(row: LwnRow) -> Vec<Token> {
    let mut tokens = vec![];

    let url = format!("https://lwn.net{}", row.url);
    for pkg in split_packages(row.packages) {
        tokens.push(Token::Header(Header {
            issue_count: 0,
            notes: None,
            package: pkg.clone(),
        }));

        tokens.push(Token::Issue(Issue {
            complete: false,
            source: SourceLink {
                id: row.id.clone(),
                url: url.clone(),
            },
            content: format!(
                "([search]({search_url}), [files]({files_url}))  {description}",
                search_url = format!(
                    "http://search.nix.gsc.io/?q={package}&i=fosho&repos=nixos-nixpkgs",
                    package=pkg
                ),
                files_url = format!(
                    "{root}?utf8=%E2%9C%93&q={package}+in%3Apath&type=Code",
                    root="https://github.com/NixOS/nixpkgs/search",
                    package=pkg
                ),
                description=row.description
            ),
        }));
    }

    return tokens;
}

fn split_packages(packages: String) -> Vec<String> {
    let pkgs = packages.split_whitespace()
        .map(|s| s.replace(",", "").to_string())
        .map(|s| if s == "tiff" {
            "libtiff".to_string()
        } else if s == "libgc" {
            "boehm-gc".to_string()
        } else if s == "perl-DBD-MySQL" {
            "DBD-mysql".to_string()
        } else {
            s
        })
        .collect();

    return pkgs;
}
<commit_msg>httpd: rename to apache-httpd<commit_after>extern crate lwnvulns;
extern crate scraper;
extern crate curl;

use lwnvulns::tokenize::{tokenize, Token, Header, Issue, SourceLink};
use lwnvulns::parse::parse;
use lwnvulns::writer::write;
use lwnvulns::transform::collapse_anemic;
use scraper::{Html, Selector};
use curl::easy::Easy;

use std::io::{Read, Write};
use std::fs::File;

struct LwnRow {
    url: String,
    id: String,
    packages: String,
    description: String,
}

fn main() {
    let mut db: Vec<String> = vec![];
    {
        let mut f = File::open("./db").unwrap();
        let mut s = String::new();
        f.read_to_string(&mut s).unwrap();
        for entry in s.split_whitespace() {
            db.push(entry.to_string());
        }
    }

    let mut tokens: Vec<Token> = tokenize(include_str!("./instructions.md").to_string());

    let mut page = 0;
    let mut pages_with_nothing = 0;

    loop {
        let mut found = 0;
        for token in tokens_from_html(fetch_page(page)) {
            if let Token::Issue(ref issue) = token {
                if db.contains(&issue.source.id) {
                    continue;
                } else {
                    found += 1;
                }
            }
            tokens.push(token);
        }

        writeln!(&mut std::io::stderr(),
                 "Page {page} yielded {found} issues, after {none} pages with nothing",
                 page = page,
                 found = found,
                 none = pages_with_nothing)
            .unwrap();

        page += 1;

        if found == 0 {
            pages_with_nothing += 1;
        } else {
            pages_with_nothing = 0;
        }

        if pages_with_nothing >= 2 {
            break;
        }

        if page > 10 {
            panic!("Not sure we should be hitting 10 pages!");
        }
    }

    print!("{}", write(&collapse_anemic(parse(tokens).unwrap())));

}

fn fetch_page(page: usize) -> scraper::Html {
    let url = format!("https://lwn.net/Vulnerabilities/?n={}&offset={}",
                      100,
                      page * 100);
    let mut data = Vec::new();

    let mut easy = Easy::new();
    easy.url(&url).unwrap();

    {
        let mut transfer = easy.transfer();
        transfer.write_function(|chunk| {
                data.extend_from_slice(chunk);
                Ok(chunk.len())
            })
            .unwrap();
        transfer.perform().unwrap();
    }

    return Html::parse_document(&String::from_utf8(data).unwrap());
}

fn tokens_from_html(html: scraper::Html) -> Vec<Token> {
    let mut tokens: Vec<Token> = vec![];

    let rows = Selector::parse("tr").unwrap();
    let links = Selector::parse("td:nth-child(1) a").unwrap();
    let packages = Selector::parse("td:nth-child(3)").unwrap();
    let descriptions = Selector::parse("td:nth-child(4)").unwrap();
    for row in html.select(&rows) {
        if let Some(link) = row.select(&links).next() {
            let package = row.select(&packages).next().unwrap().text().next().unwrap();
            let description = row.select(&descriptions).next().unwrap().text().next().unwrap();

            let lwnrow = LwnRow {
                url: link.value().attr("href").unwrap().to_string(),
                id: link.inner_html(),
                packages: package.to_string(),
                description: description.to_string(),
            };
            for token in row_to_tokens(lwnrow) {
                tokens.push(token);
            }
        }
    }

    return tokens;
}

fn row_to_tokens(row: LwnRow) -> Vec<Token> {
    let mut tokens = vec![];

    let url = format!("https://lwn.net{}", row.url);
    for pkg in split_packages(row.packages) {
        tokens.push(Token::Header(Header {
            issue_count: 0,
            notes: None,
            package: pkg.clone(),
        }));

        tokens.push(Token::Issue(Issue {
            complete: false,
            source: SourceLink {
                id: row.id.clone(),
                url: url.clone(),
            },
            content: format!(
                "([search]({search_url}), [files]({files_url}))  {description}",
                search_url = format!(
                    "http://search.nix.gsc.io/?q={package}&i=fosho&repos=nixos-nixpkgs",
                    package=pkg
                ),
                files_url = format!(
                    "{root}?utf8=%E2%9C%93&q={package}+in%3Apath&type=Code",
                    root="https://github.com/NixOS/nixpkgs/search",
                    package=pkg
                ),
                description=row.description
            ),
        }));
    }

    return tokens;
}

fn split_packages(packages: String) -> Vec<String> {
    let pkgs = packages.split_whitespace()
        .map(|s| s.replace(",", "").to_string())
        .map(|s| if s == "tiff" {
            "libtiff".to_string()
        } else if s == "libgc" {
            "boehm-gc".to_string()
        } else if s == "perl-DBD-MySQL" {
            "DBD-mysql".to_string()
        } else if s == "httpd" {
            "apache-httpd".to_string()
        } else {
            s
        })
        .collect();

    return pkgs;
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::rc::{Rc, Weak};
use log::LogLevel;
use mio::{Event, PollOpt, Ready, Token};
use mio::net::TcpStream;

use super::binary;
use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packetizer::Packetizer;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;
use super::tcp_header::{self, TCPHeader};
use super::transport_header::{TransportHeader, TransportHeaderMut};

const TAG: &'static str = "TCPConnection";

const MAX_PAYLOAD_LENGTH: u16 = 1400;

pub struct TCPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    stream: TcpStream,
    token: Token,
    client_to_network: StreamBuffer,
    network_to_client: Packetizer,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: u32,
    acknowledgement_number: u32,
    their_acknowledgement_number: u32,
    client_window: u16,
    remote_closed: bool,
}

#[derive(PartialEq, Eq)]
enum TCPState {
    Init,
    SynSent,
//    SynReceived,
    Established,
    CloseWait,
    LastAck,
}

impl TCB {
    fn new() -> Self {
        Self {
            state: TCPState::Init,
            syn_sequence_number: 0,
            sequence_number: 0,
            acknowledgement_number: 0,
            their_acknowledgement_number: 0,
            client_window: 0,
            remote_closed: false,
        }
    }

    fn numbers(&self) -> String {
        format!("(seq={}, ack={})", self.sequence_number, self.acknowledgement_number)
    }
}

impl TCPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, ipv4_header: &IPv4Header, transport_header: &TransportHeader) -> io::Result<Rc<RefCell<Self>>> {
        let stream = TCPConnection::create_stream(&id)?;

        if let TransportHeader::TCP(ref tcp_header) = *transport_header {
            // shrink the TCP options to pass a minimal refrence header to the packetizer
            let mut shrinked_tcp_header_raw = [0u8; 20];
            shrinked_tcp_header_raw.copy_from_slice(&tcp_header.raw()[..20]);
            let mut shrinked_tcp_header_data = tcp_header.data().clone();
            {
                let mut shrinked_tcp_header = shrinked_tcp_header_data.bind_mut(&mut shrinked_tcp_header_raw);
                shrinked_tcp_header.shrink_options();
                assert_eq!(20, shrinked_tcp_header.header_length());
            }

            let shrinked_transport_header = shrinked_tcp_header_data.bind(&shrinked_tcp_header_raw).into();

            let packetizer = Packetizer::new(&ipv4_header, &shrinked_transport_header);

            let rc = Rc::new(RefCell::new(Self {
                id: id,
                client: client,
                stream: stream,
                token: Token(0), // default value, will be set afterwards
                client_to_network: StreamBuffer::new(4 * MAX_PACKET_LENGTH),
                network_to_client: packetizer,
                packet_for_client_length: None,
                closed: false,
                tcb: TCB::new(),
            }));

            {
                let rc_clone = rc.clone();
                let handler = move |selector: &mut Selector, ready| {
                    let mut self_ref = rc_clone.borrow_mut();
                    self_ref.on_ready(selector, ready);
                };
                let mut self_ref = rc.borrow_mut();
                let token = selector.register(&self_ref.stream, handler, Ready::readable(), PollOpt::level())?;
                self_ref.token = token;
            }
            Ok(rc)
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_stream(id: &ConnectionId) -> io::Result<TcpStream> {
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        TcpStream::connect(&rewritten_destination)
    }

    fn close(&mut self, selector: &mut Selector) {

    }

    fn process_send(&mut self, selector: &mut Selector) {
        match self.client_to_network.write_to(&mut self.stream) {
            Ok(w) => if w == 0 {
                self.close(selector);
            },
            Err(err) => {
                error!(target: TAG, "{} Cannot write: {}", self.id, err);
                self.reset_connection(selector);
            }
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        assert!(self.packet_for_client_length.is_none(), "A pending packet was not sent");
        let remaining_client_window = self.get_remaining_client_window();
        assert!(remaining_client_window > 0, "process_received() must not be called when window == 0");
        let max_payload_length = cmp::min(remaining_client_window, MAX_PAYLOAD_LENGTH) as usize;
        TCPConnection::update_headers(&mut self.network_to_client, &self.tcb, tcp_header::FLAG_ACK | tcp_header::FLAG_PSH);
        // the packet is bound to the lifetime of self, so we cannot borrow self to call methods
        // defer the other branches in a separate match-block
        let non_lexical_lifetime_workaround = match self.network_to_client.packetize_read(&mut self.stream, Some(max_payload_length)) {
            Ok(Some(ipv4_packet)) => {
                // TODO packet source pull
                TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                error!(target: TAG, "{} Cannot read: {}", self.id, err);
                self.reset_connection(selector);
            },
            Ok(Some(_)) => () // already handled
        }
    }

    fn send_to_client(client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        let client_rc = client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.send_to_client(selector, &ipv4_packet)
    }

    fn send_to_client_ignore(id: &ConnectionId, client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        if let Err(err) = TCPConnection::send_to_client(client, selector, ipv4_packet) {
            warn!(target: TAG, "{} Cannot send packet to client: {}", id, err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += 1; // FIN counts for 1 byte
            TCPConnection::send_to_client_ignore(&self. id, &self.client, selector, &ipv4_packet);
        }
    }

    fn update_headers(packetizer: &mut Packetizer, tcb: &TCB, flags: u16) {
        if let TransportHeaderMut::TCP(ref mut tcp_header) = packetizer.transport_header_mut() {
            tcp_header.set_sequence_number(tcb.sequence_number);
            tcp_header.set_acknowledgement_number(tcb.acknowledgement_number);
            tcp_header.set_flags(flags);
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_empty_response_packet<'a>(id: &ConnectionId, packetizer: &'a mut Packetizer, tcb: &TCB, flags: u16) -> IPv4Packet<'a> {
        TCPConnection::update_headers(packetizer, tcb, flags);
        debug!(target: TAG, "{} Forging empty response (flags={}) {}", id, flags, tcb.numbers());
        if (flags & tcp_header::FLAG_ACK) != 0 {
            debug!(target: TAG, "{} Acking {}", id, tcb.numbers());
        }
        let ipv4_packet = packetizer.packetize_empty_payload();
        if log_enabled!(target: TAG, LogLevel::Trace) {
            binary::to_string(ipv4_packet.raw());
        }
        ipv4_packet
    }

    fn reset_connection(&mut self, selector: &mut Selector) {
        {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_RST);
            TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
        }
        self.close(selector);
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        if !self.closed {
            let mut ready = Ready::empty();
            if self.may_read() {
                ready = ready | Ready::readable()
            }
            if self.may_write() {
                ready = ready | Ready::writable()
            }
            selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
        }
    }

    fn may_read(&self) -> bool {
        !self.tcb.remote_closed &&
                self.packet_for_client_length.is_some() &&
                self.get_remaining_client_window() > 0
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn get_remaining_client_window(&self) -> u16 {
        // TODO
        42
    }
}

impl Connection for TCPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        info!(target: TAG, "{} Close", self.id);
        selector.deregister(&self.stream, self.token).unwrap();
        // socket will be closed by RAII
    }

    fn is_expired(&self) -> bool {
        // no external timeout expiration
        false
    }
}
<commit_msg>Implement TCPConnection.close()<commit_after>use std::cell::RefCell;
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::rc::{Rc, Weak};
use log::LogLevel;
use mio::{Event, PollOpt, Ready, Token};
use mio::net::TcpStream;

use super::binary;
use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packetizer::Packetizer;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;
use super::tcp_header::{self, TCPHeader};
use super::transport_header::{TransportHeader, TransportHeaderMut};

const TAG: &'static str = "TCPConnection";

const MAX_PAYLOAD_LENGTH: u16 = 1400;

pub struct TCPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    stream: TcpStream,
    token: Token,
    client_to_network: StreamBuffer,
    network_to_client: Packetizer,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: u32,
    acknowledgement_number: u32,
    their_acknowledgement_number: u32,
    client_window: u16,
    remote_closed: bool,
}

#[derive(PartialEq, Eq)]
enum TCPState {
    Init,
    SynSent,
//    SynReceived,
    Established,
    CloseWait,
    LastAck,
}

impl TCB {
    fn new() -> Self {
        Self {
            state: TCPState::Init,
            syn_sequence_number: 0,
            sequence_number: 0,
            acknowledgement_number: 0,
            their_acknowledgement_number: 0,
            client_window: 0,
            remote_closed: false,
        }
    }

    fn numbers(&self) -> String {
        format!("(seq={}, ack={})", self.sequence_number, self.acknowledgement_number)
    }
}

impl TCPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, ipv4_header: &IPv4Header, transport_header: &TransportHeader) -> io::Result<Rc<RefCell<Self>>> {
        let stream = TCPConnection::create_stream(&id)?;

        if let TransportHeader::TCP(ref tcp_header) = *transport_header {
            // shrink the TCP options to pass a minimal refrence header to the packetizer
            let mut shrinked_tcp_header_raw = [0u8; 20];
            shrinked_tcp_header_raw.copy_from_slice(&tcp_header.raw()[..20]);
            let mut shrinked_tcp_header_data = tcp_header.data().clone();
            {
                let mut shrinked_tcp_header = shrinked_tcp_header_data.bind_mut(&mut shrinked_tcp_header_raw);
                shrinked_tcp_header.shrink_options();
                assert_eq!(20, shrinked_tcp_header.header_length());
            }

            let shrinked_transport_header = shrinked_tcp_header_data.bind(&shrinked_tcp_header_raw).into();

            let packetizer = Packetizer::new(&ipv4_header, &shrinked_transport_header);

            let rc = Rc::new(RefCell::new(Self {
                id: id,
                client: client,
                stream: stream,
                token: Token(0), // default value, will be set afterwards
                client_to_network: StreamBuffer::new(4 * MAX_PACKET_LENGTH),
                network_to_client: packetizer,
                packet_for_client_length: None,
                closed: false,
                tcb: TCB::new(),
            }));

            {
                let rc_clone = rc.clone();
                let handler = move |selector: &mut Selector, ready| {
                    let mut self_ref = rc_clone.borrow_mut();
                    self_ref.on_ready(selector, ready);
                };
                let mut self_ref = rc.borrow_mut();
                let token = selector.register(&self_ref.stream, handler, Ready::readable(), PollOpt::level())?;
                self_ref.token = token;
            }
            Ok(rc)
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_stream(id: &ConnectionId) -> io::Result<TcpStream> {
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        TcpStream::connect(&rewritten_destination)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        self.disconnect(selector);

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        match self.client_to_network.write_to(&mut self.stream) {
            Ok(w) => if w == 0 {
                self.close(selector);
            },
            Err(err) => {
                error!(target: TAG, "{} Cannot write: {}", self.id, err);
                self.reset_connection(selector);
            }
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        assert!(self.packet_for_client_length.is_none(), "A pending packet was not sent");
        let remaining_client_window = self.get_remaining_client_window();
        assert!(remaining_client_window > 0, "process_received() must not be called when window == 0");
        let max_payload_length = cmp::min(remaining_client_window, MAX_PAYLOAD_LENGTH) as usize;
        TCPConnection::update_headers(&mut self.network_to_client, &self.tcb, tcp_header::FLAG_ACK | tcp_header::FLAG_PSH);
        // the packet is bound to the lifetime of self, so we cannot borrow self to call methods
        // defer the other branches in a separate match-block
        let non_lexical_lifetime_workaround = match self.network_to_client.packetize_read(&mut self.stream, Some(max_payload_length)) {
            Ok(Some(ipv4_packet)) => {
                // TODO packet source pull
                TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                error!(target: TAG, "{} Cannot read: {}", self.id, err);
                self.reset_connection(selector);
            },
            Ok(Some(_)) => () // already handled
        }
    }

    fn send_to_client(client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        let client_rc = client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.send_to_client(selector, &ipv4_packet)
    }

    fn send_to_client_ignore(id: &ConnectionId, client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        if let Err(err) = TCPConnection::send_to_client(client, selector, ipv4_packet) {
            warn!(target: TAG, "{} Cannot send packet to client: {}", id, err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += 1; // FIN counts for 1 byte
            TCPConnection::send_to_client_ignore(&self. id, &self.client, selector, &ipv4_packet);
        }
    }

    fn update_headers(packetizer: &mut Packetizer, tcb: &TCB, flags: u16) {
        if let TransportHeaderMut::TCP(ref mut tcp_header) = packetizer.transport_header_mut() {
            tcp_header.set_sequence_number(tcb.sequence_number);
            tcp_header.set_acknowledgement_number(tcb.acknowledgement_number);
            tcp_header.set_flags(flags);
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_empty_response_packet<'a>(id: &ConnectionId, packetizer: &'a mut Packetizer, tcb: &TCB, flags: u16) -> IPv4Packet<'a> {
        TCPConnection::update_headers(packetizer, tcb, flags);
        debug!(target: TAG, "{} Forging empty response (flags={}) {}", id, flags, tcb.numbers());
        if (flags & tcp_header::FLAG_ACK) != 0 {
            debug!(target: TAG, "{} Acking {}", id, tcb.numbers());
        }
        let ipv4_packet = packetizer.packetize_empty_payload();
        if log_enabled!(target: TAG, LogLevel::Trace) {
            binary::to_string(ipv4_packet.raw());
        }
        ipv4_packet
    }

    fn reset_connection(&mut self, selector: &mut Selector) {
        {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_RST);
            TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
        }
        self.close(selector);
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        if !self.closed {
            let mut ready = Ready::empty();
            if self.may_read() {
                ready = ready | Ready::readable()
            }
            if self.may_write() {
                ready = ready | Ready::writable()
            }
            selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
        }
    }

    fn may_read(&self) -> bool {
        !self.tcb.remote_closed &&
                self.packet_for_client_length.is_some() &&
                self.get_remaining_client_window() > 0
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn get_remaining_client_window(&self) -> u16 {
        // TODO
        42
    }
}

impl Connection for TCPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        info!(target: TAG, "{} Close", self.id);
        selector.deregister(&self.stream, self.token).unwrap();
        // socket will be closed by RAII
    }

    fn is_expired(&self) -> bool {
        // no external timeout expiration
        false
    }
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::rc::{Rc, Weak};
use log::LogLevel;
use mio::{Event, PollOpt, Ready, Token};
use mio::net::TcpStream;

use super::binary;
use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packetizer::Packetizer;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;
use super::tcp_header::{self, TCPHeader};
use super::transport_header::{TransportHeader, TransportHeaderMut};

const TAG: &'static str = "TCPConnection";

const MAX_PAYLOAD_LENGTH: u16 = 1400;

pub struct TCPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    stream: TcpStream,
    token: Token,
    client_to_network: StreamBuffer,
    network_to_client: Packetizer,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: u32,
    acknowledgement_number: u32,
    their_acknowledgement_number: u32,
    client_window: u16,
    remote_closed: bool,
}

#[derive(PartialEq, Eq)]
enum TCPState {
    Init,
    SynSent,
//    SynReceived,
    Established,
    CloseWait,
    LastAck,
}

impl TCB {
    fn new() -> Self {
        Self {
            state: TCPState::Init,
            syn_sequence_number: 0,
            sequence_number: 0,
            acknowledgement_number: 0,
            their_acknowledgement_number: 0,
            client_window: 0,
            remote_closed: false,
        }
    }

    fn numbers(&self) -> String {
        format!("(seq={}, ack={})", self.sequence_number, self.acknowledgement_number)
    }
}

impl TCPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, ipv4_header: &IPv4Header, transport_header: &TransportHeader) -> io::Result<Rc<RefCell<Self>>> {
        let stream = TCPConnection::create_stream(&id)?;

        if let TransportHeader::TCP(ref tcp_header) = *transport_header {
            // shrink the TCP options to pass a minimal refrence header to the packetizer
            let mut shrinked_tcp_header_raw = [0u8; 20];
            shrinked_tcp_header_raw.copy_from_slice(&tcp_header.raw()[..20]);
            let mut shrinked_tcp_header_data = tcp_header.data().clone();
            {
                let mut shrinked_tcp_header = shrinked_tcp_header_data.bind_mut(&mut shrinked_tcp_header_raw);
                shrinked_tcp_header.shrink_options();
                assert_eq!(20, shrinked_tcp_header.header_length());
            }

            let shrinked_transport_header = shrinked_tcp_header_data.bind(&shrinked_tcp_header_raw).into();

            let packetizer = Packetizer::new(&ipv4_header, &shrinked_transport_header);

            let rc = Rc::new(RefCell::new(Self {
                id: id,
                client: client,
                stream: stream,
                token: Token(0), // default value, will be set afterwards
                client_to_network: StreamBuffer::new(4 * MAX_PACKET_LENGTH),
                network_to_client: packetizer,
                packet_for_client_length: None,
                closed: false,
                tcb: TCB::new(),
            }));

            {
                let rc_clone = rc.clone();
                let handler = move |selector: &mut Selector, ready| {
                    let mut self_ref = rc_clone.borrow_mut();
                    self_ref.on_ready(selector, ready);
                };
                let mut self_ref = rc.borrow_mut();
                let token = selector.register(&self_ref.stream, handler, Ready::readable(), PollOpt::level())?;
                self_ref.token = token;
            }
            Ok(rc)
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_stream(id: &ConnectionId) -> io::Result<TcpStream> {
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        TcpStream::connect(&rewritten_destination)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        self.disconnect(selector);

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        match self.client_to_network.write_to(&mut self.stream) {
            Ok(w) => if w == 0 {
                self.close(selector);
            },
            Err(err) => {
                error!(target: TAG, "{} Cannot write: {}", self.id, err);
                self.reset_connection(selector);
            },
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        assert!(self.packet_for_client_length.is_none(), "A pending packet was not sent");
        let remaining_client_window = self.get_remaining_client_window();
        assert!(remaining_client_window > 0, "process_received() must not be called when window == 0");
        let max_payload_length = cmp::min(remaining_client_window, MAX_PAYLOAD_LENGTH) as usize;
        TCPConnection::update_headers(&mut self.network_to_client, &self.tcb, tcp_header::FLAG_ACK | tcp_header::FLAG_PSH);
        // the packet is bound to the lifetime of self, so we cannot borrow self to call methods
        // defer the other branches in a separate match-block
        let non_lexical_lifetime_workaround = match self.network_to_client.packetize_read(&mut self.stream, Some(max_payload_length)) {
            Ok(Some(ipv4_packet)) => {
                // TODO packet source pull
                TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                error!(target: TAG, "{} Cannot read: {}", self.id, err);
                self.reset_connection(selector);
            },
            Ok(Some(_)) => (), // already handled
        }
    }

    fn send_to_client(client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        let client_rc = client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.send_to_client(selector, &ipv4_packet)
    }

    fn send_to_client_ignore(id: &ConnectionId, client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        if let Err(err) = TCPConnection::send_to_client(client, selector, ipv4_packet) {
            warn!(target: TAG, "{} Cannot send packet to client: {}", id, err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += 1; // FIN counts for 1 byte
            TCPConnection::send_to_client_ignore(&self. id, &self.client, selector, &ipv4_packet);
        }
    }

    fn update_headers(packetizer: &mut Packetizer, tcb: &TCB, flags: u16) {
        if let TransportHeaderMut::TCP(ref mut tcp_header) = packetizer.transport_header_mut() {
            tcp_header.set_sequence_number(tcb.sequence_number);
            tcp_header.set_acknowledgement_number(tcb.acknowledgement_number);
            tcp_header.set_flags(flags);
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_empty_response_packet<'a>(id: &ConnectionId, packetizer: &'a mut Packetizer, tcb: &TCB, flags: u16) -> IPv4Packet<'a> {
        TCPConnection::update_headers(packetizer, tcb, flags);
        debug!(target: TAG, "{} Forging empty response (flags={}) {}", id, flags, tcb.numbers());
        if (flags & tcp_header::FLAG_ACK) != 0 {
            debug!(target: TAG, "{} Acking {}", id, tcb.numbers());
        }
        let ipv4_packet = packetizer.packetize_empty_payload();
        if log_enabled!(target: TAG, LogLevel::Trace) {
            binary::to_string(ipv4_packet.raw());
        }
        ipv4_packet
    }

    fn reset_connection(&mut self, selector: &mut Selector) {
        {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_RST);
            TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
        }
        self.close(selector);
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        if !self.closed {
            let mut ready = Ready::empty();
            if self.may_read() {
                ready = ready | Ready::readable()
            }
            if self.may_write() {
                ready = ready | Ready::writable()
            }
            selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
        }
    }

    fn may_read(&self) -> bool {
        !self.tcb.remote_closed &&
                self.packet_for_client_length.is_some() &&
                self.get_remaining_client_window() > 0
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn get_remaining_client_window(&self) -> u16 {
        // TODO
        42
    }
}

impl Connection for TCPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        info!(target: TAG, "{} Close", self.id);
        selector.deregister(&self.stream, self.token).unwrap();
        // socket will be closed by RAII
    }

    fn is_expired(&self) -> bool {
        // no external timeout expiration
        false
    }
}
<commit_msg>Introduce PacketStorage<commit_after>use std::cell::{Ref, RefCell};
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::rc::{Rc, Weak};
use log::LogLevel;
use mio::{Event, PollOpt, Ready, Token};
use mio::net::TcpStream;

use super::binary;
use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packetizer::Packetizer;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;
use super::tcp_header::{self, TCPHeader};
use super::transport_header::{TransportHeader, TransportHeaderMut};

const TAG: &'static str = "TCPConnection";

const MAX_PAYLOAD_LENGTH: u16 = 1400;

pub struct TCPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    stream: TcpStream,
    token: Token,
    client_to_network: StreamBuffer,
    network_to_client: Packetizer,
    pending_packet_for_client: PacketStorage,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: u32,
    acknowledgement_number: u32,
    their_acknowledgement_number: u32,
    client_window: u16,
    remote_closed: bool,
}

#[derive(PartialEq, Eq)]
enum TCPState {
    Init,
    SynSent,
//    SynReceived,
    Established,
    CloseWait,
    LastAck,
}

#[derive(Clone)]
struct PacketStorage(Rc<RefCell<Option<Box<[u8]>>>>);

impl PacketStorage {
    fn new() -> Self {
        PacketStorage(Rc::new(RefCell::new(None)))
    }

    fn set(&mut self, raw: &[u8]) {
        let data = raw.to_vec().into_boxed_slice();
        *self.0.borrow_mut() = Some(data);
    }

    //fn set_packet(&mut self, ipv4_packet: &IPv4Packet) {
    //   self.set(ipv4_packet.raw());
    //}

    fn get(&self) -> Ref<Option<Box<[u8]>>> {
        self.0.borrow()
    }

    fn has(&self) -> bool {
        self.get().is_some()
    }
}

impl TCB {
    fn new() -> Self {
        Self {
            state: TCPState::Init,
            syn_sequence_number: 0,
            sequence_number: 0,
            acknowledgement_number: 0,
            their_acknowledgement_number: 0,
            client_window: 0,
            remote_closed: false,
        }
    }

    fn numbers(&self) -> String {
        format!("(seq={}, ack={})", self.sequence_number, self.acknowledgement_number)
    }
}

impl TCPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, ipv4_header: &IPv4Header, transport_header: &TransportHeader) -> io::Result<Rc<RefCell<Self>>> {
        let stream = TCPConnection::create_stream(&id)?;

        if let TransportHeader::TCP(ref tcp_header) = *transport_header {
            // shrink the TCP options to pass a minimal refrence header to the packetizer
            let mut shrinked_tcp_header_raw = [0u8; 20];
            shrinked_tcp_header_raw.copy_from_slice(&tcp_header.raw()[..20]);
            let mut shrinked_tcp_header_data = tcp_header.data().clone();
            {
                let mut shrinked_tcp_header = shrinked_tcp_header_data.bind_mut(&mut shrinked_tcp_header_raw);
                shrinked_tcp_header.shrink_options();
                assert_eq!(20, shrinked_tcp_header.header_length());
            }

            let shrinked_transport_header = shrinked_tcp_header_data.bind(&shrinked_tcp_header_raw).into();

            let packetizer = Packetizer::new(&ipv4_header, &shrinked_transport_header);

            let rc = Rc::new(RefCell::new(Self {
                id: id,
                client: client,
                stream: stream,
                token: Token(0), // default value, will be set afterwards
                client_to_network: StreamBuffer::new(4 * MAX_PACKET_LENGTH),
                network_to_client: packetizer,
                pending_packet_for_client: PacketStorage::new(),
                packet_for_client_length: None,
                closed: false,
                tcb: TCB::new(),
            }));

            {
                let rc_clone = rc.clone();
                let handler = move |selector: &mut Selector, ready| {
                    let mut self_ref = rc_clone.borrow_mut();
                    self_ref.on_ready(selector, ready);
                };
                let mut self_ref = rc.borrow_mut();
                let token = selector.register(&self_ref.stream, handler, Ready::readable(), PollOpt::level())?;
                self_ref.token = token;
            }
            Ok(rc)
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_stream(id: &ConnectionId) -> io::Result<TcpStream> {
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        TcpStream::connect(&rewritten_destination)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        self.disconnect(selector);

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        match self.client_to_network.write_to(&mut self.stream) {
            Ok(w) => if w == 0 {
                self.close(selector);
            },
            Err(err) => {
                error!(target: TAG, "{} Cannot write: {}", self.id, err);
                self.reset_connection(selector);
            },
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        assert!(self.packet_for_client_length.is_none(), "A pending packet was not sent");
        let remaining_client_window = self.get_remaining_client_window();
        assert!(remaining_client_window > 0, "process_received() must not be called when window == 0");
        let max_payload_length = cmp::min(remaining_client_window, MAX_PAYLOAD_LENGTH) as usize;
        TCPConnection::update_headers(&mut self.network_to_client, &self.tcb, tcp_header::FLAG_ACK | tcp_header::FLAG_PSH);
        // the packet is bound to the lifetime of self, so we cannot borrow self to call methods
        // defer the other branches in a separate match-block
        let non_lexical_lifetime_workaround = match self.network_to_client.packetize_read(&mut self.stream, Some(max_payload_length)) {
            Ok(Some(ipv4_packet)) => {
                // TODO packet source pull
                TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                error!(target: TAG, "{} Cannot read: {}", self.id, err);
                self.reset_connection(selector);
            },
            Ok(Some(_)) => (), // already handled
        }
    }

    fn send_to_client(client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        let client_rc = client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.send_to_client(selector, &ipv4_packet)
    }

    fn send_to_client_ignore(id: &ConnectionId, client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        if let Err(err) = TCPConnection::send_to_client(client, selector, ipv4_packet) {
            warn!(target: TAG, "{} Cannot send packet to client: {}", id, err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += 1; // FIN counts for 1 byte
            TCPConnection::send_to_client_ignore(&self. id, &self.client, selector, &ipv4_packet);
        }
    }

    fn update_headers(packetizer: &mut Packetizer, tcb: &TCB, flags: u16) {
        if let TransportHeaderMut::TCP(ref mut tcp_header) = packetizer.transport_header_mut() {
            tcp_header.set_sequence_number(tcb.sequence_number);
            tcp_header.set_acknowledgement_number(tcb.acknowledgement_number);
            tcp_header.set_flags(flags);
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_empty_response_packet<'a>(id: &ConnectionId, packetizer: &'a mut Packetizer, tcb: &TCB, flags: u16) -> IPv4Packet<'a> {
        TCPConnection::update_headers(packetizer, tcb, flags);
        debug!(target: TAG, "{} Forging empty response (flags={}) {}", id, flags, tcb.numbers());
        if (flags & tcp_header::FLAG_ACK) != 0 {
            debug!(target: TAG, "{} Acking {}", id, tcb.numbers());
        }
        let ipv4_packet = packetizer.packetize_empty_payload();
        if log_enabled!(target: TAG, LogLevel::Trace) {
            binary::to_string(ipv4_packet.raw());
        }
        ipv4_packet
    }

    fn reset_connection(&mut self, selector: &mut Selector) {
        {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_RST);
            TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
        }
        self.close(selector);
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        if !self.closed {
            let mut ready = Ready::empty();
            if self.may_read() {
                ready = ready | Ready::readable()
            }
            if self.may_write() {
                ready = ready | Ready::writable()
            }
            selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
        }
    }

    fn may_read(&self) -> bool {
        !self.tcb.remote_closed &&
                self.packet_for_client_length.is_some() &&
                self.get_remaining_client_window() > 0
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn get_remaining_client_window(&self) -> u16 {
        // TODO
        42
    }
}

impl Connection for TCPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        info!(target: TAG, "{} Close", self.id);
        selector.deregister(&self.stream, self.token).unwrap();
        // socket will be closed by RAII
    }

    fn is_expired(&self) -> bool {
        // no external timeout expiration
        false
    }
}
<|endoftext|>


<commit_before>//! Defines `DataContext`.

use cretonne_codegen::entity::PrimaryMap;
use cretonne_codegen::binemit::{CodeOffset, Addend};
use cretonne_codegen::ir;

/// This specifies how data is to be initialized.
#[derive(PartialEq, Eq, Debug)]
pub enum Init {
    /// This indicates that no initialization has been specified yet.
    Uninitialized,
    /// Initialize the data with all zeros.
    Zeros {
        /// The size of the data.
        size: usize,
    },
    /// Initialize the data with the specified contents.
    Bytes {
        /// The contents, which also implies the size of the data.
        contents: Box<[u8]>,
    },
}

impl Init {
    /// Return the size of the data to be initialized.
    pub fn size(&self) -> usize {
        match *self {
            Init::Uninitialized => panic!("data size not initialized yet"),
            Init::Zeros { size } => size,
            Init::Bytes { ref contents } => contents.len(),
        }
    }
}

/// A flag specifying whether data is readonly or may be written to.
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum Writability {
    /// Data is readonly, meaning writes to it will trap.
    Readonly,
    /// Data is writable.
    Writable,
}

/// A description of a data object.
pub struct DataDescription {
    /// Whether the data readonly or writable.
    pub writable: Writability,
    /// How the data should be initialized.
    pub init: Init,
    /// External function declarations.
    pub function_decls: PrimaryMap<ir::FuncRef, ir::ExternalName>,
    /// External data object declarations.
    pub data_decls: PrimaryMap<ir::GlobalVar, ir::ExternalName>,
    /// Function addresses to write at specified offsets.
    pub function_relocs: Vec<(CodeOffset, ir::FuncRef)>,
    /// Data addresses to write at specified offsets.
    pub data_relocs: Vec<(CodeOffset, ir::GlobalVar, Addend)>,
}

/// This is to data objects what cretonne_codegen::Context is to functions.
pub struct DataContext {
    description: DataDescription,
}

impl DataContext {
    /// Allocate a new context.
    pub fn new() -> Self {
        Self {
            description: DataDescription {
                writable: Writability::Readonly,
                init: Init::Uninitialized,
                function_decls: PrimaryMap::new(),
                data_decls: PrimaryMap::new(),
                function_relocs: Vec::new(),
                data_relocs: Vec::new(),
            },
        }
    }

    /// Clear all data structures in this context.
    pub fn clear(&mut self) {
        self.description.writable = Writability::Readonly;
        self.description.init = Init::Uninitialized;
        self.description.function_decls.clear();
        self.description.data_decls.clear();
        self.description.function_relocs.clear();
        self.description.data_relocs.clear();
    }

    /// Define a zero-initialized object with the given size.
    pub fn define_zeroinit(&mut self, size: usize, writable: Writability) {
        debug_assert_eq!(self.description.init, Init::Uninitialized);
        self.description.writable = writable;
        self.description.init = Init::Zeros { size };
    }

    /// Define a zero-initialized object with the given size.
    ///
    /// TODO: Can we avoid a Box here?
    pub fn define(&mut self, contents: Box<[u8]>, writable: Writability) {
        debug_assert_eq!(self.description.init, Init::Uninitialized);
        self.description.writable = writable;
        self.description.init = Init::Bytes { contents };
    }

    /// Declare an external function import.
    pub fn import_function(&mut self, name: ir::ExternalName) -> ir::FuncRef {
        self.description.function_decls.push(name)
    }

    /// Declares a global variable import.
    ///
    /// TODO: Rename to import_data?
    pub fn import_global_var(&mut self, name: ir::ExternalName) -> ir::GlobalVar {
        self.description.data_decls.push(name)
    }

    /// Write the address of `func` into the data at offset `offset`.
    pub fn write_function_addr(&mut self, offset: CodeOffset, func: ir::FuncRef) {
        self.description.function_relocs.push((offset, func))
    }

    /// Write the address of `data` into the data at offset `offset`.
    pub fn write_data_addr(&mut self, offset: CodeOffset, data: ir::GlobalVar, addend: Addend) {
        self.description.data_relocs.push((offset, data, addend))
    }

    /// Reference the initializer data.
    pub fn description(&self) -> &DataDescription {
        debug_assert!(
            self.description.init != Init::Uninitialized,
            "data must be initialized first"
        );
        &self.description
    }
}

#[cfg(test)]
mod tests {
    use {DataContext, Writability, Init};
    use cretonne_codegen::ir;

    #[test]
    fn basic_data_context() {
        let mut data_ctx = DataContext::new();
        {
            let description = &data_ctx.description;
            assert_eq!(description.writable, Writability::Readonly);
            assert_eq!(description.init, Init::Uninitialized);
            assert!(description.function_decls.is_empty());
            assert!(description.data_decls.is_empty());
            assert!(description.function_relocs.is_empty());
            assert!(description.data_relocs.is_empty());
        }

        data_ctx.define_zeroinit(256, Writability::Writable);

        let _func_a = data_ctx.import_function(ir::ExternalName::user(0, 0));
        let func_b = data_ctx.import_function(ir::ExternalName::user(0, 1));
        let func_c = data_ctx.import_function(ir::ExternalName::user(1, 0));
        let _data_a = data_ctx.import_global_var(ir::ExternalName::user(2, 2));
        let data_b = data_ctx.import_global_var(ir::ExternalName::user(2, 3));

        data_ctx.write_function_addr(8, func_b);
        data_ctx.write_function_addr(16, func_c);
        data_ctx.write_data_addr(32, data_b, 27);

        {
            let description = data_ctx.description();
            assert_eq!(description.writable, Writability::Writable);
            assert_eq!(description.init, Init::Zeros { size: 256 });
            assert_eq!(description.function_decls.len(), 3);
            assert_eq!(description.data_decls.len(), 2);
            assert_eq!(description.function_relocs.len(), 2);
            assert_eq!(description.data_relocs.len(), 1);
        }

        data_ctx.clear();
        {
            let description = &data_ctx.description;
            assert_eq!(description.writable, Writability::Readonly);
            assert_eq!(description.init, Init::Uninitialized);
            assert!(description.function_decls.is_empty());
            assert!(description.data_decls.is_empty());
            assert!(description.function_relocs.is_empty());
            assert!(description.data_relocs.is_empty());
        }

        let contents = vec![33, 34, 35, 36];
        let contents_clone = contents.clone();
        data_ctx.define(contents.into_boxed_slice(), Writability::Readonly);
        {
            let description = data_ctx.description();
            assert_eq!(description.writable, Writability::Readonly);
            assert_eq!(
                description.init,
                Init::Bytes { contents: contents_clone.into_boxed_slice() }
            );
            assert_eq!(description.function_decls.len(), 0);
            assert_eq!(description.data_decls.len(), 0);
            assert_eq!(description.function_relocs.len(), 0);
            assert_eq!(description.data_relocs.len(), 0);
        }
    }
}
<commit_msg>Add comments to DataContext's import_function and import_global_var.<commit_after>//! Defines `DataContext`.

use cretonne_codegen::entity::PrimaryMap;
use cretonne_codegen::binemit::{CodeOffset, Addend};
use cretonne_codegen::ir;

/// This specifies how data is to be initialized.
#[derive(PartialEq, Eq, Debug)]
pub enum Init {
    /// This indicates that no initialization has been specified yet.
    Uninitialized,
    /// Initialize the data with all zeros.
    Zeros {
        /// The size of the data.
        size: usize,
    },
    /// Initialize the data with the specified contents.
    Bytes {
        /// The contents, which also implies the size of the data.
        contents: Box<[u8]>,
    },
}

impl Init {
    /// Return the size of the data to be initialized.
    pub fn size(&self) -> usize {
        match *self {
            Init::Uninitialized => panic!("data size not initialized yet"),
            Init::Zeros { size } => size,
            Init::Bytes { ref contents } => contents.len(),
        }
    }
}

/// A flag specifying whether data is readonly or may be written to.
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum Writability {
    /// Data is readonly, meaning writes to it will trap.
    Readonly,
    /// Data is writable.
    Writable,
}

/// A description of a data object.
pub struct DataDescription {
    /// Whether the data readonly or writable.
    pub writable: Writability,
    /// How the data should be initialized.
    pub init: Init,
    /// External function declarations.
    pub function_decls: PrimaryMap<ir::FuncRef, ir::ExternalName>,
    /// External data object declarations.
    pub data_decls: PrimaryMap<ir::GlobalVar, ir::ExternalName>,
    /// Function addresses to write at specified offsets.
    pub function_relocs: Vec<(CodeOffset, ir::FuncRef)>,
    /// Data addresses to write at specified offsets.
    pub data_relocs: Vec<(CodeOffset, ir::GlobalVar, Addend)>,
}

/// This is to data objects what cretonne_codegen::Context is to functions.
pub struct DataContext {
    description: DataDescription,
}

impl DataContext {
    /// Allocate a new context.
    pub fn new() -> Self {
        Self {
            description: DataDescription {
                writable: Writability::Readonly,
                init: Init::Uninitialized,
                function_decls: PrimaryMap::new(),
                data_decls: PrimaryMap::new(),
                function_relocs: Vec::new(),
                data_relocs: Vec::new(),
            },
        }
    }

    /// Clear all data structures in this context.
    pub fn clear(&mut self) {
        self.description.writable = Writability::Readonly;
        self.description.init = Init::Uninitialized;
        self.description.function_decls.clear();
        self.description.data_decls.clear();
        self.description.function_relocs.clear();
        self.description.data_relocs.clear();
    }

    /// Define a zero-initialized object with the given size.
    pub fn define_zeroinit(&mut self, size: usize, writable: Writability) {
        debug_assert_eq!(self.description.init, Init::Uninitialized);
        self.description.writable = writable;
        self.description.init = Init::Zeros { size };
    }

    /// Define a zero-initialized object with the given size.
    ///
    /// TODO: Can we avoid a Box here?
    pub fn define(&mut self, contents: Box<[u8]>, writable: Writability) {
        debug_assert_eq!(self.description.init, Init::Uninitialized);
        self.description.writable = writable;
        self.description.init = Init::Bytes { contents };
    }

    /// Declare an external function import.
    ///
    /// Users of the `Module` API generally should call
    /// `Module::declare_func_in_data` instead, as it takes care of generating
    /// the appropriate `ExternalName`.
    pub fn import_function(&mut self, name: ir::ExternalName) -> ir::FuncRef {
        self.description.function_decls.push(name)
    }

    /// Declares a global variable import.
    ///
    /// TODO: Rename to import_data?
    ///
    /// Users of the `Module` API generally should call
    /// `Module::declare_data_in_data` instead, as it takes care of generating
    /// the appropriate `ExternalName`.
    pub fn import_global_var(&mut self, name: ir::ExternalName) -> ir::GlobalVar {
        self.description.data_decls.push(name)
    }

    /// Write the address of `func` into the data at offset `offset`.
    pub fn write_function_addr(&mut self, offset: CodeOffset, func: ir::FuncRef) {
        self.description.function_relocs.push((offset, func))
    }

    /// Write the address of `data` into the data at offset `offset`.
    pub fn write_data_addr(&mut self, offset: CodeOffset, data: ir::GlobalVar, addend: Addend) {
        self.description.data_relocs.push((offset, data, addend))
    }

    /// Reference the initializer data.
    pub fn description(&self) -> &DataDescription {
        debug_assert!(
            self.description.init != Init::Uninitialized,
            "data must be initialized first"
        );
        &self.description
    }
}

#[cfg(test)]
mod tests {
    use {DataContext, Writability, Init};
    use cretonne_codegen::ir;

    #[test]
    fn basic_data_context() {
        let mut data_ctx = DataContext::new();
        {
            let description = &data_ctx.description;
            assert_eq!(description.writable, Writability::Readonly);
            assert_eq!(description.init, Init::Uninitialized);
            assert!(description.function_decls.is_empty());
            assert!(description.data_decls.is_empty());
            assert!(description.function_relocs.is_empty());
            assert!(description.data_relocs.is_empty());
        }

        data_ctx.define_zeroinit(256, Writability::Writable);

        let _func_a = data_ctx.import_function(ir::ExternalName::user(0, 0));
        let func_b = data_ctx.import_function(ir::ExternalName::user(0, 1));
        let func_c = data_ctx.import_function(ir::ExternalName::user(1, 0));
        let _data_a = data_ctx.import_global_var(ir::ExternalName::user(2, 2));
        let data_b = data_ctx.import_global_var(ir::ExternalName::user(2, 3));

        data_ctx.write_function_addr(8, func_b);
        data_ctx.write_function_addr(16, func_c);
        data_ctx.write_data_addr(32, data_b, 27);

        {
            let description = data_ctx.description();
            assert_eq!(description.writable, Writability::Writable);
            assert_eq!(description.init, Init::Zeros { size: 256 });
            assert_eq!(description.function_decls.len(), 3);
            assert_eq!(description.data_decls.len(), 2);
            assert_eq!(description.function_relocs.len(), 2);
            assert_eq!(description.data_relocs.len(), 1);
        }

        data_ctx.clear();
        {
            let description = &data_ctx.description;
            assert_eq!(description.writable, Writability::Readonly);
            assert_eq!(description.init, Init::Uninitialized);
            assert!(description.function_decls.is_empty());
            assert!(description.data_decls.is_empty());
            assert!(description.function_relocs.is_empty());
            assert!(description.data_relocs.is_empty());
        }

        let contents = vec![33, 34, 35, 36];
        let contents_clone = contents.clone();
        data_ctx.define(contents.into_boxed_slice(), Writability::Readonly);
        {
            let description = data_ctx.description();
            assert_eq!(description.writable, Writability::Readonly);
            assert_eq!(
                description.init,
                Init::Bytes { contents: contents_clone.into_boxed_slice() }
            );
            assert_eq!(description.function_decls.len(), 0);
            assert_eq!(description.data_decls.len(), 0);
            assert_eq!(description.function_relocs.len(), 0);
            assert_eq!(description.data_relocs.len(), 0);
        }
    }
}
<|endoftext|>