From 9b6769548b6971577a1f6a311625e4ec9f82197b Mon Sep 17 00:00:00 2001
From: Eugene Pankov <e@ajenti.org>
Date: Thu, 28 Apr 2022 17:20:32 -0700
Subject: [PATCH] use poem's new embed endpoints

---
 Cargo.lock                  |  2 +
 warpgate-admin/Cargo.toml   |  2 +-
 warpgate-admin/src/embed.rs | 96 -------------------------------------
 warpgate-admin/src/lib.rs   |  3 +-
 4 files changed, 4 insertions(+), 99 deletions(-)
 delete mode 100644 warpgate-admin/src/embed.rs

diff --git a/Cargo.lock b/Cargo.lock
index 2b06799..8ca7001 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2234,6 +2234,7 @@ dependencies = [
  "cookie",
  "futures-util",
  "headers",
+ "hex",
  "http",
  "httpdate",
  "hyper",
@@ -2247,6 +2248,7 @@ dependencies = [
  "priority-queue",
  "rand",
  "regex",
+ "rust-embed",
  "rustls-pemfile",
  "serde",
  "serde_json",
diff --git a/warpgate-admin/Cargo.toml b/warpgate-admin/Cargo.toml
index 7d694d8..60af97f 100644
--- a/warpgate-admin/Cargo.toml
+++ b/warpgate-admin/Cargo.toml
@@ -12,7 +12,7 @@ chrono = "0.4"
 futures = "0.3"
 hex = "0.4"
 mime_guess = {version = "2.0", default_features = false}
-poem = {version = "^1.3.24", features = ["cookie", "session", "anyhow", "rustls", "websocket"]}
+poem = {version = "^1.3.24", features = ["cookie", "session", "anyhow", "rustls", "websocket", "embed"]}
 poem-openapi = {version = "^1.3.24", features = ["swagger-ui", "chrono", "uuid", "static-files"]}
 russh-keys = {version = "0.22.0-beta.1", features = ["openssl"]}
 rust-embed = "6.3"
diff --git a/warpgate-admin/src/embed.rs b/warpgate-admin/src/embed.rs
deleted file mode 100644
index 4b3d4ae..0000000
--- a/warpgate-admin/src/embed.rs
+++ /dev/null
@@ -1,96 +0,0 @@
-//! Usage:
-//!
-//! ```
-//! #[derive(RustEmbed)]
-//! #[folder = "app/dist"]
-//! pub struct Assets;
-//!
-//! Route::new()
-//!     .at("/", EmbeddedFileEndpoint::<Assets>::new("index.html"))
-//!     .nest_no_strip("/assets", EmbeddedFilesEndpoint::<Assets>::new())
-//! ```
-
-use async_trait::async_trait;
-use poem::http::{header, Method, StatusCode};
-use poem::{Endpoint, Request, Response};
-use rust_embed::RustEmbed;
-use std::marker::PhantomData;
-
-pub struct EmbeddedFileEndpoint<E: RustEmbed + Send + Sync> {
-    _embed: PhantomData<E>,
-    path: String,
-}
-
-impl<E: RustEmbed + Send + Sync> EmbeddedFileEndpoint<E> {
-    pub fn new(path: &str) -> Self {
-        EmbeddedFileEndpoint {
-            _embed: PhantomData,
-            path: path.to_owned(),
-        }
-    }
-}
-
-#[async_trait]
-impl<E: RustEmbed + Send + Sync> Endpoint for EmbeddedFileEndpoint<E> {
-    type Output = Response;
-
-    async fn call(&self, req: Request) -> Result<Self::Output, poem::Error> {
-        if req.method() != Method::GET {
-            return Err(StatusCode::METHOD_NOT_ALLOWED.into());
-        }
-
-        match E::get(&self.path) {
-            Some(content) => {
-                let hash = hex::encode(content.metadata.sha256_hash());
-                if req
-                    .headers()
-                    .get(header::IF_NONE_MATCH)
-                    .map(|etag| etag.to_str().unwrap_or("000000").eq(&hash))
-                    .unwrap_or(false)
-                {
-                    return Err(StatusCode::NOT_MODIFIED.into());
-                }
-
-                // otherwise, return 200 with etag hash
-                let body: Vec<u8> = content.data.into();
-                let mime = mime_guess::from_path(&self.path).first_or_octet_stream();
-                Ok(Response::builder()
-                    .header(header::CONTENT_TYPE, mime.as_ref())
-                    .header(header::ETAG, hash)
-                    .body(body))
-            }
-            None => Err(StatusCode::NOT_FOUND.into()),
-        }
-    }
-}
-
-pub struct EmbeddedFilesEndpoint<E: RustEmbed + Send + Sync> {
-    _embed: PhantomData<E>,
-}
-
-impl<E: RustEmbed + Send + Sync> EmbeddedFilesEndpoint<E> {
-    pub fn new() -> Self {
-        EmbeddedFilesEndpoint {
-            _embed: PhantomData,
-        }
-    }
-}
-
-#[async_trait]
-impl<E: RustEmbed + Send + Sync> Endpoint for EmbeddedFilesEndpoint<E> {
-    type Output = Response;
-
-    async fn call(&self, req: Request) -> Result<Self::Output, poem::Error> {
-        let mut path = req
-            .uri()
-            .path()
-            .trim_start_matches('/')
-            .trim_end_matches('/')
-            .to_string();
-        if path.is_empty() {
-            path = "index.html".to_string();
-        }
-        let path = path.as_ref();
-        EmbeddedFileEndpoint::<E>::new(path).call(req).await
-    }
-}
diff --git a/warpgate-admin/src/lib.rs b/warpgate-admin/src/lib.rs
index a24d576..da55a7c 100644
--- a/warpgate-admin/src/lib.rs
+++ b/warpgate-admin/src/lib.rs
@@ -1,9 +1,8 @@
 #![feature(decl_macro, proc_macro_hygiene, let_else)]
 mod api;
-mod embed;
 mod helpers;
-use crate::embed::{EmbeddedFileEndpoint, EmbeddedFilesEndpoint};
 use anyhow::{Context, Result};
+use poem::endpoint::{EmbeddedFilesEndpoint, EmbeddedFileEndpoint};
 use poem::listener::{Listener, RustlsConfig, TcpListener};
 use poem::middleware::{AddData, SetHeader};
 use poem::session::{CookieConfig, MemoryStorage, ServerSession};