dhomer 0.1.0

Simple and easy to use, a proxy server based on Pingora
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

pub mod argument;
pub mod c;
pub mod config;
pub mod proxy;
pub mod util;

use std::path::PathBuf;
use std::time::Duration;
use argument::Argument;
use clap::Parser;
use pingora::{proxy::http_proxy_service, server::Server};
use pingora::lb::LoadBalancer;
use pingora::prelude::{background_service, RoundRobin, TcpHealthCheck};
use tracing::{Level, debug, info};
use tracing_appender::rolling::daily;
use tracing_subscriber::{
    EnvFilter,
    fmt::{self, time::OffsetTime, writer::MakeWriterExt},
    layer::SubscriberExt,
    util::SubscriberInitExt,
};
use util::Util;

use crate::proxy::Proxy;

fn init_log(log_dir: &PathBuf) {
    // 配置时间格式
    let time_fmt = time::format_description::parse(
        "[year]-[month]-[day] [hour]:[minute]:[second].[subsecond digits:3]",
    )
    .expect("Failed to parse time format");
    let time_offset = time::UtcOffset::current_local_offset().unwrap_or(time::UtcOffset::UTC);
    let timer = OffsetTime::new(time_offset, time_fmt);

    // 文件输出 - 无ANSI颜色
    let file_appender = daily(log_dir, "err.log")
        .with_max_level(Level::WARN)
        .or_else(daily(log_dir, "out.log"));

    let file_layer = fmt::layer()
        .with_writer(file_appender)
        .with_ansi(false) // 文件输出禁用ANSI颜色
        .with_timer(timer.clone())
        .with_target(true);

    // 控制台输出 - 有ANSI颜色
    let std_appender = std::io::stderr
        .with_max_level(Level::WARN)
        .or_else(std::io::stdout);

    let console_layer = fmt::layer()
        .with_writer(std_appender)
        .with_ansi(true) // 控制台输出启用ANSI颜色
        .with_timer(timer)
        .with_target(true);

    // 组合两层日志配置并初始化
    tracing_subscriber::registry()
        .with(EnvFilter::from_default_env())
        .with(file_layer)
        .with(console_layer)
        .init()
}

fn main() {
    let log_dir = Util::app_dirs().data_dir().join("log");
    init_log(&log_dir);
    debug!(?log_dir);

    let argument = Argument::parse();
    info!(?argument);

    let config = config::parse(&argument)
        .expect("Failed to parse configuration file");
    info!(?config);

    let mut bg_services = Vec::new();
    let mut proxy_entries = Vec::new();

    for (k, setting) in config.proxies() {
        let mut lb = LoadBalancer::<RoundRobin>::try_from_iter(setting.upstreams())
            .expect("Failed to build load balancer");

        let hc = TcpHealthCheck::new();
        lb.set_health_check(hc);
        lb.health_check_frequency = Some(Duration::from_secs(1));

        let name = format!("health check for {}", k.to_string());
        let background = background_service(&name, lb);
        let lb = background.task();

        bg_services.push(background);
        proxy_entries.push((k.clone(), lb));
    }

    let proxy = Proxy::new(proxy_entries.into_iter());
    
    let mut server = Server::new(None).expect("Failed to create server");
    server.bootstrap();

    let addr = format!("0.0.0.0:{}", config.port());
    let mut proxy_service = http_proxy_service(&server.configuration, proxy);
    if let Some((cert_path, key_path)) = config.tls_config() {
        let cert_path = cert_path.to_str().expect("Invalid cert path");
        let key_path = key_path.to_str().expect("Invalid key path");
        proxy_service.add_tls(&addr, cert_path, key_path)
            .expect("Failed to add TLS to proxy service");

        info!("Listening on {} with tls", addr);
    } else {
        proxy_service.add_tcp(&format!("0.0.0.0:{}", config.port()));
        info!("Listening on {} with tls", addr);
    }

    server.add_service(proxy_service);
    for s in bg_services {
        server.add_service(s);
    }

    server.run_forever();
}