<commit_before>extern crate clap;
extern crate git2;

use git2::{Repository, Commit, Oid, Revwalk, Index};
use std::collections::HashMap;

#[macro_use]
mod macros;

const E_SUCCESS: i32 = 0;
const E_NO_GIT_REPO: i32 = 1;

fn main() {
    let exit_code = real_main();
    std::process::exit(exit_code);
}

fn real_main() -> i32 {
    let submodule_dir = parse_cli_arguments();
    println!("Merging {}...", submodule_dir);

    let repo = match Repository::open(".") {
        Ok(repo) => repo,
        Err(e) => {
            eprintln!("Couldn't find Git repo in the current directory: {}",
                      e.message());
            return E_NO_GIT_REPO;
        }
    };

    let mut submodule_revwalk = get_submodule_revwalk(&repo, &submodule_dir);

    let mut old_id_to_new = HashMap::new();

    rewrite_submodule_history(&repo,
                              &mut submodule_revwalk,
                              &mut old_id_to_new,
                              &submodule_dir);

    let mut repo_revwalk = get_repo_revwalk(&repo);
    rewrite_repo_history(&repo, &mut repo_revwalk, &mut old_id_to_new, &submodule_dir);

    checkout_rewritten_history(&repo, &old_id_to_new);

    E_SUCCESS
}

fn parse_cli_arguments() -> String {
    let options = clap::App::new("git-submerge")
                          .version("0.1")
                          .author("Alexander Batischev <eual.jp@gmail.com>")
                          .about("Merge Git submodule into the main repo \
                                  as if they've never been separate at all")
                          .arg(clap::Arg::with_name("SUBMODULE_DIR")
                               .help("The submodule to merge")
                               .required(true)
                               .index(1))
                          .get_matches();

    // We can safely use unwrap() here because the argument is marked as "required" and Clap checks
    // its presence for us.
    String::from( options.value_of("SUBMODULE_DIR").unwrap() )
}

fn get_submodule_revwalk<'repo>(repo: &'repo Repository, submodule_dir: &str) -> Revwalk<'repo> {
    let submodule_url = String::from("./") + submodule_dir;
    let mut remote = repo.remote_anonymous(&submodule_url).expect("Couldn't create an anonymous remote");
    remote.fetch(&[], None, None).expect("Couldn't fetch submodule's history");
    let submodule = repo.find_submodule(submodule_dir)
        .expect("Couldn't find the submodule with expected path");
    let submodule_head = submodule.head_id()
        .expect("Couldn't obtain submodule's HEAD");

    let mut revwalk = repo.revwalk().expect("Couldn't obtain RevWalk object for the repo");
    // "Topological" and reverse means "parents are always visited before their children".
    // We need that in order to be sure that our old-to-new-ids map always contains everything we
    // need it to contain.
    revwalk.set_sorting(git2::SORT_REVERSE | git2::SORT_TOPOLOGICAL);
    // TODO (#6): push all branches and tags, not just HEAD
    revwalk.push(submodule_head).expect("Couldn't add submodule's HEAD to RevWalk list");

    revwalk
}

fn rewrite_submodule_history(repo: &Repository,
                             revwalk: &mut Revwalk,
                             old_id_to_new: &mut HashMap<Oid, Oid>,
                             submodule_dir: &str) {
    for maybe_oid in revwalk {
        match maybe_oid {
            Ok(oid) => {
                let commit = repo.find_commit(oid)
                    .expect(&format!("Couldn't get a commit with ID {}", oid));
                let tree = commit.tree()
                    .expect(&format!("Couldn't obtain the tree of a commit with ID {}", oid));
                // TODO (#10): consider using TreeBuilder instead of in-memory index
                let mut old_index = Index::new()
                    .expect("Couldn't create an in-memory index for commit");
                let mut new_index = Index::new().expect("Couldn't create an in-memory index");
                old_index.read_tree(&tree)
                    .expect(&format!("Couldn't read the commit {} into index", oid));

                // Obtain the new tree, where everything from the old one is moved under
                // a directory named after the submodule
                for entry in old_index.iter() {
                    let mut new_entry = entry;

                    // TODO (#11): what mode, owner, mtime etc. does the newly created dir get?
                    let mut new_path = String::from(submodule_dir);
                    new_path += "/";
                    new_path += &String::from_utf8(new_entry.path)
                        .expect("Failed to convert a path to str");

                    new_entry.path = new_path.into_bytes();
                    new_index.add(&new_entry).expect("Couldn't add an entry to the index");
                }
                let tree_id = new_index.write_tree_to(&repo)
                    .expect("Couldn't write the index into a tree");
                old_id_to_new.insert(tree.id(), tree_id);
                let tree = repo.find_tree(tree_id)
                    .expect("Couldn't retrieve the tree we just created");

                let parents = {
                    let mut p: Vec<Commit> = Vec::new();
                    for parent_id in commit.parent_ids() {
                        let new_parent_id = old_id_to_new[&parent_id];
                        let parent = repo.find_commit(new_parent_id)
                            .expect("Couldn't find parent commit by its id");
                        p.push(parent);
                    }
                    p
                };

                let mut parents_refs: Vec<&Commit> = Vec::new();
                for i in 0..parents.len() {
                    parents_refs.push(&parents[i]);
                }
                let new_commit_id = repo.commit(None,
                            &commit.author(),
                            &commit.committer(),
                            &commit.message().expect("Couldn't retrieve commit's message"),
                            &tree,
                            &parents_refs[..])
                    .expect("Failed to commit");

                old_id_to_new.insert(oid, new_commit_id);
            }
            Err(e) => eprintln!("Error walking the submodule's history: {:?}", e),
        }
    }
}

fn get_repo_revwalk<'repo>(repo: &'repo Repository) -> Revwalk<'repo> {
    let mut revwalk = repo.revwalk().expect("Couldn't obtain RevWalk object for the repo");
    revwalk.set_sorting(git2::SORT_REVERSE | git2::SORT_TOPOLOGICAL);
    let head = repo.head().expect("Couldn't obtain repo's HEAD");
    let head_id = head.target().expect("Couldn't resolve repo's HEAD to a commit ID");
    // TODO (#6): push all branches and tags, not just HEAD
    revwalk.push(head_id).expect("Couldn't add repo's HEAD to RevWalk list");

    revwalk
}

fn rewrite_repo_history(repo: &Repository,
                        revwalk: &mut Revwalk,
                        old_id_to_new: &mut HashMap<Oid, Oid>,
                        submodule_dir: &str) {
    let submodule_path = std::path::Path::new(submodule_dir);

    for maybe_oid in revwalk {
        match maybe_oid {
            Ok(oid) => {
                let commit = repo.find_commit(oid)
                    .expect(&format!("Couldn't get a commit with ID {}", oid));
                let tree = commit.tree()
                    .expect(&format!("Couldn't obtain the tree of a commit with ID {}", oid));

                let submodule_subdir = match tree.get_path(submodule_path) {
                    Ok(tree) => tree,
                    Err(e) => {
                        if e.code() == git2::ErrorCode::NotFound &&
                           e.class() == git2::ErrorClass::Tree {
                            // It's okay. The tree lacks the subtree corresponding to the
                            // submodule. In other words, the commit doesn't include the submodule.
                            // That's totally fine. Let's map it into itself and move on.
                            old_id_to_new.insert(oid, oid);
                            continue;
                        } else {
                            // Unexpected error; let's report it and abort the program
                            panic!("Error getting submodule's subdir from the tree: {:?}", e);
                        };
                    }
                };

                // **INVARIANT**: if we got this far, current commit contains a submodule and
                // should be rewritten

                let submodule_commit_id = submodule_subdir.id();
                let new_submodule_commit_id = old_id_to_new[&submodule_commit_id];
                let submodule_commit = repo.find_commit(new_submodule_commit_id)
                    .expect("Couldn't obtain submodule's commit");
                let subtree_id = submodule_commit.tree()
                    .and_then(|t| t.get_path(submodule_path))
                    .and_then(|te| Ok(te.id()))
                    .expect("Couldn't obtain submodule's subtree ID");

                let mut treebuilder = repo.treebuilder(Some(&tree))
                    .expect("Couldn't create TreeBuilder");
                treebuilder.remove(submodule_path)
                    .expect("Couldn't remove submodule path from TreeBuilder");
                treebuilder.insert(submodule_path, subtree_id, 0o040000)
                    .expect("Couldn't add submodule as a subdir to TreeBuilder");
                let new_tree_id = treebuilder.write()
                    .expect("Couldn't write TreeBuilder into a Tree");
                let new_tree = repo.find_tree(new_tree_id)
                    .expect("Couldn't read back the Tree we just wrote");

                // In commits that used to update the submodule, add a parent pointing to
                // appropriate commit in new submodule history
                let mut parent_subtree_ids = std::collections::HashSet::new();
                for parent in commit.parents() {
                    let parent_tree = parent.tree().expect("Couldn't obtain parent's tree");
                    let parent_subdir_tree_id = parent_tree.get_path(submodule_path)
                        .and_then(|x| Ok(x.id()));

                    match parent_subdir_tree_id {
                        Ok(id) => {
                            parent_subtree_ids.insert(id);
                            ()
                        }
                        Err(e) => {
                            if e.code() == git2::ErrorCode::NotFound &&
                               e.class() == git2::ErrorClass::Tree
                            {
                                continue;
                            } else {
                                panic!("Error getting submodule's subdir from the tree: {:?}", e);
                            };
                        }
                    }
                }

                // Here's a few pictures to help you understand how we figure out if current commit
                // updated the submodule. If we draw a DAG and name submodule states, the following
                // situations will mean that the submodule wasn't updated:
                //
                //     o--o--o--A--
                //                 `,-A
                //      o--o--o--B-
                //
                // or
                //
                //     o--o--o--A--
                //                 `,-B
                //      o--o--o--B-
                //
                // And in the following graphs the submodule was updated:
                //
                //     o--o--o--A--
                //                 `,-C
                //      o--o--o--B-
                //
                // or
                //
                //     o--o--o--o--A--B
                //
                // Put into words, the rule will be "the submodule state in current commit is
                // different from states in all its parents". Or, more formally, the current state
                // doesn't belong to the set of states in parents.
                let submodule_updated: bool = !parent_subtree_ids.contains(&submodule_commit_id);

                // Rewrite the parents if the submodule was updated
                let parents = {
                    let mut p: Vec<Commit> = Vec::new();
                    for parent_id in commit.parent_ids() {
                        let actual_parent_id = old_id_to_new[&parent_id];
                        let parent = repo.find_commit(actual_parent_id)
                            .expect("Couldn't find parent commit by its id");
                        p.push(parent);
                    }

                    if submodule_updated {
                        p.push(submodule_commit);
                    }

                    p
                };

                let mut parents_refs: Vec<&Commit> = Vec::new();
                for i in 0..parents.len() {
                    parents_refs.push(&parents[i]);
                }
                let new_commit_id = repo.commit(None,
                            &commit.author(),
                            &commit.committer(),
                            &commit.message().expect("Couldn't retrieve commit's message"),
                            &new_tree,
                            &parents_refs[..])
                    .expect("Failed to commit");

                old_id_to_new.insert(oid, new_commit_id);
            }
            Err(e) => eprintln!("Error walking the submodule's history: {:?}", e),
        }
    }
}

fn checkout_rewritten_history(repo: &Repository, old_id_to_new: &HashMap<Oid, Oid>) {
    let mut checkoutbuilder = git2::build::CheckoutBuilder::new();
    checkoutbuilder.force();

    let head = repo.head().expect("Couldn't obtain repo's HEAD");
    let head_id = head.target().expect("Couldn't resolve repo's HEAD to a commit ID");
    let updated_id = old_id_to_new[&head_id];
    let object = repo.find_object(updated_id, None)
        .expect("Couldn't look up an object at which HEAD points");
    repo.reset(&object, git2::ResetType::Hard, Some(&mut checkoutbuilder))
        .expect("Couldn't run force-reset");
}
<commit_msg>Drop invalid TODO item<commit_after>extern crate clap;
extern crate git2;

use git2::{Repository, Commit, Oid, Revwalk, Index};
use std::collections::HashMap;

#[macro_use]
mod macros;

const E_SUCCESS: i32 = 0;
const E_NO_GIT_REPO: i32 = 1;

fn main() {
    let exit_code = real_main();
    std::process::exit(exit_code);
}

fn real_main() -> i32 {
    let submodule_dir = parse_cli_arguments();
    println!("Merging {}...", submodule_dir);

    let repo = match Repository::open(".") {
        Ok(repo) => repo,
        Err(e) => {
            eprintln!("Couldn't find Git repo in the current directory: {}",
                      e.message());
            return E_NO_GIT_REPO;
        }
    };

    let mut submodule_revwalk = get_submodule_revwalk(&repo, &submodule_dir);

    let mut old_id_to_new = HashMap::new();

    rewrite_submodule_history(&repo,
                              &mut submodule_revwalk,
                              &mut old_id_to_new,
                              &submodule_dir);

    let mut repo_revwalk = get_repo_revwalk(&repo);
    rewrite_repo_history(&repo, &mut repo_revwalk, &mut old_id_to_new, &submodule_dir);

    checkout_rewritten_history(&repo, &old_id_to_new);

    E_SUCCESS
}

fn parse_cli_arguments() -> String {
    let options = clap::App::new("git-submerge")
                          .version("0.1")
                          .author("Alexander Batischev <eual.jp@gmail.com>")
                          .about("Merge Git submodule into the main repo \
                                  as if they've never been separate at all")
                          .arg(clap::Arg::with_name("SUBMODULE_DIR")
                               .help("The submodule to merge")
                               .required(true)
                               .index(1))
                          .get_matches();

    // We can safely use unwrap() here because the argument is marked as "required" and Clap checks
    // its presence for us.
    String::from( options.value_of("SUBMODULE_DIR").unwrap() )
}

fn get_submodule_revwalk<'repo>(repo: &'repo Repository, submodule_dir: &str) -> Revwalk<'repo> {
    let submodule_url = String::from("./") + submodule_dir;
    let mut remote = repo.remote_anonymous(&submodule_url).expect("Couldn't create an anonymous remote");
    remote.fetch(&[], None, None).expect("Couldn't fetch submodule's history");
    let submodule = repo.find_submodule(submodule_dir)
        .expect("Couldn't find the submodule with expected path");
    let submodule_head = submodule.head_id()
        .expect("Couldn't obtain submodule's HEAD");

    let mut revwalk = repo.revwalk().expect("Couldn't obtain RevWalk object for the repo");
    // "Topological" and reverse means "parents are always visited before their children".
    // We need that in order to be sure that our old-to-new-ids map always contains everything we
    // need it to contain.
    revwalk.set_sorting(git2::SORT_REVERSE | git2::SORT_TOPOLOGICAL);
    // TODO (#6): push all branches and tags, not just HEAD
    revwalk.push(submodule_head).expect("Couldn't add submodule's HEAD to RevWalk list");

    revwalk
}

fn rewrite_submodule_history(repo: &Repository,
                             revwalk: &mut Revwalk,
                             old_id_to_new: &mut HashMap<Oid, Oid>,
                             submodule_dir: &str) {
    for maybe_oid in revwalk {
        match maybe_oid {
            Ok(oid) => {
                let commit = repo.find_commit(oid)
                    .expect(&format!("Couldn't get a commit with ID {}", oid));
                let tree = commit.tree()
                    .expect(&format!("Couldn't obtain the tree of a commit with ID {}", oid));
                let mut old_index = Index::new()
                    .expect("Couldn't create an in-memory index for commit");
                let mut new_index = Index::new().expect("Couldn't create an in-memory index");
                old_index.read_tree(&tree)
                    .expect(&format!("Couldn't read the commit {} into index", oid));

                // Obtain the new tree, where everything from the old one is moved under
                // a directory named after the submodule
                for entry in old_index.iter() {
                    let mut new_entry = entry;

                    // TODO (#11): what mode, owner, mtime etc. does the newly created dir get?
                    let mut new_path = String::from(submodule_dir);
                    new_path += "/";
                    new_path += &String::from_utf8(new_entry.path)
                        .expect("Failed to convert a path to str");

                    new_entry.path = new_path.into_bytes();
                    new_index.add(&new_entry).expect("Couldn't add an entry to the index");
                }
                let tree_id = new_index.write_tree_to(&repo)
                    .expect("Couldn't write the index into a tree");
                old_id_to_new.insert(tree.id(), tree_id);
                let tree = repo.find_tree(tree_id)
                    .expect("Couldn't retrieve the tree we just created");

                let parents = {
                    let mut p: Vec<Commit> = Vec::new();
                    for parent_id in commit.parent_ids() {
                        let new_parent_id = old_id_to_new[&parent_id];
                        let parent = repo.find_commit(new_parent_id)
                            .expect("Couldn't find parent commit by its id");
                        p.push(parent);
                    }
                    p
                };

                let mut parents_refs: Vec<&Commit> = Vec::new();
                for i in 0..parents.len() {
                    parents_refs.push(&parents[i]);
                }
                let new_commit_id = repo.commit(None,
                            &commit.author(),
                            &commit.committer(),
                            &commit.message().expect("Couldn't retrieve commit's message"),
                            &tree,
                            &parents_refs[..])
                    .expect("Failed to commit");

                old_id_to_new.insert(oid, new_commit_id);
            }
            Err(e) => eprintln!("Error walking the submodule's history: {:?}", e),
        }
    }
}

fn get_repo_revwalk<'repo>(repo: &'repo Repository) -> Revwalk<'repo> {
    let mut revwalk = repo.revwalk().expect("Couldn't obtain RevWalk object for the repo");
    revwalk.set_sorting(git2::SORT_REVERSE | git2::SORT_TOPOLOGICAL);
    let head = repo.head().expect("Couldn't obtain repo's HEAD");
    let head_id = head.target().expect("Couldn't resolve repo's HEAD to a commit ID");
    // TODO (#6): push all branches and tags, not just HEAD
    revwalk.push(head_id).expect("Couldn't add repo's HEAD to RevWalk list");

    revwalk
}

fn rewrite_repo_history(repo: &Repository,
                        revwalk: &mut Revwalk,
                        old_id_to_new: &mut HashMap<Oid, Oid>,
                        submodule_dir: &str) {
    let submodule_path = std::path::Path::new(submodule_dir);

    for maybe_oid in revwalk {
        match maybe_oid {
            Ok(oid) => {
                let commit = repo.find_commit(oid)
                    .expect(&format!("Couldn't get a commit with ID {}", oid));
                let tree = commit.tree()
                    .expect(&format!("Couldn't obtain the tree of a commit with ID {}", oid));

                let submodule_subdir = match tree.get_path(submodule_path) {
                    Ok(tree) => tree,
                    Err(e) => {
                        if e.code() == git2::ErrorCode::NotFound &&
                           e.class() == git2::ErrorClass::Tree {
                            // It's okay. The tree lacks the subtree corresponding to the
                            // submodule. In other words, the commit doesn't include the submodule.
                            // That's totally fine. Let's map it into itself and move on.
                            old_id_to_new.insert(oid, oid);
                            continue;
                        } else {
                            // Unexpected error; let's report it and abort the program
                            panic!("Error getting submodule's subdir from the tree: {:?}", e);
                        };
                    }
                };

                // **INVARIANT**: if we got this far, current commit contains a submodule and
                // should be rewritten

                let submodule_commit_id = submodule_subdir.id();
                let new_submodule_commit_id = old_id_to_new[&submodule_commit_id];
                let submodule_commit = repo.find_commit(new_submodule_commit_id)
                    .expect("Couldn't obtain submodule's commit");
                let subtree_id = submodule_commit.tree()
                    .and_then(|t| t.get_path(submodule_path))
                    .and_then(|te| Ok(te.id()))
                    .expect("Couldn't obtain submodule's subtree ID");

                let mut treebuilder = repo.treebuilder(Some(&tree))
                    .expect("Couldn't create TreeBuilder");
                treebuilder.remove(submodule_path)
                    .expect("Couldn't remove submodule path from TreeBuilder");
                treebuilder.insert(submodule_path, subtree_id, 0o040000)
                    .expect("Couldn't add submodule as a subdir to TreeBuilder");
                let new_tree_id = treebuilder.write()
                    .expect("Couldn't write TreeBuilder into a Tree");
                let new_tree = repo.find_tree(new_tree_id)
                    .expect("Couldn't read back the Tree we just wrote");

                // In commits that used to update the submodule, add a parent pointing to
                // appropriate commit in new submodule history
                let mut parent_subtree_ids = std::collections::HashSet::new();
                for parent in commit.parents() {
                    let parent_tree = parent.tree().expect("Couldn't obtain parent's tree");
                    let parent_subdir_tree_id = parent_tree.get_path(submodule_path)
                        .and_then(|x| Ok(x.id()));

                    match parent_subdir_tree_id {
                        Ok(id) => {
                            parent_subtree_ids.insert(id);
                            ()
                        }
                        Err(e) => {
                            if e.code() == git2::ErrorCode::NotFound &&
                               e.class() == git2::ErrorClass::Tree
                            {
                                continue;
                            } else {
                                panic!("Error getting submodule's subdir from the tree: {:?}", e);
                            };
                        }
                    }
                }

                // Here's a few pictures to help you understand how we figure out if current commit
                // updated the submodule. If we draw a DAG and name submodule states, the following
                // situations will mean that the submodule wasn't updated:
                //
                //     o--o--o--A--
                //                 `,-A
                //      o--o--o--B-
                //
                // or
                //
                //     o--o--o--A--
                //                 `,-B
                //      o--o--o--B-
                //
                // And in the following graphs the submodule was updated:
                //
                //     o--o--o--A--
                //                 `,-C
                //      o--o--o--B-
                //
                // or
                //
                //     o--o--o--o--A--B
                //
                // Put into words, the rule will be "the submodule state in current commit is
                // different from states in all its parents". Or, more formally, the current state
                // doesn't belong to the set of states in parents.
                let submodule_updated: bool = !parent_subtree_ids.contains(&submodule_commit_id);

                // Rewrite the parents if the submodule was updated
                let parents = {
                    let mut p: Vec<Commit> = Vec::new();
                    for parent_id in commit.parent_ids() {
                        let actual_parent_id = old_id_to_new[&parent_id];
                        let parent = repo.find_commit(actual_parent_id)
                            .expect("Couldn't find parent commit by its id");
                        p.push(parent);
                    }

                    if submodule_updated {
                        p.push(submodule_commit);
                    }

                    p
                };

                let mut parents_refs: Vec<&Commit> = Vec::new();
                for i in 0..parents.len() {
                    parents_refs.push(&parents[i]);
                }
                let new_commit_id = repo.commit(None,
                            &commit.author(),
                            &commit.committer(),
                            &commit.message().expect("Couldn't retrieve commit's message"),
                            &new_tree,
                            &parents_refs[..])
                    .expect("Failed to commit");

                old_id_to_new.insert(oid, new_commit_id);
            }
            Err(e) => eprintln!("Error walking the submodule's history: {:?}", e),
        }
    }
}

fn checkout_rewritten_history(repo: &Repository, old_id_to_new: &HashMap<Oid, Oid>) {
    let mut checkoutbuilder = git2::build::CheckoutBuilder::new();
    checkoutbuilder.force();

    let head = repo.head().expect("Couldn't obtain repo's HEAD");
    let head_id = head.target().expect("Couldn't resolve repo's HEAD to a commit ID");
    let updated_id = old_id_to_new[&head_id];
    let object = repo.find_object(updated_id, None)
        .expect("Couldn't look up an object at which HEAD points");
    repo.reset(&object, git2::ResetType::Hard, Some(&mut checkoutbuilder))
        .expect("Couldn't run force-reset");
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate qrcode;
extern crate termion;
extern crate image;

use clap::{App, Arg};
use qrcode::{QrCode, EcLevel};
use termion::color;

/*TODO: add arguments for:
- safe zone rendering
- ec level
- size of qr code
- implement different types of qr payloads
*/
fn main() {
    let matches = App::new("qrterm")
        .version("0.1")
        .author("Lukas R. <lukas@bootsmann-games.de>")
        .about("Generates and displays terminal friendly QR-Codes from input strings")
        .arg(Arg::with_name("type")
            .short("t")
            .long("type")
            .value_name("TYPE")
            .help("Sets if a special qr type should be used")
            .takes_value(true))
        .arg(Arg::with_name("INPUT")
            .help("The input string to use")
            .required(true)
            .index(1))
        .get_matches();

    let input = matches.value_of("INPUT").unwrap();

    let code = QrCode::with_error_correction_level(input, EcLevel::H).unwrap();

    let bit_array = code.to_vec();

    let w = code.width();
    let wide = w + 6;

    //Draw the first white safe zone
    for a in 1..(wide * 3) + 1 {
        print!("{}  ", color::Bg(color::LightWhite));
        if a % wide == 0 {
            println!("{}", color::Bg(color::Reset));
        }
    }

    for (i, item) in bit_array.iter().enumerate() {
        //left safe zone
        if i % w == 0 {
            print!("{}      ", color::Bg(color::LightWhite));
        }

        if *item {
            print!("{}  ", color::Bg(color::Black)); //█ ▜
        } else {
            print!("{}  ", color::Bg(color::LightWhite));
        }

        if (i + 1) % w == 0 {
            println!("{}      {}",
                     color::Bg(color::LightWhite),
                     color::Bg(color::Reset));
        }
    }

    //Draw the last white safe zone
    for a in 1..(wide * 3) + 1 {
        print!("{}  ", color::Bg(color::LightWhite));
        if a % wide == 0 {
            println!("{}", color::Bg(color::Reset));
        }
    }
}
<commit_msg>added safe zone switch<commit_after>extern crate clap;
extern crate qrcode;
extern crate termion;
extern crate image;

use clap::{App, Arg};
use qrcode::{QrCode, EcLevel};
use termion::color;

/*TODO: add arguments for:
- ec level
- size of qr code
- implement different types of qr payloads
*/
fn main() {
    let matches = App::new("qrterm")
        .version("0.1")
        .author("Lukas R. <lukas@bootsmann-games.de>")
        .about("Generates and displays terminal friendly QR-Codes from input strings")
        .arg(Arg::with_name("safe_zone")
            .short("s")
            .long("safe-zone")
            .value_name("BOOL")
            .help("Sets wether the safe zone around the code should be drawn or not.")
            .takes_value(true)
            .default_value("true")
        )
        //.arg(Arg::with_name("type")
        //    .short("t")
        //    .long("type")
        //    .value_name("TYPE")
        //    .help("Sets if a special qr type should be used")
        //    .takes_value(true))
        .arg(Arg::with_name("INPUT")
            .help("The input string to use")
            .required(true)
            .index(1))
        .get_matches();

    let input = matches.value_of("INPUT").unwrap();

    let safe: bool = matches.value_of("safe_zone").unwrap() == "true";

    let code = QrCode::with_error_correction_level(input, EcLevel::H).unwrap();

    let bit_array = code.to_vec();

    let w = code.width();
    let wide = w + 6;

    //Draw the first white safe zone
    if safe {
        for a in 1..(wide * 3) + 1 {
            print!("{}  ", color::Bg(color::LightWhite));
            if a % wide == 0 {
                println!("{}", color::Bg(color::Reset));
            }
        }
    }

    for (i, item) in bit_array.iter().enumerate() {
        //left safe zone
        if safe {
            if i % w == 0 {
                print!("{}      ", color::Bg(color::LightWhite));
            }
        }

        if *item {
            print!("{}  ", color::Bg(color::Black)); //█ ▜
        } else {
            print!("{}  ", color::Bg(color::LightWhite));
        }

        if (i + 1) % w == 0 {
            if safe {
                //draw right safe zone
                println!("{}      {}",
                         color::Bg(color::LightWhite),
                         color::Bg(color::Reset));
            } else {
                println!("{}", color::Bg(color::Reset));
            }
        }
    }

    //Draw the last white safe zone
    if safe {
        for a in 1..(wide * 3) + 1 {
            print!("{}  ", color::Bg(color::LightWhite));
            if a % wide == 0 {
                println!("{}", color::Bg(color::Reset));
            }
        }
    }
}
<|endoftext|>


<commit_before>//TODO(remove)
#![allow(dead_code)]

use std::io::{File, BufWriter, Reader};
use error::Ch8Error;

const RAM_SIZE: usize = 4096;
const PROGRAM_START: usize = 0x200;

mod error;

struct Vm {
    reg: [u8; 16],
    index: u16,
    pc: u16,
    sp: u8,
    stack: [u16; 256],
    ram: [u8; RAM_SIZE],

    timer: u8,
    tone: u8,

    screen: [u8; 64 * 32],
    keys: [u8; 16],
}

impl Vm {
    fn new() -> Vm {
        Vm {
            reg: [0; 16],
            index: 0,
            pc: PROGRAM_START as u16,
            sp: 0,
            stack: [0; 256],
            ram: [0; RAM_SIZE],

            timer: 0,
            tone: 0,

            screen: [0; 64 * 32],
            keys: [0; 16],
        }
    }

    fn load_rom(&mut self, reader: &mut Reader) -> Result<usize, Ch8Error> {
        let rom = try!(reader.read_to_end());
        if rom.len() > (RAM_SIZE - PROGRAM_START) {
           return Err(Ch8Error::Io("Rom was larger than available RAM", None))
        }
        let mut ram = BufWriter::new(&mut self.ram[PROGRAM_START..RAM_SIZE]);
        try!(ram.write(rom.as_slice()));
        return Ok(rom.len());
    }

    fn dump_ram(&self, writer: &mut Writer) {
        writer.write(&self.ram);
    }
}

struct Op {
    raw: u16
}

impl Op {
    fn addr(&self) -> u16 {
        self.raw & 0x0FFF
    }

    fn x(&self) -> u8 {
        ((self.raw & 0x0F00) >> 8) as u8
    }

    fn y(&self) -> u8 {
        ((self.raw & 0x00F0) >> 4) as u8
    }

    fn n(&self) -> u8 {
        ((self.raw & 0xF000) >> 12) as u8
    }

    fn kk(&self) -> u8 {
        (self.raw & 0x00FF) as u8
    }
}

fn main() {
    let mut vm = Vm::new();

    let mut rom_file = File::open(&Path::new("/Users/jakerr/Downloads/IBM Logo.ch8")).unwrap();
    match vm.load_rom(&mut rom_file) {
        Ok(size) => println!("Loaded rom size: {}", size),
        Err(e) => println!("Error loading rom: {}", e)
    }

    let mut dump_file = File::create(&Path::new("/Users/jakerr/tmp/dump.ch8ram")).unwrap();
    vm.dump_ram(&mut dump_file);

    let op = Op{raw: 0x8cd1};
    println!("addr: 0x{:X}", op.addr());
    println!("x: 0x{:X}", op.x());
    println!("y: 0x{:X}", op.y());
    println!("n: 0x{:X}", op.n());
    println!("kk: 0x{:X}", op.kk());
}
<commit_msg>Print ops from ram<commit_after>//TODO(remove)
#![allow(dead_code)]

use std::io::{File, BufWriter, Reader};
use error::Ch8Error;

const RAM_SIZE: usize = 4096;
const PROGRAM_START: usize = 0x200;

mod error;

struct Vm {
    reg: [u8; 16],
    index: u16,
    pc: u16,
    sp: u8,
    stack: [u16; 256],
    ram: [u8; RAM_SIZE],

    timer: u8,
    tone: u8,

    screen: [u8; 64 * 32],
    keys: [u8; 16],
}

impl Vm {
    fn new() -> Vm {
        Vm {
            reg: [0; 16],
            index: 0,
            pc: PROGRAM_START as u16,
            sp: 0,
            stack: [0; 256],
            ram: [0; RAM_SIZE],

            timer: 0,
            tone: 0,

            screen: [0; 64 * 32],
            keys: [0; 16],
        }
    }

    fn load_rom(&mut self, reader: &mut Reader) -> Result<usize, Ch8Error> {
        let rom = try!(reader.read_to_end());
        if rom.len() > (RAM_SIZE - PROGRAM_START) {
           return Err(Ch8Error::Io("Rom was larger than available RAM", None))
        }
        let mut ram = BufWriter::new(&mut self.ram[PROGRAM_START..RAM_SIZE]);
        try!(ram.write(rom.as_slice()));
        return Ok(rom.len());
    }

    fn dump_ram(&self, writer: &mut Writer) {
        writer.write(&self.ram);
    }
}

struct Op {
    raw: u16
}

impl Op {
    fn addr(&self) -> u16 {
        self.raw & 0x0FFF
    }

    fn x(&self) -> u8 {
        ((self.raw & 0x0F00) >> 8) as u8
    }

    fn y(&self) -> u8 {
        ((self.raw & 0x00F0) >> 4) as u8
    }

    fn n(&self) -> u8 {
        ((self.raw & 0xF000) >> 12) as u8
    }

    fn kk(&self) -> u8 {
        (self.raw & 0x00FF) as u8
    }
}

fn main() {
    let mut vm = Vm::new();

    let mut rom_file = File::open(&Path::new("/Users/jakerr/Downloads/IBM Logo.ch8")).unwrap();
    match vm.load_rom(&mut rom_file) {
        Ok(size) => println!("Loaded rom size: {}", size),
        Err(e) => println!("Error loading rom: {}", e)
    }

    let mut dump_file = File::create(&Path::new("/Users/jakerr/tmp/dump.ch8ram")).unwrap();
    vm.dump_ram(&mut dump_file);

    for i in vm.ram.chunks(2) {
        match i {
            [0, 0] => continue,
            [h, l] => {
                let op = Op{raw:((h as u16) << 8) | l as u16};
                println!("raw: 0x{:X}", op.raw);
                println!("addr: 0x{:X}", op.addr());
                println!("x: 0x{:X}", op.x());
                println!("y: 0x{:X}", op.y());
                println!("n: 0x{:X}", op.n());
                println!("kk: 0x{:X}\n", op.kk());
            },
            _ => continue
        }
    }
}
<|endoftext|>


<commit_before>#![feature(path_ext)]
#![feature(collections)]

#[macro_use]
extern crate log;
extern crate git2;
extern crate iron;
extern crate toml;

use std::env;
use std::io::Read;
use std::fs::{self, PathExt, File};
use std::path::PathBuf;
use std::error::Error;

use git2::Repository;

use iron::prelude::*;

/// Contains various logger implementations for Grist
mod loggers;
/// Contains Iron server functions
mod servers;
/// Contains code related to configuration file parsing
mod config;

fn main() {

    // try to parse the config file at ~/.gristconfig.toml; otherwise, use defaults
    let configs = File::open("~/.gristconfig.toml")
        .and_then(|mut f: File| {
            let mut s = String::new();
            try!(f.read_to_string(&mut s));
            Ok(s) })
        .map_err(|err| String::from_str(err.description()) )
        .and_then(|s: String| config::parse_toml(s.as_ref()))
        .unwrap_or(config::DEFAULTS.clone());

    // set log level from configs
    let _ = if configs.debug {
        log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Debug);
            Box::new(loggers::DebugLogger)
        });
    } else {
        log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Info);
            Box::new(loggers::DefaultLogger)
        });
    };

    let roots = configs.roots.unwrap_or(
        // get the current working dir, assume that we are in a valid directory
        vec![env::current_dir().unwrap()]
    );

    info!("Starting Grist in {:?}", roots);

    let repos: Vec<Repository> = roots.iter().flat_map(|ref root: &PathBuf| {
        fs::read_dir(root) // walk the root dir
            .unwrap()
            .filter_map( |entry| match entry {
                Ok(ref dir) if dir.path().is_dir() => { // found a dir
                    debug!("Found directory: {:?}", dir.path());
                    match Repository::init(dir.path()) { // attempt to open dir as git repo
                        Ok(repo) => { info!("Found repository: {:?}", dir.path()); Some(repo) },
                        Err(why) => { warn!("Failed to open {:?}: {}", dir.path(), why); None }
                    }
                },
                Err(why) => { warn!("Could not read entry: {}", why); None },
                Ok(ref e)=> { debug!("{:?} is not a directory.", e.path()); None }
            })
    }).collect();

    info!("Discovered {} repositories.", repos.len());

    // serve sample hello world page for now
    Iron::new(servers::hello_world).http("localhost:3000").unwrap();

}
<commit_msg>Add documentation to main<commit_after>#![feature(path_ext)]
#![feature(collections)]

#[macro_use]
extern crate log;
extern crate git2;
extern crate iron;
extern crate toml;

use std::env;
use std::io::Read;
use std::fs::{self, PathExt, File};
use std::path::PathBuf;
use std::error::Error;

use git2::Repository;

use iron::prelude::*;

/// Contains various logger implementations for Grist.
///
/// TODO: log to a file in addition to the console
/// TODO: log to /var/log/
pub mod loggers;

/// Contains Iron server functions.
pub mod servers;

/// Contains code related to setting Grist's configuration.
///
/// Currently, configurations are set from a TOML configuration
/// file stored in `~/.gristconfig.toml`. Eventually, functionality
/// will be added to also allow configurations to be set from the
/// command-line, falling back to `~/.gristconfig.toml` and finally
/// to the defaults.
///
/// TODO: add command-line argument parsing
/// TODO: allow log file to be set from configs
pub mod config;

fn main() {

    // try to parse the config file at ~/.gristconfig.toml; otherwise, use defaults
    let configs = File::open("~/.gristconfig.toml")
        .and_then(|mut f: File| {
            let mut s = String::new();
            try!(f.read_to_string(&mut s));
            Ok(s) })
        .map_err(|err| String::from_str(err.description()) )
        .and_then(|s: String| config::parse_toml(s.as_ref()))
        .unwrap_or(config::DEFAULTS.clone());

    // set log level from configs
    let _ = if configs.debug {
        log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Debug);
            Box::new(loggers::DebugLogger)
        });
    } else {
        log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Info);
            Box::new(loggers::DefaultLogger)
        });
    };

    let roots = configs.roots.unwrap_or(
        // get the current working dir, assume that we are in a valid directory
        vec![env::current_dir().unwrap()]
    );

    info!("Starting Grist in {:?}", roots);

    let repos: Vec<Repository> = roots.iter().flat_map(|ref root: &PathBuf| {
        fs::read_dir(root) // walk the root dir
            .unwrap()
            .filter_map( |entry| match entry {
                Ok(ref dir) if dir.path().is_dir() => { // found a dir
                    debug!("Found directory: {:?}", dir.path());
                    match Repository::init(dir.path()) { // attempt to open dir as git repo
                        Ok(repo) => { info!("Found repository: {:?}", dir.path()); Some(repo) },
                        Err(why) => { warn!("Failed to open {:?}: {}", dir.path(), why); None }
                    }
                },
                Err(why) => { warn!("Could not read entry: {}", why); None },
                Ok(ref e)=> { debug!("{:?} is not a directory.", e.path()); None }
            })
    }).collect();

    info!("Discovered {} repositories.", repos.len());

    // serve sample hello world page for now
    Iron::new(servers::hello_world).http("localhost:3000").unwrap();

}
<|endoftext|>


<commit_before>use std::io;

enum Mode {
    Command,
    Input,
}

fn command_line(line: &String, mode: &mut Mode) {
    match line.trim() {
        "a" => *mode = Mode::Input,
        "q" => std::process::exit(0),
        _ => println!("Command: {}", line),
    }
}

fn input_line(line: &String, mode: &mut Mode) {
    match line.trim() {
        "." => *mode = Mode::Command,
        _ => println!("Input: {}", line),
    }
}

fn main() {
    let mut mode = Mode::Command;
    loop {
        let mut line = String::new();
        match io::stdin().read_line(&mut line) {
            Ok(_) => (),
            Err(_) => {
                println!("?");
                continue
            },
        }
        match mode {
            Mode::Command => command_line(&line, &mut mode),
            Mode::Input => input_line(&line, &mut mode),
        }
    }
}
<commit_msg>Add Range type<commit_after>use std::io;

enum Mode {
    Command,
    Input,
}

struct Range {
    start: i32,
    end: i32,
}

fn command_line(line: &String, mode: &mut Mode) {
    match line.trim() {
        "a" => *mode = Mode::Input,
        "q" => std::process::exit(0),
        _ => println!("Command: {}", line),
    }
}

fn input_line(line: &String, mode: &mut Mode) {
    match line.trim() {
        "." => *mode = Mode::Command,
        _ => println!("Input: {}", line),
    }
}

fn append_after(lines: Vec<line>, range: Range, diff: &Diff) {
    diff.add_lines(range.start, lines);
}

fn main() {
    let mut mode = Mode::Command;
    loop {
        let mut line = String::new();
        match io::stdin().read_line(&mut line) {
            Ok(_) => (),
            Err(_) => {
                println!("?");
                continue
            },
        }
        match mode {
            Mode::Command => command_line(&line, &mut mode),
            Mode::Input => input_line(&line, &mut mode),
        }
    }
}
<|endoftext|>


<commit_before>extern crate argparse;

use std::fs::File;
use argparse::{ArgumentParser, Store, StoreTrue};
use satyrs::cnf::CNF;
use satyrs::dpll;

// extern crate satyrs;
mod satyrs;

fn main() {
    let mut filename = String::new();
    let mut verbose: bool = false;
    {
        // this block limits scope of borrows by ap.refer() method
        let mut ap = ArgumentParser::new();
        ap.set_description("Satyrs: A lustful, drunken SAT solver");
        ap.refer(&mut filename)
          .add_argument("filename", Store, "SAT filename to read")
          .required();
        ap.refer(&mut verbose)
          .add_option(&["-v", "--verbose"], StoreTrue, "Be verbose");
        ap.parse_args_or_exit();
    }

    // Read the file
    let f: File = File::open(filename).expect("Could not open file");

    // TODO: This is definitely not the correct way to handle errors
    let cnf: CNF = satyrs::cnf::parse_dimacs_file(f).expect("Dimacs Error");
    let solvable = dpll::DPLL(&cnf, verbose);
    match solvable {
        Some((solution, partial)) => {
            let solution_str = satyrs::cnf::format_output(&solution);
            println!("Satisfiable!");
            println!("SAT Format: {}\nSolution: {:?}\nPartial: {}\n",  solution_str, solution, partial);
        }
        None => println!("Unsatisfiable"),
    }
}
<commit_msg>format update. Minor<commit_after>extern crate argparse;

use std::fs::File;
use argparse::{ArgumentParser, Store, StoreTrue};
use satyrs::cnf::CNF;
use satyrs::dpll;

// extern crate satyrs;
mod satyrs;

fn main() {
    let mut filename = String::new();
    let mut verbose: bool = false;
    {
        // this block limits scope of borrows by ap.refer() method
        let mut ap = ArgumentParser::new();
        ap.set_description("Satyrs: A lustful, drunken SAT solver");
        ap.refer(&mut filename)
          .add_argument("filename", Store, "SAT filename to read")
          .required();
        ap.refer(&mut verbose)
          .add_option(&["-v", "--verbose"], StoreTrue, "Be verbose");
        ap.parse_args_or_exit();
    }

    // Read the file
    let f: File = File::open(filename).expect("Could not open file");

    // TODO: This is definitely not the correct way to handle errors
    let cnf: CNF = satyrs::cnf::parse_dimacs_file(f).expect("Dimacs Error");
    let solvable = dpll::DPLL(&cnf, verbose);
    match solvable {
        Some((solution, partial)) => {
            let solution_str = satyrs::cnf::format_output(&solution);
            println!("Satisfiable!");
            println!("SAT Format: {}\nSolution: {:?}\nPartial {}\n",  solution_str, solution, partial);
        }
        None => println!("Unsatisfiable"),
    }
}
<|endoftext|>


<commit_before>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
extern crate clap;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame, Server as WsServer};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        ::std::io::Error, Io;
        log::SetLoggerError, Log;
        ::url::ParseError, Url;
        ::websocket::result::WebSocketError, Ws;
        ::std::env::VarError, Ev;
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::Builder::new().spawn(move || -> Result<()> {
            // Actual data transfer happens here
            ::std::io::copy(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        })?;

        // Actual data transfer happens here
        ::std::io::copy(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let request = Client::connect(url)?;

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_tcp_peer(addr: &str) -> Result<
        Peer<
            ::std::net::TcpStream,
            ::std::net::TcpStream,
        >> {
    let sock = ::std::net::TcpStream::connect(addr)?;

    let peer = Peer {
        reader : sock.try_clone()?,
        writer : sock.try_clone()?,
    };
    info!("Connected to TCP {}", addr);
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

struct TcpServer(::std::net::TcpListener);

impl TcpServer {
    fn new(addr: &str) -> Result<Self> {
        Ok(TcpServer(::std::net::TcpListener::bind(addr)?))
    }
}

impl Server for TcpServer {
    fn accept_client(&mut self) -> Result<IPeer> {
        let (sock, addr) = self.0.accept()?;
        info!("TCP client connection from {}", addr);
        let peer = Peer {
            reader : sock.try_clone()?,
            writer : sock.try_clone()?,
        };
        Ok(peer.upcast())
    }
}

struct WebsockServer<'a>(WsServer<'a>);

impl<'a> WebsockServer<'a> {
    fn new(addr: &str) -> Result<Self> {
        Ok(WebsockServer(WsServer::bind(addr)?))
    }
}

impl<'a> Server for WebsockServer<'a> {
    fn accept_client(&mut self) -> Result<IPeer> {
        let connection = self.0.accept()?;
        info!("WebSocket client connection ...");
        let request = connection.read_request()?;
        request.validate()?;
        let response = request.accept(); // Form a response
        let mut client = response.send()?; // Send the response

        let ip = client.get_mut_sender()
            .get_mut()
            .peer_addr()
            .unwrap();

        info!("... from IP {}", ip);

        let (sender, receiver) = client.split();

        let peer = Peer {
            reader : ReceiverWrapper(receiver),
            writer : SenderWrapper(sender),
        };
        Ok(peer.upcast())
    }
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

trait Server
{
    fn accept_client(&mut self) -> Result<IPeer>;

    fn start_serving(&mut self, spec2: &str, once: bool) -> Result<()> {
        let spec2s = spec2.to_string();
        let closure = move |peer, spec2 : String|{
            let spec2_ = get_peer_by_spec(spec2.as_str())?;
            let peer2 = match spec2_ {
                Spec::Server(mut x) => {
                    x.accept_client()?
                }
                Spec::Client(p1) => {
                    p1
                }
            };
            let des = DataExchangeSession {
                peer1 : peer,
                peer2 : peer2,
            };

            des.data_exchange()
        };
        if once {
            let peer = self.accept_client()?;
            closure(peer, spec2s)
        } else {
            let cl2 = ::std::sync::Arc::new(closure);
            loop {
                let ret = self.accept_client();
                let peer = match ret {
                    Ok(x) => x,
                    Err(er) => {
                        warn!("Can't accept client: {}", er);
                        continue;
                    }
                };
                let cl3 = cl2.clone();
                let spec2s2 = spec2s.clone();
                if let Err(x) = thread::Builder::new().spawn(move|| {
                    if let Err(x) = cl3(peer, spec2s2) {
                        warn!("Error while serving: {}", x);
                    }
                }) {
                    warn!("Error creating thread: {}", x);
                    thread::sleep(::std::time::Duration::from_millis(200));
                }
            }
        }
    }

    fn upcast(self) -> Box<Server+Send>
        where Self : Sized + Send + 'static
        { Box::new(self) as Box<Server+Send> }
}

fn main2(spec1: &str, spec2: &str, once: bool) -> Result<()> {
    let spec1_ = get_peer_by_spec(spec1)?;

    match spec1_ {
        Spec::Server(mut x) => {
            x.start_serving(spec2, once)
        }
        Spec::Client(p1) => {
            let spec2_ = get_peer_by_spec(spec2)?;

            let otherpeer = match spec2_ {
                Spec::Server(mut x) => {
                    let t = x.accept_client()?;
                    t
                }
                Spec::Client(p2) => {
                    p2
                }
            };

            let des = DataExchangeSession {
                peer1 : p1,
                peer2 : otherpeer,
            };

            des.data_exchange()
        }
    }
}

enum Spec {
    Server(Box<Server + Send>),
    Client(IPeer)
}

fn get_peer_by_spec(specifier: &str) -> Result<Spec> {
    use Spec::{Server,Client};
    match specifier {
        x if x == "-"               => Ok(Client(get_stdio_peer()?.upcast())),
        x if x.starts_with("ws:")   => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("wss:")  => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("tcp:")  => Ok(Client(get_tcp_peer(&x[4..])?.upcast())),
        x if x.starts_with("l-tcp:")  => Ok(Server(TcpServer::new(&x[6..])?.upcast())),
        x if x.starts_with("l-ws:")  => Ok(Server(WebsockServer::new(&x[5..])?.upcast())),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("websocat")
        .version("0.1")
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Exchange binary data between binary websocket and something.\nSocat analogue with websockets.")
        .arg(::clap::Arg::with_name("spec1")
             .help("First specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("spec2")
             .help("Second specifier.")
             .required(true)
             .index(2))
        .after_help(r#"
Specifiers can be:
  ws[s]://<rest of websocket URL>   Connect to websocket
  l-ws:host:port                    Listen unencrypted websocket
  tcp:host:port                     Connect to TCP
  l-tcp:host:port                   Listen TCP connections
  -                                 stdin/stdout
  (more to be implemented)

Examples:
  websocat l-tcp:0.0.0.0:9559 ws://echo.websocket.org/
    Listen port 9959 on address :: and forward
    all connections to a public loopback websocket
  websocat l-ws:127.0.0.1:7878 tcp:127.0.0.1:1194
    Listen websocket and forward connections to local tcp
    Use nginx proxy for SSL if you want
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to a secure web socket.
    Like netcat, but for websocket.
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.

Specify listening part first, unless you want websocat to serve once.

IPv6 supported, just use specs like `l-ws:::1:4567`

Web socket usage is not obligatory, you any specs on both sides.
"#)
        .get_matches();

    let spec1  = matches.value_of("spec1") .ok_or("no listener_spec" )?;
    let spec2 = matches.value_of("spec2").ok_or("no connector_spec")?;

    main2(spec1, spec2, false)?;

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<commit_msg>Use modified std::io::copy<commit_after>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
extern crate clap;

const BUFSIZ : usize = 8192;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame, Server as WsServer};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        ::std::io::Error, Io;
        log::SetLoggerError, Log;
        ::url::ParseError, Url;
        ::websocket::result::WebSocketError, Ws;
        ::std::env::VarError, Ev;
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

// Derived from https://doc.rust-lang.org/src/std/up/src/libstd/io/util.rs.html#46-61
pub fn copy_with_flushes<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> ::std::io::Result<u64>
    where R: Read, W: Write
{
    let mut buf = [0; BUFSIZ];
    let mut written = 0;
    loop {
        let len = match reader.read(&mut buf) {
            Ok(0) => return Ok(written),
            Ok(len) => len,
            Err(ref e) if e.kind() == IoErrorKind::Interrupted => continue,
            Err(ref e) if e.kind() == IoErrorKind::WouldBlock => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(&buf[..len])?;
        writer.flush()?;
        written += len as u64;
    }
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::Builder::new().spawn(move || -> Result<()> {
            // Actual data transfer happens here
            copy_with_flushes(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        })?;

        // Actual data transfer happens here
        copy_with_flushes(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let request = Client::connect(url)?;

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_tcp_peer(addr: &str) -> Result<
        Peer<
            ::std::net::TcpStream,
            ::std::net::TcpStream,
        >> {
    let sock = ::std::net::TcpStream::connect(addr)?;

    let peer = Peer {
        reader : sock.try_clone()?,
        writer : sock.try_clone()?,
    };
    info!("Connected to TCP {}", addr);
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

struct TcpServer(::std::net::TcpListener);

impl TcpServer {
    fn new(addr: &str) -> Result<Self> {
        Ok(TcpServer(::std::net::TcpListener::bind(addr)?))
    }
}

impl Server for TcpServer {
    fn accept_client(&mut self) -> Result<IPeer> {
        let (sock, addr) = self.0.accept()?;
        info!("TCP client connection from {}", addr);
        let peer = Peer {
            reader : sock.try_clone()?,
            writer : sock.try_clone()?,
        };
        Ok(peer.upcast())
    }
}

struct WebsockServer<'a>(WsServer<'a>);

impl<'a> WebsockServer<'a> {
    fn new(addr: &str) -> Result<Self> {
        Ok(WebsockServer(WsServer::bind(addr)?))
    }
}

impl<'a> Server for WebsockServer<'a> {
    fn accept_client(&mut self) -> Result<IPeer> {
        let connection = self.0.accept()?;
        info!("WebSocket client connection ...");
        let request = connection.read_request()?;
        request.validate()?;
        let response = request.accept(); // Form a response
        let mut client = response.send()?; // Send the response

        let ip = client.get_mut_sender()
            .get_mut()
            .peer_addr()
            .unwrap();

        info!("... from IP {}", ip);

        let (sender, receiver) = client.split();

        let peer = Peer {
            reader : ReceiverWrapper(receiver),
            writer : SenderWrapper(sender),
        };
        Ok(peer.upcast())
    }
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

trait Server
{
    fn accept_client(&mut self) -> Result<IPeer>;

    fn start_serving(&mut self, spec2: &str, once: bool) -> Result<()> {
        let spec2s = spec2.to_string();
        let closure = move |peer, spec2 : String|{
            let spec2_ = get_peer_by_spec(spec2.as_str())?;
            let peer2 = match spec2_ {
                Spec::Server(mut x) => {
                    x.accept_client()?
                }
                Spec::Client(p1) => {
                    p1
                }
            };
            let des = DataExchangeSession {
                peer1 : peer,
                peer2 : peer2,
            };

            des.data_exchange()
        };
        if once {
            let peer = self.accept_client()?;
            closure(peer, spec2s)
        } else {
            let cl2 = ::std::sync::Arc::new(closure);
            loop {
                let ret = self.accept_client();
                let peer = match ret {
                    Ok(x) => x,
                    Err(er) => {
                        warn!("Can't accept client: {}", er);
                        continue;
                    }
                };
                let cl3 = cl2.clone();
                let spec2s2 = spec2s.clone();
                if let Err(x) = thread::Builder::new().spawn(move|| {
                    if let Err(x) = cl3(peer, spec2s2) {
                        warn!("Error while serving: {}", x);
                    }
                }) {
                    warn!("Error creating thread: {}", x);
                    thread::sleep(::std::time::Duration::from_millis(200));
                }
            }
        }
    }

    fn upcast(self) -> Box<Server+Send>
        where Self : Sized + Send + 'static
        { Box::new(self) as Box<Server+Send> }
}

fn main2(spec1: &str, spec2: &str, once: bool) -> Result<()> {
    let spec1_ = get_peer_by_spec(spec1)?;

    match spec1_ {
        Spec::Server(mut x) => {
            x.start_serving(spec2, once)
        }
        Spec::Client(p1) => {
            let spec2_ = get_peer_by_spec(spec2)?;

            let otherpeer = match spec2_ {
                Spec::Server(mut x) => {
                    let t = x.accept_client()?;
                    t
                }
                Spec::Client(p2) => {
                    p2
                }
            };

            let des = DataExchangeSession {
                peer1 : p1,
                peer2 : otherpeer,
            };

            des.data_exchange()
        }
    }
}

enum Spec {
    Server(Box<Server + Send>),
    Client(IPeer)
}

fn get_peer_by_spec(specifier: &str) -> Result<Spec> {
    use Spec::{Server,Client};
    match specifier {
        x if x == "-"               => Ok(Client(get_stdio_peer()?.upcast())),
        x if x.starts_with("ws:")   => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("wss:")  => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("tcp:")  => Ok(Client(get_tcp_peer(&x[4..])?.upcast())),
        x if x.starts_with("l-tcp:")  => Ok(Server(TcpServer::new(&x[6..])?.upcast())),
        x if x.starts_with("l-ws:")  => Ok(Server(WebsockServer::new(&x[5..])?.upcast())),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("websocat")
        .version("0.1")
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Exchange binary data between binary websocket and something.\nSocat analogue with websockets.")
        .arg(::clap::Arg::with_name("spec1")
             .help("First specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("spec2")
             .help("Second specifier.")
             .required(true)
             .index(2))
        .after_help(r#"
Specifiers can be:
  ws[s]://<rest of websocket URL>   Connect to websocket
  l-ws:host:port                    Listen unencrypted websocket
  tcp:host:port                     Connect to TCP
  l-tcp:host:port                   Listen TCP connections
  -                                 stdin/stdout
  (more to be implemented)

Examples:
  websocat l-tcp:0.0.0.0:9559 ws://echo.websocket.org/
    Listen port 9959 on address :: and forward
    all connections to a public loopback websocket
  websocat l-ws:127.0.0.1:7878 tcp:127.0.0.1:1194
    Listen websocket and forward connections to local tcp
    Use nginx proxy for SSL if you want
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to a secure web socket.
    Like netcat, but for websocket.
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.

Specify listening part first, unless you want websocat to serve once.

IPv6 supported, just use specs like `l-ws:::1:4567`

Web socket usage is not obligatory, you any specs on both sides.
"#)
        .get_matches();

    let spec1  = matches.value_of("spec1") .ok_or("no listener_spec" )?;
    let spec2 = matches.value_of("spec2").ok_or("no connector_spec")?;

    main2(spec1, spec2, false)?;

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;

mod args;
mod macros;

use crate::args::AppArgs;
use monolith::html::{html_to_dom, stringify_document, walk_and_embed_assets};
use monolith::http::retrieve_asset;
use monolith::utils::is_valid_url;
use reqwest::blocking::Client;
use reqwest::header::{HeaderMap, HeaderValue, USER_AGENT};
use std::collections::HashMap;
use std::fs::{remove_file, File};
use std::io::{Error, Write};
use std::time::Duration;

fn create_file(file_path: &String, content: String) -> Result<(), Error> {
    let file = File::create(file_path.as_str());

    let mut file = match file {
        Ok(file) => file,
        Err(error) => return Err(error),
    };

    if content != str!() {
        file.write_all(content.as_bytes())?;
        file.write_all("\n".as_bytes())?;
        file.sync_all()?;
    } else {
        // Remove the file right away if it had no content
        remove_file(file_path.as_str())?;
    }

    Ok(())
}

fn main() {
    let app_args = AppArgs::get();
    let cache = &mut HashMap::new();

    // Attempt to create output file
    if app_args.output != str!() {
        create_file(&app_args.output, str!()).unwrap();
    }

    if is_valid_url(app_args.url_target.as_str()) {
        // Initialize client
        let mut header_map = HeaderMap::new();
        match HeaderValue::from_str(&app_args.user_agent) {
            Ok(header) => header_map.insert(USER_AGENT, header),
            Err(err) => {
                eprintln!("Invalid user agent! {}", err);
                return;
            }
        };
        let client = Client::builder()
            .timeout(Duration::from_secs(10))
            .danger_accept_invalid_certs(app_args.insecure)
            .default_headers(header_map)
            .build()
            .expect("Failed to initialize HTTP client");

        // Retrieve root document
        let (data, final_url) = retrieve_asset(
            cache,
            &client,
            app_args.url_target.as_str(),
            false,
            "",
            app_args.silent,
        )
        .unwrap();
        let dom = html_to_dom(&data);

        walk_and_embed_assets(
            cache,
            &client,
            &final_url,
            &dom.document,
            app_args.no_css,
            app_args.no_js,
            app_args.no_images,
            app_args.silent,
            app_args.no_frames,
        );

        let html: String = stringify_document(
            &dom.document,
            app_args.no_css,
            app_args.no_frames,
            app_args.no_js,
            app_args.no_images,
            app_args.isolate,
        );

        if app_args.output == str!() {
            println!("{}", html);
        } else {
            create_file(&app_args.output, html).unwrap();
        }
    }
}
<commit_msg>refactor main to address several issues<commit_after>#[macro_use]
extern crate clap;

mod args;
mod macros;

use crate::args::AppArgs;
use monolith::html::{html_to_dom, stringify_document, walk_and_embed_assets};
use monolith::http::retrieve_asset;
use monolith::utils::is_valid_url;
use reqwest::blocking::Client;
use reqwest::header::{HeaderMap, HeaderValue, USER_AGENT};
use std::collections::HashMap;
use std::fs::File;
use std::io::{self, Error, Write};
use std::process;
use std::time::Duration;

enum Output {
    Stdout(io::Stdout),
    File(File),
}

impl Output {
    fn new(file_path: &str) -> Result<Output, Error> {
        if file_path.is_empty() {
            Ok(Output::Stdout(io::stdout()))
        } else {
            Ok(Output::File(File::create(file_path)?))
        }
    }

    fn writeln_str(&mut self, s: &str) -> Result<(), Error> {
        match self {
            Output::Stdout(stdout) => {
                writeln!(stdout, "{}", s)?;
                stdout.flush()
            }
            Output::File(f) => {
                writeln!(f, "{}", s)?;
                f.flush()
            }
        }
    }
}

fn main() {
    let app_args = AppArgs::get();

    if !is_valid_url(app_args.url_target.as_str()) {
        eprintln!(
            "Only HTTP and HTTPS URLs are allowed but got: {}",
            &app_args.url_target
        );
        process::exit(1);
    }

    let mut output = Output::new(&app_args.output).expect("Could not prepare output");

    // Initialize client
    let mut cache = HashMap::new();
    let mut header_map = HeaderMap::new();
    header_map.insert(
        USER_AGENT,
        HeaderValue::from_str(&app_args.user_agent).expect("Invalid User-Agent header specified"),
    );

    let client = Client::builder()
        .timeout(Duration::from_secs(10))
        .danger_accept_invalid_certs(app_args.insecure)
        .default_headers(header_map)
        .build()
        .expect("Failed to initialize HTTP client");

    // Retrieve root document
    let (data, final_url) = retrieve_asset(
        &mut cache,
        &client,
        app_args.url_target.as_str(),
        false,
        "",
        app_args.silent,
    )
    .expect("Could not retrieve assets in HTML");
    let dom = html_to_dom(&data);

    walk_and_embed_assets(
        &mut cache,
        &client,
        &final_url,
        &dom.document,
        app_args.no_css,
        app_args.no_js,
        app_args.no_images,
        app_args.silent,
        app_args.no_frames,
    );

    let html: String = stringify_document(
        &dom.document,
        app_args.no_css,
        app_args.no_frames,
        app_args.no_js,
        app_args.no_images,
        app_args.isolate,
    );

    output
        .writeln_str(&html)
        .expect("Could not write HTML output");
}
<|endoftext|>


<commit_before>use std::io::Write;
use std::io::Read;
use std::net::TcpListener;
use std::thread;
use std::sync::{Arc, Mutex};

mod server;
mod broker;
mod message;

struct TcpClient<'a> {
    stream: &'a mut std::net::TcpStream,
}

impl<'a> TcpClient<'a> {
    fn new(stream: &'a mut std::net::TcpStream) -> Self {
        TcpClient { stream : stream }
    }
}

impl<'a> server::Client for TcpClient<'a> {
    fn send(&mut self, bytes: &[u8]) {
        let _ = self.stream.write(bytes);
    }
}

fn main() {
    let server = Arc::new(Mutex::new(server::Server::new()));
    let listener = TcpListener::bind(("0.0.0.0", 1883)).unwrap();

    for byte_stream in listener.incoming() {
        let server = server.clone();
        let mut mqtt_stream = server::Stream::new();

        thread::spawn(move || {
            let mut byte_stream = byte_stream.unwrap();
            loop {
                let read_result = byte_stream.read(mqtt_stream.buffer());
                let mut client = TcpClient::new(&mut byte_stream);

                match read_result {
                    Ok(length) => {
                        let mut server = server.lock().unwrap();
                        mqtt_stream.handle_messages(length, &mut server, &mut client);
                    },
                    _ => panic!("Error reading bytes from stream"),
                }
            }
        });
    }
}
<commit_msg>Added a comment to the main function<commit_after>use std::io::Write;
use std::io::Read;
use std::net::TcpListener;
use std::thread;
use std::sync::{Arc, Mutex};

mod server;
mod broker;
mod message;

struct TcpClient<'a> {
    stream: &'a mut std::net::TcpStream,
}

impl<'a> TcpClient<'a> {
    fn new(stream: &'a mut std::net::TcpStream) -> Self {
        TcpClient { stream : stream }
    }
}

impl<'a> server::Client for TcpClient<'a> {
    fn send(&mut self, bytes: &[u8]) {
        let _ = self.stream.write(bytes);
    }
}

fn main() {
    let server = Arc::new(Mutex::new(server::Server::new()));
    let listener = TcpListener::bind(("0.0.0.0", 1883)).unwrap();

    for byte_stream in listener.incoming() {
        let server = server.clone();
        let mut mqtt_stream = server::Stream::new();

        thread::spawn(move || {
            let mut byte_stream = byte_stream.unwrap();
            loop {
                let read_result = byte_stream.read(mqtt_stream.buffer());
                //the client must be created here due to scoping issues with byte_stream
                let mut client = TcpClient::new(&mut byte_stream);

                match read_result {
                    Ok(length) => {
                        let mut server = server.lock().unwrap();
                        mqtt_stream.handle_messages(length, &mut server, &mut client);
                    },
                    _ => panic!("Error reading bytes from stream"),
                }
            }
        });
    }
}
<|endoftext|>