From 19564e2df22525222e92f688e423753c5eda2a5f Mon Sep 17 00:00:00 2001
From: theMackabu <theMackabu@gmail.com>
Date: Sun, 3 Mar 2024 15:09:44 -0800
Subject: [PATCH] add proxy function

---
 Cargo.lock         |  1 +
 Cargo.toml         |  1 +
 app.routes         |  5 ++++
 src/database/kv.rs |  3 +++
 src/main.rs        | 58 +++++++++++++++++++++++++++++-----------------
 5 files changed, 47 insertions(+), 21 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4734d0d..d2c429a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2388,6 +2388,7 @@ dependencies = [
  "home",
  "lazy_static",
  "macros-rs",
+ "mime",
  "mongodb",
  "peg",
  "pickledb",
diff --git a/Cargo.toml b/Cargo.toml
index 05c93d8..b3530d6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -32,3 +32,4 @@ mongodb = { version = "2.8.0", features = ["sync"], default-features = false }
 pickledb = { version = "0.5.1", features = ["json", "bincode", "cbor", "yaml"] }
 redis = "0.24.0"
 askama = "0.12.1"
+mime = "0.3.17"
diff --git a/app.routes b/app.routes
index 7aeb3c8..9ff2134 100644
--- a/app.routes
+++ b/app.routes
@@ -84,6 +84,11 @@ name(id, name) {
    json(#{name: name, id: id})
 }
 
+#[route("/proxy/{arg}")]
+fetch_proxy(arg) {
+   proxy(`https://internal.themackabu.dev/${ arg }`)
+}
+
 example/json() {
    let res = http::get("https://httpbin.org/json");
    let body = #{ 
diff --git a/src/database/kv.rs b/src/database/kv.rs
index 28f2343..8b35876 100644
--- a/src/database/kv.rs
+++ b/src/database/kv.rs
@@ -13,3 +13,6 @@ pub fn load(path: String) -> PickleDb {
 
     PickleDb::new(path, PickleDbDumpPolicy::AutoDump, method)
 }
+
+// use load to not erase db on every load
+// add .iter() method
diff --git a/src/main.rs b/src/main.rs
index 03ea491..8e36c38 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,6 +7,7 @@ use askama::Template;
 use config::structs::Config;
 use lazy_static::lazy_static;
 use macros_rs::{crashln, str, string, ternary};
+use mime::Mime;
 use pickledb::PickleDb;
 use redis::{Client as RedisClient, Commands};
 use regex::{Captures, Error, Regex};
@@ -77,6 +78,23 @@ pub fn response(data: String, content_type: String, status_code: i64) -> (String
     (data, content_type, helpers::convert_status(status_code))
 }
 
+pub fn proxy(url: String) -> (String, ContentType, StatusCode) {
+    let client = ReqwestClient::new();
+    let response = match client.get(url).send() {
+        Ok(res) => res,
+        Err(err) => return (err.to_string(), ContentType::plaintext(), StatusCode::GATEWAY_TIMEOUT),
+    };
+
+    let status = response.status();
+    let content_type = response.headers().get("Content-Type").unwrap().to_str().unwrap_or("text/plain").parse::<Mime>().unwrap();
+
+    if status.is_success() {
+        (response.text().unwrap(), ContentType(content_type), status)
+    } else {
+        (response.text().unwrap(), ContentType(content_type), status)
+    }
+}
+
 fn match_route(route_template: &str, placeholders: &[&str], url: &str) -> Option<Vec<String>> {
     let mut matched_placeholders = Vec::new();
 
@@ -671,18 +689,17 @@ mod http {
 
     pub fn get(url: String) -> Http {
         let client = ReqwestClient::new();
-        let response =
-            match client.get(url).send() {
-                Ok(res) => res,
-                Err(err) => {
-                    return Http {
-                        length: Some(0),
-                        status: 0,
-                        err: Some(err.to_string()),
-                        body: None,
-                    }
+        let response = match client.get(url).send() {
+            Ok(res) => res,
+            Err(err) => {
+                return Http {
+                    length: Some(0),
+                    status: 0,
+                    err: Some(err.to_string()),
+                    body: None,
                 }
-            };
+            }
+        };
 
         if response.status().is_success() {
             Http {
@@ -704,11 +721,10 @@ mod http {
     pub fn post(url: String, data: Map) -> Http {
         let client = ReqwestClient::new();
 
-        let data =
-            match serde_json::to_string(&data) {
-                Ok(result) => result,
-                Err(err) => err.to_string(),
-            };
+        let data = match serde_json::to_string(&data) {
+            Ok(result) => result,
+            Err(err) => err.to_string(),
+        };
 
         let response = match client.post(url).body(data).send() {
             Ok(res) => res,
@@ -874,6 +890,7 @@ async fn handler(url: Path<String>, req: HttpRequest, config: Data<Config>) -> i
     scope.push("request", request.to_dynamic());
 
     engine
+        .register_fn("proxy", proxy)
         .register_fn("cwd", file::cwd)
         .register_fn("response", response)
         .register_fn("text", default::text)
@@ -912,11 +929,10 @@ async fn handler(url: Path<String>, req: HttpRequest, config: Data<Config>) -> i
 
             if args != "" {
                 let r_path = Regex::new(r"(?m)_arg_(\w+)").unwrap();
-                let key =
-                    r_path.replace_all(&path, |captures: &regex::Captures| {
-                        let key = captures.get(1).map_or("", |m| m.as_str());
-                        format!("{{{key}}}")
-                    });
+                let key = r_path.replace_all(&path, |captures: &regex::Captures| {
+                    let key = captures.get(1).map_or("", |m| m.as_str());
+                    format!("{{{key}}}")
+                });
 
                 routes.insert(string!(key), args.split(",").map(|s| s.to_string().replace(" ", "")).collect());
                 format!("fmt_{path}({args})")
-- 
GitLab