TOKIO ASYNC&AWAIT 初探

• tokio async&await 初探

• 3.1.1 建立Tcp连接
• 3.1.2 https
• 3.1.3 获取网页
• 3.1.4 完整的抓网页
• 一 想解决的问题
• 工具的用法
• 二 tokio 简介
• 三 任务分解
• 3.1 获取网页
• 3.2 解析网页
• 3.3 写配置文件
• 3.4 合在一起
• 3.5 main函数
• 3.6 其他
• 四 完整的程序
  

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"))
}
}