novel_api/ciweimao/
server.rs

1use std::net::{IpAddr, Ipv4Addr, SocketAddr};
2
3use askama::Template;
4use axum::{
5    extract::{self, State},
6    http::{header, StatusCode},
7    response::{Html, IntoResponse, Response},
8    routing, Router,
9};
10use rust_embed::RustEmbed;
11use tokio::{
12    net::TcpListener,
13    sync::{
14        mpsc::{self, Sender},
15        oneshot,
16    },
17    task,
18};
19
20use super::GeetestInfoResponse;
21use crate::Error;
22
23#[derive(RustEmbed)]
24#[folder = "templates"]
25struct Asset;
26
27struct StaticFile<T>(pub T);
28
29impl<T> IntoResponse for StaticFile<T>
30where
31    T: Into<String>,
32{
33    fn into_response(self) -> Response {
34        let path = self.0.into();
35
36        match Asset::get(path.as_str()) {
37            Some(content) => {
38                let mime = mime_guess::from_path(path).first_or_octet_stream();
39                ([(header::CONTENT_TYPE, mime.as_ref())], content.data).into_response()
40            }
41            None => (StatusCode::NOT_FOUND, "404 Not Found").into_response(),
42        }
43    }
44}
45
46impl IntoResponse for Error {
47    fn into_response(self) -> Response {
48        (
49            StatusCode::INTERNAL_SERVER_ERROR,
50            format!("Something went wrong: {}", self),
51        )
52            .into_response()
53    }
54}
55
56#[derive(Template)]
57#[template(path = "index.html")]
58struct IndexTemplate {
59    gt: String,
60    challenge: String,
61    new_captcha: bool,
62}
63
64async fn captcha(
65    State(state): State<(GeetestInfoResponse, Sender<String>)>,
66) -> Result<IndexTemplate, Error> {
67    let (info, _) = state;
68
69    Ok(IndexTemplate {
70        gt: info.gt,
71        challenge: info.challenge,
72        new_captcha: info.new_captcha,
73    })
74}
75
76async fn geetest_js() -> StaticFile<&'static str> {
77    StaticFile("geetest.js")
78}
79
80async fn validate(
81    extract::Path(validate): extract::Path<String>,
82    State(state): State<(GeetestInfoResponse, Sender<String>)>,
83) -> Html<&'static str> {
84    let (_, tx) = state;
85    tx.send(validate).await.unwrap();
86
87    Html("Verification is successful, you can close the browser now")
88}
89
90pub(crate) async fn run_geetest(info: GeetestInfoResponse) -> Result<String, Error> {
91    let (tx, mut rx) = mpsc::channel(1);
92
93    let app = Router::new()
94        .route("/captcha", routing::get(captcha))
95        .route("/geetest.js", routing::get(geetest_js))
96        .route("/validate/:validate", routing::get(validate))
97        .with_state((info, tx));
98
99    let addr = SocketAddr::new(
100        IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
101        portpicker::pick_unused_port().ok_or(Error::Port(String::from("No ports free")))?,
102    );
103    let listener = TcpListener::bind(addr).await?;
104
105    let (stop_tx, stop_rx) = oneshot::channel();
106
107    task::spawn(async move {
108        axum::serve(listener, app)
109            .with_graceful_shutdown(async {
110                stop_rx.await.ok();
111            })
112            .await?;
113
114        Ok::<_, Error>(())
115    });
116
117    open::that(format!("http://{}:{}/captcha", addr.ip(), addr.port()))?;
118
119    let validate = rx.recv().await.unwrap();
120    stop_tx.send(()).unwrap();
121
122    Ok(validate)
123}