novel_api/ciweimao/
server.rs

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