novel_api/ciweimao/
server.rs

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