tokio async&await 初探

rust的async/await终于在万众瞩目之下稳定下来了,一起来尝尝鲜. 这篇文章主要是介绍基于tokio 0.2做一个服务程序员的小工具githubdns.

一 想解决的问题

github是程序员日常离不开的工具,但是国内访问实在是太慢.看到网上介绍的,通过ipaddress.com来查询github.com地址的方式,验证了一下,还有一定效果. 但是每次都要手工操作,有点麻烦,就做了这么一个小工具githubdns.同时其实也是想试试tokio,看看方便不.

工具的用法

sudo githubdns

非常简单可以不带任何参数,执行完毕后会在/etc/hosts文件中添加三行对于codeload.github.com,github.com,github.global.ssl.fastly.net的解析. 这就是我想做的全部事情.
如果大家有更好的提速方式,请告诉我,提前感谢!

二 tokio 简介

tokio现在基本上是Rust上异步编程的标配了, 用官方的话来说,他就是一个Rust的异步程序Runtime.目前的0.2版本已经完全按照async/await重构,用起来非常方便. 另外热议的Rust的零成本抽象我就不罗嗦了.

三 任务分解

3.1 获取网页

找到域名对应的ip地址,这部分看起来比较简单,就是一个https请求. 比如https://github.com.ipaddress.com
看着简单,我就在这里卡了一会儿. Rust毕竟是新兴语言,就这么一个小功能,都没有现成的,如果是go的话,一句话就搞定了.
但是既然想尝鲜,就从最基础的Tcp开干吧.

3.1.1 建立Tcp连接

这个属于最常用的功能了,非常方便. 一句话

let socket = TcpStream::connect(&addr).await.unwrap();

这里的await特性就是我们要的了,async wait,连接建立完了再继续. 不会一直堵塞当前线程.

3.1.2 https

因为是https连接,所以必须转换成tls连接. 这里用的是tokio-tls,虽然说不是很完善,但是这种基本的操作,还是足够了.

let builder = native_tls::TlsConnector::builder();
let connector = builder.build().unwrap();
let connector = tokio_tls::TlsConnector::from(connector);
let mut socket = connector.connect(domain.as_str(), socket).await?;

略显罗嗦了,四行才行. 前三行可以封装的更好一点. 第四行的connect传入了domain参数,就是为了进行tls握手,验证证书的有效性.

3.1.3 获取网页

如果有封装好的tokio-https库,这里的三个步骤应该是一步完成的. 无奈没有. 本来想自己封装一个简单的,但是嫌太罗嗦了,针对这个小工具,也没有必要,干脆裸上吧.

socket.write_all(format!("GET {} HTTP/1.0\r\nHost:{}\r\n\r\n", path, domain).as_bytes()).await?;

let mut data = Vec::new();
socket.read_to_end(&mut data).await?;

自己拼一个header发出去,然后直接抓取response. 因为我们只关心body中的html,不关心response中的header,直接扔掉.

let s = String::from_utf8(data)?;
let pos = s.find("\r\n\r\n").unwrap_or(0);
let (_, body) = s.split_at(pos);

3.1.4 完整的抓网页

//根据url,获取其地址对应的html内容
async fn get(domain: String) -> Result<String, Box<dyn Error>> {
    let (domain, path) = parse_url(domain.as_str());
    let ip_port = format!("{}:443", domain.clone());
    let addr = ip_port.to_socket_addrs().unwrap().next().unwrap();
    let socket = TcpStream::connect(&addr).await.unwrap();
    let builder = native_tls::TlsConnector::builder();
    let connector = builder.build().unwrap();
    let connector = tokio_tls::TlsConnector::from(connector);
    let mut socket = connector.connect(domain.as_str(), socket).await?;
    socket
        .write_all(format!("GET {} HTTP/1.0\r\nHost:{}\r\n\r\n", path, domain).as_bytes())
        .await?;

    let mut data = Vec::new();
    socket.read_to_end(&mut data).await?;
    let s = String::from_utf8(data)?;
    let pos = s.find("\r\n\r\n").unwrap_or(0);
    let (_, body) = s.split_at(pos);
    Ok(String::from(body))
}

首先是函数的签名async fn get(domain: String) -> Result<String, Box<dyn Error>>. 必须是async,否则函数体中是无法使用await的. 感兴趣的同学可以看看网上的教程. 简单的说就是async关键字会把我们的返回值转换为Future.
而里面的await关键字则会自动为我们保存上下文,封装成一个状态机.

3.2 解析网页

这个就简单多了,我们有现成的crate scraper,拿来用即可.因为重点是说异步,这里的代码虽然有点长,就不罗嗦了.关键可以用一句jquery来描述

$("ul.comma-separated li")

3.3 写配置文件

对于文件的异步读写,使用tokio-fs,非常方便.

let contents = fs::read(hosts_file_name).await;
//... 修改
fs::write(hosts_file_name, lines.join(enter).as_bytes()).await;

ok,一个小工具就完成了.

3.4 合在一起

工具的目标是抓取多个域名对应的ip地址,然后写入配置文件. 既然是异步,肯定要同时抓取多个.这里顺便展示一下join_all如何使用了.

let domains_str = m.value_of("domains").unwrap().to_string();
let domains: Vec<&str> = domains_str.split(",").collect();
let domain_ips = Arc::new(RwLock::new(HashMap::new()));

let mut v = Vec::new();
for domain in domains.iter() {
    let domain_ips = domain_ips.clone();
    let domain = String::from(*domain);
    v.push(async move { // 大名鼎鼎的async move,用起来真香!
        let res = get(domain.clone()).await;
        if res.is_err() {
            println!(
                "get domain {},err {}",
                domain,
                res.unwrap_err().description()
            );
            return;
        }
        let res = res.unwrap();
        let ip = get_address(&res, domain.clone());
        domain_ips.write().unwrap().insert(domain.clone(), ip);
    });
}
join_all(v).await;

let (hosts_file, enter) = get_hosts_file();
read_and_modify_hosts(domain_ips.clone(), &hosts_file, &enter).await;
Ok(())

整体看给人感觉还是略显罗嗦,不过用rust写代码,我也明显感觉到罗嗦,也可能是我功力不够,不能吐槽rust了.
如果抛开错误处理,我们可以很简洁的.

 for domain in domains.iter() {
        let domain_ips = domain_ips.clone();
        let domain = String::from(*domain);
        v.push(async move {
            //取网页内容
            let res = get(domain.clone()).await?;
            //解析ip地址
            let ip = get_address(&res, domain.clone());
            //放入map中,好写入hosts文件
            domain_ips.write().unwrap().insert(domain.clone(), ip);
        });
    }
join_all(v).await;

这样看来是不是还是比较清晰的. 多个连接同时发出,又不用像goroutine一样启动协程,总的来说还是感觉很清爽的.

3.5 main函数

为了更方便的使用tokio,避免手工使用tokio::spawn之类的,tokio提供了async main. 使用起来是真香!

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> 

3.6 其他

这里故意忽略了一些稍微复杂的内容,比如hosts文件的解析以及不同平台的差异. 这些是所有代码都无法绕开的.
不过还有一点是要特别吐槽的,rust的String设计的真是不好用,导致字符串的处理总是显得比较罗嗦.

整个下来,有230行左右, 不过我想已经把tokio异步编程要点都覆盖到了.

四 完整的程序

纯粹是为了让文章显得长一点,哈哈,完全可以忽略.或者直接到我的github上看.

use futures::future::*;
use native_tls;
use scraper::{Html, Selector};
use std::collections::HashMap;
use std::error::Error;

use std::net::{IpAddr, Ipv4Addr, ToSocketAddrs};
use std::sync::{Arc, RwLock};

use clap::{App, Arg};
use std::process::Command;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let m = App::new("githubdns")
        .arg(
            Arg::with_name("domains")
                .long("domains")
                .help("set the dns needs to write to hosts")
                .default_value("github.com,github.global.ssl.fastly.net,codeload.github.com,assets-cdn.github.com")
                .required(false),
        )
        .get_matches();
    let domains_str = m.value_of("domains").unwrap().to_string();
    let domains: Vec<&str> = domains_str.split(",").collect();
    let domain_ips = Arc::new(RwLock::new(HashMap::new()));

    let mut v = Vec::new();
    for domain in domains.iter() {
        let domain_ips = domain_ips.clone();
        let domain = String::from(*domain);
        v.push(async move {
            let res = get(domain.clone()).await;
            if res.is_err() {
                println!(
                    "get domain {},err {}",
                    domain,
                    res.unwrap_err().description()
                );
                return;
            }
            let res = res.unwrap();
            let ip = get_address(&res, domain.clone());
            domain_ips.write().unwrap().insert(domain.clone(), ip);
        });
    }
    join_all(v).await;

    //    println!("domain_ips={:?}", domain_ips.read().unwrap());
    let (hosts_file, enter) = get_hosts_file();
    read_and_modify_hosts(domain_ips.clone(), &hosts_file, &enter).await;
    Ok(())
}
//hosts文件路径,以及回车换行对应的是\r\n还是\n,\r
//这里有一个副作用,会讲hosts文件的只读属性移除,可以考虑写入后再增加上,
fn get_hosts_file() -> (String, String) {
    let info = sys_info::os_type();
    if info.is_err() {
        panic!("unsupported os");
    }
    let info = info.unwrap();
    //    Such as "Linux", "Darwin", "Windows".
    match info.as_str() {
        "Linux" => ("/etc/hosts".into(), "\n".into()),
        "Darwin" => ("/etc/hosts".into(), "\r".into()),
        "Windows" => {
            let path = r"C:\Windows\System32\drivers\etc\hosts";
            //windows下hosts文件默认是只读的,如果不修改,后续会写不进去
            Command::new("attrib")
                .args(&["-R", path])
                .output()
                .expect("remove readonly failed");

            (path.into(), "\r\n".into())
        }
        _ => panic!("not supported os {}", info),
    }
}
//读取hosts文件,如果其中已经有相关域名的设置,先删除,再添加
//原有注释保持不动
async fn read_and_modify_hosts(
    m: Arc<RwLock<HashMap<String, String>>>,
    hosts_file: &str,
    enter: &str,
) {
    use tokio::fs;
    let flags = "# ----Generated By githubdns ---";
    let hosts_file_name = hosts_file;
    let mut m = m.write().unwrap();
    let contents = fs::read(hosts_file_name).await;
    if contents.is_err() {
        println!(
            "read {} err {}",
            hosts_file_name,
            contents.err().unwrap().description()
        );
        return;
    }
    let contents = contents.unwrap();
    let s = String::from_utf8(contents).unwrap();
    let mut lines: Vec<&str> = s.split(enter).collect();
    let mut i = 0;
    while i < lines.len() {
        let l = lines.get(i).unwrap().clone();
        if l == flags {
            lines.remove(i);
            continue;
        }
        if l.trim_start().starts_with("#") {
            i += 1;
            continue; //注释行
        }
        let _ = m.iter().any(|(domain, ip)| {
            let pos = l.find(domain.as_str());
            if pos.is_some() {
                //如果是这个domain的子域名,也不关心
                let pos = pos.unwrap();
                if ip.len() > 0
                    && pos > 0
                    && (l.as_bytes()[pos - 1] == ' ' as u8 || l.as_bytes()[pos - 1] == '\t' as u8)
                {
                    //是我们要找的完整的域名
                    lines.remove(i);
                    i -= 1;
                    return true;
                }
            }
            return false;
        });
        i += 1;
    }

    let mut lines: Vec<_> = lines.iter().map(|n| String::from(*n)).collect();

    lines.push(flags.into());
    for (domain, ip) in m.iter_mut() {
        if ip.len() > 0 {
            lines.push(format!("{}\t {}", ip, domain));
        }
    }
    lines.push(flags.into());
    let r = fs::write(hosts_file_name, lines.join(enter).as_bytes()).await;
    if r.is_err() {
        panic!("write to {} ,err={}", hosts_file, r.unwrap_err());
    }
}
//解析url,返回对应的domain和path
fn parse_url(domain: &str) -> (String, String) {
    let ss: Vec<_> = domain.split(".").collect();
    let mut path = "/".into();
    let mut domain: String = domain.into();
    if ss.len() > 2 {
        path = format!("/{}", domain.clone());
        domain = ss[ss.len() - 2..].join(".");
    }
    domain = format!("{}.ipaddress.com", domain);
    return (domain, path);
}

//根据url,获取其地址对应的html内容
async fn get(domain: String) -> Result<String, Box<dyn Error>> {
    let (domain, path) = parse_url(domain.as_str());
    println!("get {},{}", domain, path);
    let ip_port = format!("{}:443", domain.clone());
    //    println!("ip_port={}", ip_port);
    let addr = ip_port.to_socket_addrs().unwrap().next().unwrap();
    //    println!("addr={}", addr);
    let socket = TcpStream::connect(&addr).await.unwrap();
    // Send off the request by first negotiating an SSL handshake, then writing
    // of our request, then flushing, then finally read off the response.
    let builder = native_tls::TlsConnector::builder();
    let connector = builder.build().unwrap();
    let connector = tokio_tls::TlsConnector::from(connector);
    let mut socket = connector.connect(domain.as_str(), socket).await?;
    socket
        .write_all(format!("GET {} HTTP/1.0\r\nHost:{}\r\n\r\n", path, domain).as_bytes())
        .await?;

    let mut data = Vec::new();
    socket.read_to_end(&mut data).await?;
    let s = String::from_utf8(data)?;
    let pos = s.find("\r\n\r\n").unwrap_or(0);
    let (_, body) = s.split_at(pos);
    //    println!("body={}", body);
    Ok(String::from(body))
}
//从html中提取domain对应的第一个ipv4地址
fn get_address(data: &str, domain: String) -> String {
    let document = Html::parse_document(data);
    let ul_selector = Selector::parse("ul.comma-separated").unwrap();
    let li_selector = Selector::parse("li").unwrap();
    let ul = document.select(&ul_selector).next();
    if ul.is_none() {
        println!("{} cannot found ul,data={}", domain, data);
        return String::new();
    }
    let ul = ul.unwrap();
    let mut ip_v4 = Ipv4Addr::new(127, 0, 0, 1);
    let found = ul.select(&li_selector).any(|n| {
        //        println!("n={}", n.inner_html().trim());
        let ip: Result<IpAddr, _> = n.inner_html().trim().parse();
        match ip {
            Err(_) => {
                return false;
            }
            Ok(ip) => match ip {
                IpAddr::V4(ipv4) => {
                    ip_v4 = ipv4;
                    return true;
                }
                _ => {
                    return false;
                }
            },
        }
    });
    if found {
        return ip_v4.to_string();
    }
    return String::new(); //没有找到就返回空
}

#[cfg(test)]
mod test {
    use super::*;
    #[test]
    fn test_parse_url() {
        let u = parse_url("github.global.ssl.fastly.net");
        assert_eq!(
            u,
            (
                String::from("fastly.net.ipaddress.com"),
                String::from("/github.global.ssl.fastly.net")
            )
        );
        let u = parse_url("github.com");
        assert_eq!(
            u,
            (String::from("github.com.ipaddress.com"), String::from("/"))
        );
    }
    #[tokio::test]
    async fn test_get() {
        assert!(true);
    }
    #[test]
    fn test_get_address() {
        let data = std::fs::read_to_string("github.com.html").unwrap();
        let ip = get_address(data.as_str(), String::from("github.com"));
        assert_eq!(ip, String::from("192.30.253.112"))
    }
}