kungfu 0.1.27

Flexible DNS hijacking and proxy tool. Read more: https://github.com/yinheli/kungfu 
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

use std::{net::SocketAddr, str::FromStr};

use bytes::Bytes;
use http_body_util::Full;
use hyper::{server::conn::http1, service::service_fn, Request, Response, StatusCode};
use hyper_util::rt::TokioIo;
use log::{info, warn};
use prometheus::{Encoder, TextEncoder};
use tokio::{net::TcpListener, spawn};

pub async fn serve(addr: Option<String>) {
    if addr.is_none() {
        return;
    }

    let addr = addr.unwrap();

    let addr = SocketAddr::from_str(&addr).unwrap();

    let listener = TcpListener::bind(addr).await.unwrap();

    info!("metrics server listening on {}", addr);

    loop {
        let (stream, _) = listener.accept().await.unwrap();
        let io = TokioIo::new(stream);
        spawn(async move {
            if let Err(err) = http1::Builder::new()
                .serve_connection(io, service_fn(metrics))
                .await
            {
                warn!("Error serving connection: {:?}", err);
            }
        });
    }
}

async fn metrics(
    req: Request<hyper::body::Incoming>,
) -> Result<Response<Full<Bytes>>, hyper::Error> {
    match req.uri().path() {
        "/metrics" => {
            let encoder = TextEncoder::new();
            let metric_families = prometheus::gather();
            let mut buffer = Vec::with_capacity(1024 * 32);
            encoder.encode(&metric_families, &mut buffer).unwrap();

            let body = Full::new(Bytes::from(buffer));

            let response = Response::builder()
                .status(200)
                .header(hyper::header::CONTENT_TYPE, encoder.format_type())
                .body(body)
                .unwrap();

            Ok(response)
        }
        _ => Ok(Response::builder()
            .status(StatusCode::NOT_FOUND)
            .body(Default::default())
            .unwrap()),
    }
}