-
Notifications
You must be signed in to change notification settings - Fork 340
/
main.rs
219 lines (189 loc) · 6.38 KB
/
main.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
use anyhow::{anyhow, Result};
use clap::Parser;
use colored::Colorize;
use mime::Mime;
use reqwest::{header, Client, Response, Url};
use std::{collections::HashMap, str::FromStr};
use syntect::{
easy::HighlightLines,
highlighting::{Style, ThemeSet},
parsing::SyntaxSet,
util::{as_24_bit_terminal_escaped, LinesWithEndings},
};
// 以下部分用于处理 CLI
// 定义 httpie 的 CLI 的主入口,它包含若干个子命令
// 下面 /// 的注释是文档,clap 会将其作为 CLI 的帮助
/// A naive httpie implementation with Rust, can you imagine how easy it is?
#[derive(Parser, Debug)]
#[clap(version = "1.0", author = "Tyr Chen <tyr@chen.com>")]
struct Opts {
#[clap(subcommand)]
subcmd: SubCommand,
}
// 子命令分别对应不同的 HTTP 方法,目前只支持 get / post
#[derive(Parser, Debug)]
enum SubCommand {
Get(Get),
Post(Post),
// 我们暂且不支持其它 HTTP 方法
}
// get 子命令
/// feed get with an url and we will retrieve the response for you
#[derive(Parser, Debug)]
struct Get {
/// HTTP 请求的 URL
#[clap(parse(try_from_str = parse_url))]
url: String,
}
// post 子命令。需要输入一个 url,和若干个可选的 key=value,用于提供 json body
/// feed post with an url and optional key=value pairs. We will post the data
/// as JSON, and retrieve the response for you
#[derive(Parser, Debug)]
struct Post {
/// HTTP 请求的 URL
#[clap(parse(try_from_str = parse_url))]
url: String,
/// HTTP 请求的 body
#[clap(parse(try_from_str=parse_kv_pair))]
body: Vec<KvPair>,
}
/// 命令行中的 key=value 可以通过 parse_kv_pair 解析成 KvPair 结构
#[derive(Debug, PartialEq)]
struct KvPair {
k: String,
v: String,
}
/// 当我们实现 FromStr trait 后,可以用 str.parse() 方法将字符串解析成 KvPair
impl FromStr for KvPair {
type Err = anyhow::Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
// 使用 = 进行 split,这会得到一个迭代器
let mut split = s.split('=');
let err = || anyhow!(format!("Failed to parse {}", s));
Ok(Self {
// 从迭代器中取第一个结果作为 key,迭代器返回 Some(T)/None
// 我们将其转换成 Ok(T)/Err(E),然后用 ? 处理错误
k: (split.next().ok_or_else(err)?).to_string(),
// 从迭代器中取第二个结果作为 value
v: (split.next().ok_or_else(err)?).to_string(),
})
}
}
/// 因为我们为 KvPair 实现了 FromStr,这里可以直接 s.parse() 得到 KvPair
fn parse_kv_pair(s: &str) -> Result<KvPair> {
s.parse()
}
fn parse_url(s: &str) -> Result<String> {
// 这里我们仅仅检查一下 URL 是否合法
let _url: Url = s.parse()?;
Ok(s.into())
}
/// 处理 get 子命令
async fn get(client: Client, args: &Get) -> Result<()> {
let resp = client.get(&args.url).send().await?;
Ok(print_resp(resp).await?)
}
/// 处理 post 子命令
async fn post(client: Client, args: &Post) -> Result<()> {
let mut body = HashMap::new();
for pair in args.body.iter() {
body.insert(&pair.k, &pair.v);
}
let resp = client.post(&args.url).json(&body).send().await?;
Ok(print_resp(resp).await?)
}
// 打印服务器版本号 + 状态码
fn print_status(resp: &Response) {
let status = format!("{:?} {}", resp.version(), resp.status()).blue();
println!("{}\n", status);
}
// 打印服务器返回的 HTTP header
fn print_headers(resp: &Response) {
for (name, value) in resp.headers() {
println!("{}: {:?}", name.to_string().green(), value);
}
println!();
}
/// 打印服务器返回的 HTTP body
fn print_body(m: Option<Mime>, body: &str) {
match m {
// 对于 "application/json" 我们 pretty print
Some(v) if v == mime::APPLICATION_JSON => print_syntect(body, "json"),
Some(v) if v == mime::TEXT_HTML => print_syntect(body, "html"),
// 其它 mime type,我们就直接输出
_ => println!("{}", body),
}
}
/// 打印整个响应
async fn print_resp(resp: Response) -> Result<()> {
print_status(&resp);
print_headers(&resp);
let mime = get_content_type(&resp);
let body = resp.text().await?;
print_body(mime, &body);
Ok(())
}
/// 将服务器返回的 content-type 解析成 Mime 类型
fn get_content_type(resp: &Response) -> Option<Mime> {
resp.headers()
.get(header::CONTENT_TYPE)
.map(|v| v.to_str().unwrap().parse().unwrap())
}
/// 程序的入口函数,因为在 http 请求时我们使用了异步处理,所以这里引入 tokio
#[tokio::main]
async fn main() -> Result<()> {
let opts: Opts = Opts::parse();
let mut headers = header::HeaderMap::new();
// 为我们的 http 客户端添加一些缺省的 HTTP 头
headers.insert("X-POWERED-BY", "Rust".parse()?);
headers.insert(header::USER_AGENT, "Rust Httpie".parse()?);
let client = reqwest::Client::builder()
.default_headers(headers)
.build()?;
let result = match opts.subcmd {
SubCommand::Get(ref args) => get(client, args).await?,
SubCommand::Post(ref args) => post(client, args).await?,
};
Ok(result)
}
fn print_syntect(s: &str, ext: &str) {
// Load these once at the start of your program
let ps = SyntaxSet::load_defaults_newlines();
let ts = ThemeSet::load_defaults();
let syntax = ps.find_syntax_by_extension(ext).unwrap();
let mut h = HighlightLines::new(syntax, &ts.themes["base16-ocean.dark"]);
for line in LinesWithEndings::from(s) {
let ranges: Vec<(Style, &str)> = h.highlight(line, &ps);
let escaped = as_24_bit_terminal_escaped(&ranges[..], true);
print!("{}", escaped);
}
}
// 仅在 cargo test 时才编译
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_url_works() {
assert!(parse_url("abc").is_err());
assert!(parse_url("http://abc.xyz").is_ok());
assert!(parse_url("https://httpbin.org/post").is_ok());
}
#[test]
fn parse_kv_pair_works() {
assert!(parse_kv_pair("a").is_err());
assert_eq!(
parse_kv_pair("a=1").unwrap(),
KvPair {
k: "a".into(),
v: "1".into()
}
);
assert_eq!(
parse_kv_pair("b=").unwrap(),
KvPair {
k: "b".into(),
v: "".into()
}
);
}
}