From 8f2747edfd0160abd81b2e6b18f883b5cb33b68f Mon Sep 17 00:00:00 2001 From: shenjack <3695888@qq.com> Date: Fri, 26 Jan 2024 14:32:13 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=E4=B8=80=E6=B3=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/config.rs | 63 +++++++++++++++++++++++++++++++++++++++++++++++++++ src/log.rs | 5 +--- src/main.rs | 5 ++-- src/serve.rs | 43 +++++++++++++++++++++++++++++++++++ src/utils.rs | 27 ++++++++++++++-------- 5 files changed, 127 insertions(+), 16 deletions(-) create mode 100644 src/config.rs create mode 100644 src/serve.rs diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..9c38dee --- /dev/null +++ b/src/config.rs @@ -0,0 +1,63 @@ + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ClusterByoc { + http, + https, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct Config { + pub cluster_port: u32, + pub cluster_id: u32, + pub cluster_secret: String, + /// CLUSTER_BYOC + pub cluster_byoc: ClusterByoc, + pub no_demaon: bool, + /// DISABLE_ACCESS_LOG + pub disable_access_log: bool, + /// FORCE_NOOPEN + pub force_noopen: bool, +} + +impl Config { + pub fn new_from_env() -> Self { + let cluster_port = std::env::var("CLUSTER_PORT") + .unwrap_or_else(|_| "8080".to_string()) + .parse::() + .expect("CLUSTER_PORT must be a number"); + let cluster_id = std::env::var("CLUSTER_ID") + .unwrap_or_else(|_| "1".to_string()) + .parse::() + .expect("CLUSTER_ID must be a number"); + let cluster_secret = std::env::var("CLUSTER_SECRET").expect("CLUSTER_SECRET must be set"); + let cluster_byoc = match std::env::var("CLUSTER_BYOC") + .unwrap_or_else(|_| "http".to_string()) + .as_str() + { + "http" => ClusterByoc::http, + "https" => ClusterByoc::https, + _ => panic!("CLUSTER_BYOC must be http or https"), + }; + let no_demaon = std::env::var("NO_DAEMON") + .unwrap_or_else(|_| "false".to_string()) + .parse::() + .expect("NO_DAEMON must be true or false"); + let disable_access_log = std::env::var("DISABLE_ACCESS_LOG") + .unwrap_or_else(|_| "false".to_string()) + .parse::() + .expect("DISABLE_ACCESS_LOG must be true or false"); + let force_noopen = std::env::var("FORCE_NOOPEN") + .unwrap_or_else(|_| "false".to_string()) + .parse::() + .expect("FORCE_NOOPEN must be true or false"); + Self { + cluster_port, + cluster_id, + cluster_secret, + cluster_byoc, + no_demaon, + disable_access_log, + force_noopen, + } + } +} diff --git a/src/log.rs b/src/log.rs index 99d2950..ed79b27 100644 --- a/src/log.rs +++ b/src/log.rs @@ -1,4 +1 @@ - -pub fn init_logging() { - -} +pub fn init_logging() {} diff --git a/src/main.rs b/src/main.rs index e287d65..7568fa1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,5 @@ mod log; +mod serve; mod utils; -fn main() { - -} +fn main() {} diff --git a/src/serve.rs b/src/serve.rs new file mode 100644 index 0000000..18743da --- /dev/null +++ b/src/serve.rs @@ -0,0 +1,43 @@ +use axum::{extract::Path, http::header::HeaderMap, http::StatusCode, response::IntoResponse}; + +/// ```ts +/// import express from 'express' +/// +/// const router = express.Router() +/// +/// const MeasureRoute = router +/// +/// router.get('/:size(\\d+)', (req, res) => { +/// if (req.get('x-openbmclapi-secret') !== process.env.CLUSTER_SECRET) { +/// return res.sendStatus(403) +/// } +/// const size = parseInt(req.params.size, 10) +/// if (isNaN(size) || size > 200) return res.sendStatus(400) +/// const buffer = Buffer.alloc(1024 * 1024, '0066ccff', 'hex') +/// res.set('content-length', (size * 1024 * 1024).toString()) +/// for (let i = 0; i < size; i++) { +/// res.write(buffer) +/// } +/// }) +/// +/// export default MeasureRoute +/// ``` +async fn measure(header: HeaderMap, Path(size): Path) -> impl IntoResponse { + match header.get("x-openbmclapi-secret") { + Some(secret) => { + if secret != "secret" { + return (StatusCode::FORBIDDEN, Vec::new()); + } + if size > 200 { + return (StatusCode::BAD_REQUEST, Vec::new()); + } + let buffer: Vec = vec![0x00, 0x66, 0xcc, 0xff]; + let mut response: Vec = Vec::new(); + for _ in 0..size { + response.extend(&buffer); + } + return (StatusCode::OK, response); + } + None => return (StatusCode::FORBIDDEN, Vec::new()), + } +} diff --git a/src/utils.rs b/src/utils.rs index 6725786..50aefc8 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,9 +1,9 @@ +use md5::{Digest, Md5}; +use sha1::{Digest as Sha1Digest, Sha1}; use std::path::PathBuf; -use md5::{Md5, Digest}; -use sha1::{Sha1, Digest as Sha1Digest}; /// import {join} from 'path' -/// +/// /// export function hashToFilename(hash: string): string { /// // eslint-disable-next-line @typescript-eslint/no-magic-numbers /// return join(hash.substring(0, 2), hash) @@ -16,7 +16,7 @@ pub fn hash_to_filename(hash: &str) -> PathBuf { } /// import {createHash, Hash} from 'crypto' -/// +/// /// export function validateFile(buffer: Buffer, checkSum: string): boolean { /// let hash: Hash /// if (checkSum.length === 32) { @@ -35,7 +35,7 @@ pub fn validate_file(buffer: &[u8], check_sum: &str) -> bool { let result = hasher.finalize(); let result_str = format!("{:x}", result); result_str == check_sum - }, + } _ => { let mut hasher = Sha1::new(); hasher.update(buffer); @@ -48,11 +48,20 @@ pub fn validate_file(buffer: &[u8], check_sum: &str) -> bool { #[test] fn test_hash_to_filename() { - assert_eq!(hash_to_filename("1234567890abcdef"), PathBuf::from("12/1234567890abcdef")); + assert_eq!( + hash_to_filename("1234567890abcdef"), + PathBuf::from("12/1234567890abcdef") + ); } #[test] fn test_validate_file() { - assert_eq!(validate_file(b"hello", "5d41402abc4b2a76b9719d911017c592"), true); - assert_eq!(validate_file(b"hello", "5d41402abc4b2a76b9719d911017c593"), false); -} \ No newline at end of file + assert_eq!( + validate_file(b"hello", "5d41402abc4b2a76b9719d911017c592"), + true + ); + assert_eq!( + validate_file(b"hello", "5d41402abc4b2a76b9719d911017c593"), + false + ); +}