Rustでmastodonのstreaming APIを叩く

Rustでmastodonのタイムラインをstreamingで取得するためのAPIを叩いてみようと思ったので、その結果をまとめます。

Cargo.toml

[package]
name = "mastodon-tl"
version = "0.1.0"
authors = ["Pocket7878 <[e-mail]>"]

[dependencies]
reqwest = "0.5.1"
serde = "1.0.2"
serde_derive = "1.0.2"
serde_json = "1.0"
dotenv = "0.10.0"

[dependencies.url]
git = "https://github.com/servo/rust-url"

リクエスト周りはreqwestがhyperのラッパーになっていて多少らくに使えそうだったので
採用しています。また、クライアントのID等を保存しておくためにdotenvを利用しています。

実装

streaming API に
接続し、一行一行BufReadで読み出すことによって取得しています。

ハートビートを除くと、イベントがあってそれからpayloadがきて空白行がきてという繰り返しなので、それを順次処理しています。
payloadのなかの処理なども行ったらターミナルにトゥートを流すことができそうです。

main.rs

extern crate serde;
extern crate serde_json;
extern crate url;
extern crate dotenv;
extern crate reqwest;

use dotenv::dotenv;
use std::env;

#[macro_use]
extern crate serde_derive;

use std::io;
use std::fs::File;
use std::io::prelude::*;

#[derive(Serialize, Deserialize, Debug)]
struct ClientInfo {
    client_id: String,
    client_secret: String,
}

#[derive(Serialize, Deserialize, Debug)]
struct AccessToken {
    access_token: String,
}

// ClientInfo
const MASTODON_HOST: &'static str = "https://mstdn.jp";
const MASTODON_CLIENT_ID_KEY: &'static str = "MASTODON_RS_CLIENT_ID";
const MASTODON_CLIENT_SECRET_KEY: &'static str = "MASTODON_RS_CLIENT_SECRET";

fn read_client_info_from_dotenv() -> Option<ClientInfo> {
    let mut client_id: Option<String> = None;
    let mut client_secret: Option<String> = None;
    for (key, value) in env::vars() {
        if key == MASTODON_CLIENT_ID_KEY {
            client_id = Some(value);
        } else if key == MASTODON_CLIENT_SECRET_KEY {
            client_secret = Some(value);
        }
    }
    match (client_id, client_secret) {
        (Some(id), Some(sec)) => {
            Some(ClientInfo {
                client_id: id,
                client_secret: sec,
            })
        }
        _ => None,
    }
}

fn register_client(base_url: &str) -> Option<ClientInfo> {
    let client = reqwest::Client::new().unwrap();
    let params = [("client_name", "mastodon-tl-rs"),
                  ("redirect_uris", "urn:ietf:wg:oauth:2.0:oob"),
                  ("scopes", "read write follow")];
    let res = client.post(base_url)
        .form(&params)
        .send()
        .unwrap();
    if res.status().is_success() {
        let client_info = serde_json::from_reader(res);
        return Some(client_info.unwrap());
    } else {
        return None;
    }
}

fn write_client_to_dotenv(client_info: &ClientInfo) {
    let mut f = File::create(".env").unwrap();
    f.write_all(format!("{}={}\n", MASTODON_CLIENT_ID_KEY, client_info.client_id).as_bytes());
    f.write_all(format!("{}={}\n",
                        MASTODON_CLIENT_SECRET_KEY,
                        client_info.client_secret)
        .as_bytes());
}

fn build_authorize_url(client_info: &ClientInfo) -> String {
    let param_str = format!("client_id={}&response_type=code&redirect_uri=urn:ietf:wg:oauth:2.0:\
                            oob&scope=read write follow",
                           client_info.client_id);
    let params =
        url::percent_encoding::utf8_percent_encode(param_str.as_str(),
                                                   url::percent_encoding::QUERY_ENCODE_SET);
    return format!("{}/oauth/authorize?{}", MASTODON_HOST, params);
}

fn get_access_token(client_info: &ClientInfo, auth_code: &str) -> Option<String> {
    let client = reqwest::Client::new().unwrap();
    let params = [("grant_type", "authorization_code".to_string()),
                  ("redirect_uri", "urn:ietf:wg:oauth:2.0:oob".to_string()),
                  ("client_id", client_info.client_id.clone()),
                  ("client_secret", client_info.client_secret.clone()),
                  ("code", auth_code.to_string()),
                  ("scope", "read write follow".to_string())];
    let res = client.post(&format!("{}/oauth/token", MASTODON_HOST))
        .form(&params)
        .send()
        .unwrap();
    if res.status().is_success() {
        let access_token: AccessToken = serde_json::from_reader(res).unwrap();
        return Some(access_token.access_token);
    } else {
        return None;
    }
}

fn watching_tl(access_token: &str) {
    use std::thread;
    use std::io::BufReader;
    use reqwest::header::{Authorization, Bearer};

    let client = reqwest::Client::new().unwrap();
    let res = client.get(&format!("{}/api/v1/streaming/public", MASTODON_HOST))
        .header(Authorization(Bearer { token: access_token.to_string() }))
        .send()
        .unwrap();
    if res.status().is_success() {
        println!("Connected");
        let receive_loop = thread::spawn(move || {
            // Receive loop
            let mut current_event_key: Option<String> = None;
            let mut current_event_value: Option<String> = None;
            let buf_resp = BufReader::new(res);
            for line in buf_resp.lines() {
                let l = line.unwrap();
                if l.is_empty() {
                    match (current_event_key.clone(), current_event_value.clone()) {
                        (Some(k), Some(v)) => {
                            println!("{}", k);
                            println!("{}", v);
                        },
                        _ => {
                            panic!("Illigal status");
                        }
                    }
                } else if l.starts_with(":") {
                    println!("-- heartbeat --");
                } else if l.starts_with("event: ") {
                    let key = &l[7..];
                    current_event_key = Some(key.to_string());
                } else if l.starts_with("data: ") {
                    let value = &l[6..];
                    current_event_value = Some(value.to_string());
                }
            }
        });
        let _ = receive_loop.join();
    }
}

fn main() {
    dotenv().ok();

    let client_info =
        read_client_info_from_dotenv().or(register_client(&format!("{}/api/v1/apps", MASTODON_HOST))).unwrap();
    write_client_to_dotenv(&client_info);
    let authorize_url = build_authorize_url(&client_info);
    println!("Auth Url:");
    println!("{}", authorize_url);
    println!("Input authorization code:");
    let mut auth_code = String::new();
    io::stdin().read_line(&mut auth_code);
    let access_token = get_access_token(&client_info, auth_code.trim()).unwrap();
    watching_tl(&access_token);
}