Skip to content

Commit

Permalink
feat: support ip limit and key auth proxy plugin
Browse files Browse the repository at this point in the history
  • Loading branch information
vicanso committed Apr 17, 2024
1 parent 47a76a5 commit cd1ebf5
Show file tree
Hide file tree
Showing 22 changed files with 394 additions and 26 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ hostname = "0.3.1"
http = "1.1.0"
humantime = "2.1.0"
humantime-serde = "1.1.1"
ipnet = "2.9.0"
log = "0.4.21"
memory-stats = { version = "1.1.0", features = ["always_use_statm"] }
mime_guess = "2.0.4"
Expand Down
7 changes: 4 additions & 3 deletions TODO.md
Original file line number Diff line number Diff line change
@@ -1,13 +1,14 @@
# TODO

- [ ] request id proxy plugin
- [ ] allow deny ip proxy plugin
- [ ] support etcd or other storage for config
- [ ] better error handler
- [ ] log rotate
- [ ] tls cert auto update
- [ ] support validate config before save(web)
- [ ] auto reload config and restart
- [x] allow none upstream for location
- [x] allow deny ip proxy plugin
- [x] auto reload config and restart
- [x] request id proxy plugin
- [x] support plugin for proxy and response
- [x] authentication for admin page
- [x] custom error for pingora error
Expand Down
8 changes: 4 additions & 4 deletions benches/bench.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ fn bench_location_filter(c: &mut Criterion) {
let lo = Location::new(
"",
&LocationConf {
upstream: upstream_name.to_string(),
upstream: Some(upstream_name.to_string()),
path: Some("/api".to_string()),
..Default::default()
},
Expand All @@ -123,7 +123,7 @@ fn bench_location_filter(c: &mut Criterion) {
let lo = Location::new(
"",
&LocationConf {
upstream: upstream_name.to_string(),
upstream: Some(upstream_name.to_string()),
path: Some("~/api".to_string()),
..Default::default()
},
Expand All @@ -139,7 +139,7 @@ fn bench_location_filter(c: &mut Criterion) {
let lo = Location::new(
"",
&LocationConf {
upstream: upstream_name.to_string(),
upstream: Some(upstream_name.to_string()),
path: Some("=/api".to_string()),
..Default::default()
},
Expand Down Expand Up @@ -170,7 +170,7 @@ fn bench_location_rewrite_path(c: &mut Criterion) {
let lo = Location::new(
"",
&LocationConf {
upstream: upstream_name.to_string(),
upstream: Some(upstream_name.to_string()),
rewrite: Some("^/users/(.*)$ /$1".to_string()),
..Default::default()
},
Expand Down
9 changes: 6 additions & 3 deletions src/config/load.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,8 @@ pub enum ProxyPluginCategory {
Directory,
Mock,
RequestId,
IpLimit,
KeyAuth,
}

#[derive(PartialEq, Debug, Default, Deserialize_repr, Clone, Copy, Serialize_repr)]
Expand Down Expand Up @@ -148,7 +150,7 @@ impl UpstreamConf {

#[derive(Debug, Default, Deserialize, Clone, Serialize)]
pub struct LocationConf {
pub upstream: String,
pub upstream: Option<String>,
pub path: Option<String>,
pub host: Option<String>,
pub proxy_headers: Option<Vec<String>>,
Expand Down Expand Up @@ -181,9 +183,10 @@ impl LocationConf {
};
validate(&self.proxy_headers)?;
validate(&self.headers)?;
if !upstream_names.contains(&self.upstream) {
let upstream = self.upstream.clone().unwrap_or_default();
if !upstream.is_empty() && !upstream_names.contains(&upstream) {
return Err(Error::Invalid {
message: format!("{} upstream is not found(location:{name})", self.upstream),
message: format!("{upstream} upstream is not found(location:{name})"),
});
}

Expand Down
16 changes: 16 additions & 0 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,15 @@ use config::{PingapConf, ProxyPluginCategory, ProxyPluginConf};
use log::{error, info, Level};
use pingora::server;
use pingora::server::configuration::Opt;
use pingora::services::background::background_service;
use proxy::{Server, ServerConf};
use state::get_start_time;
use std::error::Error;
use std::io::Write;
use std::sync::Arc;

use crate::state::AutoRestart;

mod config;
mod http_extra;
mod plugin;
Expand Down Expand Up @@ -56,6 +59,9 @@ struct Args {
/// Admin server adddr
#[arg(long)]
admin: Option<String>,
/// Whether this server should try to auto restart
#[arg(long)]
autorestart: bool,
}

fn new_server_conf(args: &Args, conf: &PingapConf) -> server::configuration::ServerConf {
Expand Down Expand Up @@ -162,6 +168,12 @@ fn run() -> Result<(), Box<dyn Error>> {
if let Some(log) = &args.log {
new_args.push(format!("--log={log}"));
}
if let Some(admin) = &args.admin {
new_args.push(format!("--admin={admin}"));
}
if args.autorestart {
new_args.push("--autorestart".to_string());
}
cmd.args = new_args;
state::set_restart_process_command(cmd);
}
Expand Down Expand Up @@ -217,6 +229,10 @@ fn run() -> Result<(), Box<dyn Error>> {
my_server.add_services(services.bg_services);
my_server.add_service(services.lb);
}
if args.autorestart {
my_server.add_service(background_service("Auto Restart", AutoRestart {}));
}

info!("server is running");
let _ = get_start_time();
my_server.run_forever();
Expand Down
3 changes: 2 additions & 1 deletion src/plugin/admin.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ use bytesize::ByteSize;
use hex::encode;
use http::Method;
use http::{header, HeaderValue, StatusCode};
use log::error;
use log::{debug, error};
use memory_stats::memory_stats;
use pingora::http::RequestHeader;
use pingora::proxy::Session;
Expand Down Expand Up @@ -93,6 +93,7 @@ pub struct AdminServe {
}
impl AdminServe {
pub fn new(value: &str, proxy_step: ProxyPluginStep) -> Result<Self> {
debug!("new admin server proxy plugin, {value}, {proxy_step:?}");
let arr: Vec<&str> = value.split(' ').collect();
let mut authorization = "".to_string();
if arr.len() >= 2 {
Expand Down
3 changes: 3 additions & 0 deletions src/plugin/compression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ use crate::config::{ProxyPluginCategory, ProxyPluginStep};
use crate::state::State;
use async_trait::async_trait;
use http::HeaderValue;
use log::debug;
use once_cell::sync::Lazy;
use pingora::proxy::Session;

Expand All @@ -35,6 +36,8 @@ pub struct Compression {

impl Compression {
pub fn new(value: &str, proxy_step: ProxyPluginStep) -> Result<Self> {
debug!("new compresson proxy plugin, {value}, {proxy_step:?}");

let mut levels: [u32; 3] = [0, 0, 0];
let mut support_compression = false;
for (index, item) in value.split(' ').enumerate() {
Expand Down
3 changes: 2 additions & 1 deletion src/plugin/directory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ use crate::util;
use async_trait::async_trait;
use bytes::Bytes;
use http::{header, HeaderValue, StatusCode};
use log::error;
use log::{debug, error};
use pingora::proxy::Session;
use std::fs::Metadata;
use std::os::unix::fs::MetadataExt;
Expand Down Expand Up @@ -92,6 +92,7 @@ fn get_cacheable_and_headers_from_meta(
impl Directory {
/// Creates a new directory upstream, which will serve static file of directory.
pub fn new(value: &str, proxy_step: ProxyPluginStep) -> Result<Self> {
debug!("new serve static file proxy plugin, {value}, {proxy_step:?}");
let mut new_path = value.to_string();
let mut chunk_size = None;
let mut max_age = None;
Expand Down
100 changes: 100 additions & 0 deletions src/plugin/ip_limit.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
// Copyright 2024 Tree xie.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use super::ProxyPlugin;
use super::Result;
use crate::config::ProxyPluginCategory;
use crate::config::ProxyPluginStep;
use crate::state::State;
use crate::util;
use async_trait::async_trait;
use ipnet::IpNet;
use log::debug;
use pingora::proxy::Session;
use std::net::IpAddr;
use std::str::FromStr;

pub struct IpLimit {
proxy_step: ProxyPluginStep,
ip_net_list: Vec<IpNet>,
ip_list: Vec<String>,
category: u8,
}

impl IpLimit {
pub fn new(value: &str, proxy_step: ProxyPluginStep) -> Result<Self> {
debug!("new ip limit proxy plugin, {value}, {proxy_step:?}");
let arr: Vec<&str> = value.split(' ').collect();
let ip = arr[0].trim().to_string();
let mut category = 0;
if arr.len() >= 2 {
let v = arr[1].parse::<u8>().unwrap();
if v > 0 {
category = v;
}
}
let mut ip_net_list = vec![];
let mut ip_list = vec![];
for item in ip.split(',') {
if let Ok(value) = IpNet::from_str(item) {
ip_net_list.push(value);
} else {
ip_list.push(item.to_string());
}
}
Ok(Self {
proxy_step,
ip_list,
ip_net_list,
category,
})
}
}

#[async_trait]
impl ProxyPlugin for IpLimit {
#[inline]
fn step(&self) -> ProxyPluginStep {
self.proxy_step
}
#[inline]
fn category(&self) -> ProxyPluginCategory {
ProxyPluginCategory::IpLimit
}
#[inline]
async fn handle(&self, session: &mut Session, ctx: &mut State) -> pingora::Result<bool> {
let ip = if let Some(ip) = &ctx.client_ip {
ip.to_string()
} else {
let ip = util::get_client_ip(session);
ctx.client_ip = Some(ip.clone());
ip
};

let found = if self.ip_list.contains(&ip) {
true
} else {
let addr = ip
.parse::<IpAddr>()
.map_err(|err| util::new_internal_error(400, err.to_string()))?;
self.ip_net_list.iter().any(|item| item.contains(&addr))
};
// deny ip
let allow = if self.category > 0 { !found } else { found };
if !allow {
return Err(util::new_internal_error(403, "Forbidden".to_string()));
}
return Ok(false);
}
}
Loading

0 comments on commit cd1ebf5

Please sign in to comment.