首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >TOKIO ASYNC&AWAIT 初探

TOKIO ASYNC&AWAIT 初探

作者头像
MikeLoveRust
发布2019-12-17 17:08:38
1.3K0
发布2019-12-17 17:08:38
举报
  • 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"))
    }
}

本文来自bai的投稿,原文地址:https://stevenbai.top/rust/tokio_async_await-%E5%88%9D%E6%8E%A2/

本文参与 腾讯云自媒体分享计划,分享自微信公众号。
原始发表:2019-12-07,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 Rust语言学习交流 微信公众号,前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体分享计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
目录
  • 一 想解决的问题
    • 工具的用法
    • 二 tokio 简介
    • 三 任务分解
      • 3.1 获取网页
        • 3.1.1 建立Tcp连接
        • 3.1.2 https
        • 3.1.3 获取网页
      • 3.2 解析网页
        • 3.4 合在一起
          • 3.5 main函数
          • 四 完整的程序
          领券
          问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档