From 6a953c100f1c7a53bd209a58a9984c3027cae152 Mon Sep 17 00:00:00 2001
From: NexVeridian <NexVeridian@gmail.com>
Date: Sun, 23 Jul 2023 01:48:43 +0000
Subject: [PATCH] 1.1.2

---
 .devcontainer/devcontainer.json |  2 +-
 Cargo.toml                      |  5 +--
 docker-compose.dev.yml          | 17 ++++++++++
 docker-compose.yml              |  5 ---
 src/main.rs                     | 60 ++++++++++++++++++++++++++-------
 5 files changed, 69 insertions(+), 20 deletions(-)
 mode change 100644 => 100755 .devcontainer/devcontainer.json
 create mode 100644 docker-compose.dev.yml

diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
old mode 100644
new mode 100755
index 3c4eea0..7aa0a93
--- a/.devcontainer/devcontainer.json
+++ b/.devcontainer/devcontainer.json
@@ -57,7 +57,7 @@
 			"extensions": [
 				"vadimcn.vscode-lldb",
 				"serayuzgur.crates",
-				"bungcip.better-toml",
+				"tamasfe.even-better-toml",
 				"rust-lang.rust-analyzer",
 				"mutantdino.resourcemonitor",
 				"christian-kohler.path-intellisense",
diff --git a/Cargo.toml b/Cargo.toml
index 4cd37b4..40c822d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-polars = { version = "0.28", features = [
+polars = { version = "0.32", features = [
 	"lazy",
 	"strings",
 	"parquet",
@@ -22,4 +22,5 @@ chrono = { version = "0.4", features = ["serde"] }
 glob = { version = "0.3" }
 strum_macros = "0.25"
 tower = { version = "0.4", features = ["limit", "buffer"] }
-tower-http = { version = "0.4", features = ["cors"] }
+tower-http = { version = "0.4", features = ["compression-zstd", "cors"] }
+tower_governor = "0.1"
diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml
new file mode 100644
index 0000000..f5349ec
--- /dev/null
+++ b/docker-compose.dev.yml
@@ -0,0 +1,17 @@
+version: "3"
+services:
+    ark-invest-api-rust:
+        image: ark-invest-api-rust
+        build:
+            context: .
+            args:
+                DOCKER_BUILDKIT: 1
+        container_name: ark-invest-api-rust
+        restart: unless-stopped
+        volumes:
+            - ./data:/ark-invest-api-rust/data
+        ports:
+            - "3000:3000"
+
+volumes:
+    data:
diff --git a/docker-compose.yml b/docker-compose.yml
index 23d94b2..25cf1f5 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -2,11 +2,6 @@ version: "3"
 services:
     ark-invest-api-rust:
         image: ghcr.io/NexVeridian/ark-invest-api-rust:latest
-        # image: ark-invest-api-rust
-        build:
-            context: .
-            args:
-                DOCKER_BUILDKIT: 1
         container_name: ark-invest-api-rust
         restart: unless-stopped
         volumes:
diff --git a/src/main.rs b/src/main.rs
index 8d6f471..5bcf8ce 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -14,7 +14,13 @@ use axum::{
 };
 use std::{net::SocketAddr, time::Duration};
 use tower::{buffer::BufferLayer, limit::RateLimitLayer, ServiceBuilder};
-use tower_http::cors::{Any, CorsLayer};
+use tower_governor::{
+    governor::GovernorConfigBuilder, key_extractor::SmartIpKeyExtractor, GovernorLayer,
+};
+use tower_http::{
+    compression::CompressionLayer,
+    cors::{Any, CorsLayer},
+};
 
 mod routes;
 
@@ -33,18 +39,40 @@ fn description_date(op: TransformOperation) -> TransformOperation {
 
 #[tokio::main]
 async fn main() {
-    let rate_limit = |req_per_sec: u64| {
+    let error_handler = || {
+        ServiceBuilder::new().layer(HandleErrorLayer::new(|err: BoxError| async move {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Unhandled error: {}", err),
+            )
+        }))
+    };
+
+    let rate_limit_global = |req_per_sec: u64| {
         ServiceBuilder::new()
-            .layer(HandleErrorLayer::new(|err: BoxError| async move {
-                (
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    format!("Unhandled error: {}", err),
-                )
-            }))
+            .layer(error_handler())
             .layer(BufferLayer::new(1024))
             .layer(RateLimitLayer::new(req_per_sec, Duration::from_secs(1)))
     };
 
+    let rate_limit_ip = || {
+        let config = Box::new(
+            GovernorConfigBuilder::default()
+                .per_second(2)
+                .burst_size(10)
+                .use_headers()
+                .key_extractor(SmartIpKeyExtractor)
+                .finish()
+                .unwrap(),
+        );
+
+        ServiceBuilder::new()
+            .layer(error_handler())
+            .layer(GovernorLayer {
+                config: Box::leak(config),
+            })
+    };
+
     let cors = || {
         ServiceBuilder::new().layer(
             CorsLayer::new()
@@ -62,8 +90,12 @@ async fn main() {
                 description_date(o)
             }),
         )
-        .layer(rate_limit(5))
-        .layer(cors())
+        .layer(
+            rate_limit_global(500)
+                .layer(rate_limit_ip())
+                .layer(cors())
+                .layer(CompressionLayer::new().zstd(true)),
+        )
         .api_route(
             "/ark_holdings",
             get_with(routes::ark_holdings, |mut o| {
@@ -71,8 +103,12 @@ async fn main() {
                 description_date(o)
             }),
         )
-        .layer(rate_limit(20))
-        .layer(cors())
+        .layer(
+            rate_limit_global(200)
+                .layer(rate_limit_ip())
+                .layer(cors())
+                .layer(CompressionLayer::new().zstd(true)),
+        )
         .route("/api.json", get(serve_api));
 
     let mut api = OpenApi {