add file routing support

Cargo.toml

@@ -2,17 +2,19 @@
 authors = ["yao <[email protected]>"]
 edition = "2021"
 name = "mock-server"
-version = "1.1.0"
+version = "1.2.0"
 
 [dependencies]
 axum = "0.5.13"
+axum-streams = "0.9.1"
 clap = {version = "3.2.17", features = ["derive"]}
 once_cell = "1.13.0"
 rand = "0.8.5"
 regex = "1.6.0"
 serde = {version = "1.0.141", features = ["derive"]}
 serde_json = "1.0.82"
 tokio = {version = "1.20.1", features = ["full"]}
+tokio-util = "0.7.9"
 tower = "0.4.13"
 tower-http = {version = "0.3.4", features = ["fs", "trace", "cors"]}
 tracing = "0.1.36"

src/app.rs

@@ -17,6 +17,7 @@ use tower_http::services::ServeDir;
 
 const DATA_QUERY_TPL: &str = "/api/([^/]*)$";
 const DATA_ID_TPL: &str = "/api/([^/]*)/([^/]*)$";
+const FILE_ID_TPL: &str = "/file/([^/]*)$";
 
 pub fn create(_: &Database) -> Router {
     let app = Router::new()
@@ -28,6 +29,7 @@ pub fn create(_: &Database) -> Router {
                 .put(put_data)
                 .delete(delete_data),
         )
+        .route("/file/:id", get(get_file).post(get_file))
         .fallback(get_service(ServeDir::new("static")).handle_error(handle_error));
 
     app
@@ -101,8 +103,10 @@ pub fn proxy(db: &Database) -> Router {
                         return Ok(res);
                     }
 
-                    let re = regex::Regex::new(DATA_ID_TPL).unwrap();
-                    if let Some(cap) = re.captures(&routing_value.to) {
+                    if let Some(cap) = regex::Regex::new(DATA_ID_TPL)
+                        .unwrap()
+                        .captures(&routing_value.to)
+                    {
                         let data = cap.get(1).unwrap().as_str();
                         let id = cap.get(2).unwrap().as_str();
                         if !data.starts_with(":") {
@@ -124,6 +128,24 @@ pub fn proxy(db: &Database) -> Router {
                         return Ok(res);
                     }
 
+                    if let Some(cap) = regex::Regex::new(FILE_ID_TPL)
+                        .unwrap()
+                        .captures(&routing_value.to)
+                    {
+                        let id = cap.get(1).unwrap().as_str();
+                        if !id.starts_with(":") {
+                            path.insert("id".to_string(), id.to_string());
+                        }
+                        let res = match method {
+                            &Method::GET => get_file(path, db).await.into_response(),
+                            &Method::POST => get_file(path, db).await.into_response(),
+                            _ => (StatusCode::METHOD_NOT_ALLOWED, "method not support")
+                                .into_response(),
+                        };
+
+                        return Ok(res);
+                    }
+
                     let res = (StatusCode::BAD_REQUEST, "bad request").into_response();
                     Ok(res)
                 }),

src/db.rs

@@ -8,10 +8,12 @@ use serde_json::Value;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::sync::RwLock;
+use tracing::error;
 
 #[derive(Clone)]
 pub struct Database {
     collections: Arc<RwLock<HashMap<String, Value>>>,
+    files: HashMap<String, String>,
     config: Arc<DataConfig>,
 }
 
@@ -27,6 +29,7 @@ impl<'a> Database {
         Database {
             config: Arc::new(DataConfig::default()),
             collections: Arc::new(RwLock::new(HashMap::new())),
+            files: HashMap::new(),
         }
     }
 
@@ -47,6 +50,18 @@ impl<'a> Database {
                     collections.insert(name, data);
                 }
 
+                if let Some(file) = json.get("file") {
+                    if let Some(file_list) = file.as_object() {
+                        for (key, value) in file_list {
+                            if let Some(path) = value.as_str() {
+                                self.files.insert(key.to_string(), path.to_string());
+                            } else {
+                                error!("file.{} must be string type", key);
+                            }
+                        }
+                    }
+                }
+
                 self.config = Arc::new(DataConfig::new(&config));
                 DATA_CONFIG.set(self.config.clone()).unwrap();
             }
@@ -237,4 +252,15 @@ impl<'a> Database {
 
         Ok(Value::Bool(true))
     }
+
+    pub fn get_file(
+        &mut self,
+        path_map: &HashMap<std::string::String, std::string::String>,
+    ) -> Result<String, String> {
+        let file_id = path_map.get("id").unwrap();
+        match self.files.get(file_id) {
+            Some(v) => Ok(v.clone()),
+            None => Err(format!("not found item by id {}", file_id)),
+        }
+    }
 }

src/service.rs

@@ -1,12 +1,17 @@
 use crate::models::Wrapper;
 use crate::Database;
 use crate::HashMap;
+use axum::body::StreamBody;
 use axum::extract::Path;
 use axum::extract::Query;
+use axum::http::header;
+use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::Extension;
 use axum::Json;
+use regex::Regex;
 use serde_json::Value;
+use tokio_util::io::ReaderStream;
 
 macro_rules! wrapping {
     ($result: expr) => {
@@ -59,3 +64,42 @@ pub async fn delete_data(
 ) -> impl IntoResponse {
     wrapping!(db.delete_data(&query), wrap)
 }
+
+pub async fn get_file(
+    Path(query): Path<HashMap<String, String>>,
+    Extension(mut db): Extension<Database>,
+) -> impl IntoResponse {
+    match db.get_file(&query) {
+        Ok(ref path) => {
+            let file = match tokio::fs::File::open(path).await {
+                Ok(file) => file,
+                Err(err) => {
+                    return Err((StatusCode::NOT_FOUND, format!("File not found: {}", err)))
+                }
+            };
+            let stream = ReaderStream::new(file);
+            let body = StreamBody::new(stream);
+
+            let mut content_type = "application/octet-stream".to_string();
+            let file_path = std::path::Path::new(path);
+            if let Some(ext) = file_path.extension().and_then(|s| s.to_str()) {
+                if Regex::new("jpg|png|gif|jpeg").unwrap().is_match(ext) {
+                    content_type = format!("image/{}; charset=utf-8", ext);
+                }
+
+                if Regex::new("json").unwrap().is_match(ext) {
+                    content_type = format!("application/{}; charset=utf-8", ext);
+                }
+
+                if Regex::new("txt").unwrap().is_match(ext) {
+                    content_type = format!("text/plain; charset=utf-8");
+                }
+            }
+
+            let headers = [(header::CONTENT_TYPE, content_type)];
+            Ok((headers, body))
+        }
+
+        Err(err) => return Err((StatusCode::NOT_FOUND, err)),
+    }
+}

static/db.json

@@ -3,8 +3,8 @@
   "config": {
 
     "routing": {
-      "/fbpm_demo/api/expense/listByPageTask": {
-        "to": "/api/data2"
+      "/my_image/:id": {
+        "to": "/file/:id"
       },
       "/a/b/c/:id": { 
         "to": "/api/data/:id"
@@ -96,5 +96,11 @@
       "age|1-100": 20,
       "id|+1": 1
     }]
+  },
+
+  "file": {
+    "image1": "./static/image1.jpg",
+    "text1": "./static/text1.txt",
+    "json1": "./static/json1.json"
   }
 }

static/json1.json

@@ -0,0 +1,3 @@
+{
+    "name": "Yao"
+}

static/text1.txt

@@ -0,0 +1 @@
+hello world