vm_tools/chunnel/src/bin/chunnel.rs

// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::env;
use std::ffi::CStr;
use std::fmt;
use std::process;
use std::result;

use getopts::Options;
use libchromeos::syslog;
use log::warn;
use sys_util::{self, block_signal, PollContext, PollToken};

use chunnel::forwarder::{ForwarderError, ForwarderSession};
use chunnel::stream::{StreamSocket, StreamSocketError};

// Program name.
const IDENT: &[u8] = b"chunnel\0";

#[remain::sorted]
#[derive(Debug)]
enum Error {
    BlockSigpipe(sys_util::signal::Error),
    ConnectSocket(StreamSocketError),
    Forward(ForwarderError),
    PollContextDelete(sys_util::Error),
    PollContextNew(sys_util::Error),
    PollWait(sys_util::Error),
    Syslog(log::SetLoggerError),
}

type Result<T> = result::Result<T, Error>;

impl fmt::Display for Error {
    #[remain::check]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        #[remain::sorted]
        match self {
            BlockSigpipe(e) => write!(f, "failed to block SIGPIPE: {}", e),
            ConnectSocket(e) => write!(f, "failed to connnect socket: {}", e),
            Forward(e) => write!(f, "failed to forward traffic: {}", e),
            PollContextDelete(e) => write!(f, "failed to delete fd from poll context: {}", e),
            PollContextNew(e) => write!(f, "failed to create poll context: {}", e),
            PollWait(e) => write!(f, "failed to wait for poll: {}", e),
            Syslog(e) => write!(f, "failed to initialize syslog: {}", e),
        }
    }
}

fn print_usage(program: &str, opts: &Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

fn run_forwarder(local_stream: StreamSocket, remote_stream: StreamSocket) -> Result<()> {
    block_signal(libc::SIGPIPE).map_err(Error::BlockSigpipe)?;

    #[derive(PollToken)]
    enum Token {
        LocalStreamReadable,
        RemoteStreamReadable,
    }
    let poll_ctx: PollContext<Token> = PollContext::build_with(&[
        (&local_stream, Token::LocalStreamReadable),
        (&remote_stream, Token::RemoteStreamReadable),
    ])
    .map_err(Error::PollContextNew)?;

    let mut forwarder = ForwarderSession::new(local_stream, remote_stream);

    loop {
        let events = poll_ctx.wait().map_err(Error::PollWait)?;

        for event in events.iter_readable() {
            match event.token() {
                Token::LocalStreamReadable => {
                    let shutdown = forwarder.forward_from_local().map_err(Error::Forward)?;
                    if shutdown {
                        poll_ctx
                            .delete(forwarder.local_stream())
                            .map_err(Error::PollContextDelete)?;
                    }
                }
                Token::RemoteStreamReadable => {
                    let shutdown = forwarder.forward_from_remote().map_err(Error::Forward)?;
                    if shutdown {
                        poll_ctx
                            .delete(forwarder.remote_stream())
                            .map_err(Error::PollContextDelete)?;
                    }
                }
            }
        }
        if forwarder.is_shut_down() {
            return Ok(());
        }
    }
}

fn main() -> Result<()> {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    opts.reqopt("l", "local", "local socket to forward", "SOCKADDR");
    opts.reqopt("r", "remote", "remote socket to forward to", "SOCKADDR");
    opts.optopt("t", "type", "type of traffic to forward", "stream|datagram");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(e) => {
            warn!("failed to parse arg: {}", e);
            print_usage(&program, &opts);
            process::exit(1);
        }
    };
    if matches.opt_present("h") {
        print_usage(&program, &opts);
        return Ok(());
    }

    // Safe because this string is defined above in this file and it contains exactly
    // one nul byte, which appears at the end.
    let ident = CStr::from_bytes_with_nul(IDENT).unwrap();
    syslog::init(ident).map_err(Error::Syslog)?;

    let local_sockaddr = match matches.opt_str("l") {
        Some(sockaddr) => sockaddr,
        None => {
            warn!("local socket must be defined");
            print_usage(&program, &opts);
            process::exit(1);
        }
    };

    let remote_sockaddr = match matches.opt_str("r") {
        Some(sockaddr) => sockaddr,
        None => {
            warn!("remote socket must be defined");
            print_usage(&program, &opts);
            process::exit(1);
        }
    };

    // Default to "stream" if traffic type is not defined.
    let traffic_type = matches.opt_str("t");
    if let Some(t) = traffic_type {
        match t.as_ref() {
            "stream" => {}
            "datagram" => {
                warn!("datagram sockets are not yet supported");
                process::exit(1);
            }
            s => {
                warn!("not a valid type of traffic: {}", s);
                print_usage(&program, &opts);
                process::exit(1);
            }
        }
    }

    let local_stream = StreamSocket::connect(&local_sockaddr).map_err(Error::ConnectSocket)?;
    let remote_stream = StreamSocket::connect(&remote_sockaddr).map_err(Error::ConnectSocket)?;

    run_forwarder(local_stream, remote_stream)
}