fix(webserver): support resolve all types of repos (#1145)

* fix(webserver): support resolve all types of repos * resolve comment * Update handler.rs --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Meng Zhang <[email protected]>
TabbyML · Jan 1, 2024 · 34f9c7e · 34f9c7e
1 parent 1e9ce3e
commit 34f9c7e
Show file tree

Hide file tree

Showing 7 changed files with 90 additions and 79 deletions.
diff --git a/crates/tabby-common/src/config.rs b/crates/tabby-common/src/config.rs
@@ -42,7 +42,7 @@ impl Config {
     }
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct RepositoryConfig {
     #[serde(skip_serializing_if = "Option::is_none")]
     name: Option<String>,

diff --git a/crates/tabby/src/serve.rs b/crates/tabby/src/serve.rs
@@ -122,7 +122,7 @@ pub async fn main(config: &Config, args: &ServeArgs) {
 
     #[cfg(feature = "ee")]
     let (api, ui) = if args.webserver {
-        tabby_webserver::attach_webserver(api, ui, logger, code).await
+        tabby_webserver::public::attach_webserver(api, ui, logger, code, config).await
     } else {
         let ui = ui.fallback(|| async { axum::response::Redirect::temporary("/swagger-ui") });
         (api, ui)

diff --git a/ee/tabby-webserver/Cargo.toml b/ee/tabby-webserver/Cargo.toml
@@ -34,8 +34,8 @@ tower = { version = "0.4", features = ["util"] }
 tower-http = { version = "0.4.0", features = ["fs", "trace"] }
 tracing.workspace = true
 unicase = "2.7.0"
-validator = { version = "0.16.1", features = ["derive"] }
 uuid.workspace = true
+validator = { version = "0.16.1", features = ["derive"] }
 
 [dev-dependencies]
 assert_matches = "1.5.0"

diff --git a/ee/tabby-webserver/src/handler.rs b/ee/tabby-webserver/src/handler.rs
@@ -8,7 +8,10 @@ use axum::{
 };
 use hyper::Body;
 use juniper_axum::{graphiql, graphql, playground};
-use tabby_common::api::{code::CodeSearch, event::RawEventLogger};
+use tabby_common::{
+    api::{code::CodeSearch, event::RawEventLogger},
+    config::Config,
+};
 
 use crate::{
     hub, repositories,
@@ -22,9 +25,13 @@ pub async fn attach_webserver(
     ui: Router,
     logger: Arc<dyn RawEventLogger>,
     code: Arc<dyn CodeSearch>,
+    config: &Config,
 ) -> (Router, Router) {
     let ctx = create_service_locator(logger, code).await;
     let schema = Arc::new(create_schema());
+    let rs = Arc::new(repositories::ResolveState {
+        repositories: config.repositories.clone(),
+    });
 
     let api = api
         .layer(from_fn_with_state(ctx.clone(), distributed_tabby_layer))
@@ -38,7 +45,10 @@ pub async fn attach_webserver(
             "/hub",
             routing::get(hub::ws_handler).with_state(ctx.clone()),
         )
-        .nest("/repositories", repositories::routes(ctx.clone()));
+        .nest(
+            "/repositories",
+            repositories::routes(rs.clone(), ctx.clone()),
+        );
 
     let ui = ui
         .route("/graphiql", routing::get(graphiql("/graphql", None)))

diff --git a/ee/tabby-webserver/src/lib.rs b/ee/tabby-webserver/src/lib.rs
@@ -13,55 +13,3 @@ pub mod public {
         /* used by examples/update-schema.rs */ schema::create_schema,
     };
 }
-
-use std::sync::Arc;
-
-use axum::{
-    extract::State,
-    http::Request,
-    middleware::{from_fn_with_state, Next},
-    routing, Extension, Router,
-};
-use hyper::Body;
-use juniper_axum::{graphiql, graphql, playground};
-use schema::{Schema, ServiceLocator};
-use service::create_service_locator;
-use tabby_common::api::{code::CodeSearch, event::RawEventLogger};
-
-pub async fn attach_webserver(
-    api: Router,
-    ui: Router,
-    logger: Arc<dyn RawEventLogger>,
-    code: Arc<dyn CodeSearch>,
-) -> (Router, Router) {
-    let ctx = create_service_locator(logger, code).await;
-    let schema = Arc::new(schema::create_schema());
-
-    let api = api
-        .layer(from_fn_with_state(ctx.clone(), distributed_tabby_layer))
-        .route(
-            "/graphql",
-            routing::post(graphql::<Arc<Schema>, Arc<dyn ServiceLocator>>).with_state(ctx.clone()),
-        )
-        .route("/graphql", routing::get(playground("/graphql", None)))
-        .layer(Extension(schema))
-        .route(
-            "/hub",
-            routing::get(hub::ws_handler).with_state(ctx.clone()),
-        )
-        .nest("/repositories", repositories::routes(ctx.clone()));
-
-    let ui = ui
-        .route("/graphiql", routing::get(graphiql("/graphql", None)))
-        .fallback(ui::handler);
-
-    (api, ui)
-}
-
-async fn distributed_tabby_layer(
-    State(ws): State<Arc<dyn ServiceLocator>>,
-    request: Request<Body>,
-    next: Next<Body>,
-) -> axum::response::Response {
-    ws.worker().dispatch_request(request, next).await
-}
diff --git a/ee/tabby-webserver/src/repositories/mod.rs b/ee/tabby-webserver/src/repositories/mod.rs
@@ -12,23 +12,41 @@ use axum::{
 };
 use hyper::Body;
 use juniper_axum::extract::AuthBearer;
-use tabby_common::path::repositories_dir;
+use tabby_common::config::RepositoryConfig;
 use tracing::{instrument, warn};
 
 use crate::{
-    repositories,
-    repositories::resolve::{resolve_dir, resolve_file, resolve_meta, Meta, ResolveParams},
+    repositories::resolve::{
+        resolve_all, resolve_dir, resolve_file, resolve_meta, Meta, ResolveParams,
+    },
     schema::ServiceLocator,
 };
 
-pub fn routes(locator: Arc<dyn ServiceLocator>) -> Router {
+#[derive(Debug)]
+pub struct ResolveState {
+    pub repositories: Vec<RepositoryConfig>,
+}
+
+impl ResolveState {
+    pub fn find_repository(&self, name: &str) -> Option<&RepositoryConfig> {
+        self.repositories.iter().find(|repo| repo.name() == name)
+    }
+}
+
+pub fn routes(rs: Arc<ResolveState>, locator: Arc<dyn ServiceLocator>) -> Router {
     Router::new()
+        .route("/resolve", routing::get(resolve))
+        .with_state(rs.clone())
+        .route("/resolve/", routing::get(resolve))
+        .with_state(rs.clone())
         .route("/:name/resolve/.git/", routing::get(not_found))
         .route("/:name/resolve/.git/*path", routing::get(not_found))
-        .route("/:name/resolve/", routing::get(repositories::resolve))
-        .route("/:name/resolve/*path", routing::get(repositories::resolve))
-        .route("/:name/meta/", routing::get(repositories::meta))
-        .route("/:name/meta/*path", routing::get(repositories::meta))
+        .route("/:name/resolve/", routing::get(resolve_path))
+        .with_state(rs.clone())
+        .route("/:name/resolve/*path", routing::get(resolve_path))
+        .with_state(rs.clone())
+        .route("/:name/meta/", routing::get(meta))
+        .route("/:name/meta/*path", routing::get(meta))
         .fallback(not_found)
         .layer(from_fn_with_state(locator, require_login_middleware))
 }
@@ -61,9 +79,15 @@ async fn not_found() -> StatusCode {
 }
 
 #[instrument(skip(repo))]
-async fn resolve(Path(repo): Path<ResolveParams>) -> Result<Response, StatusCode> {
-    let root = repositories_dir().join(repo.name_str());
-    let full_path = root.join(repo.path_str());
+async fn resolve_path(
+    State(rs): State<Arc<ResolveState>>,
+    Path(repo): Path<ResolveParams>,
+) -> Result<Response, StatusCode> {
+    let Some(conf) = rs.find_repository(repo.name_str()) else {
+        return Err(StatusCode::NOT_FOUND);
+    };
+    let root = conf.dir();
+    let full_path = root.join(repo.os_path());
     let is_dir = tokio::fs::metadata(full_path.clone())
         .await
         .map(|m| m.is_dir())
@@ -96,3 +120,7 @@ async fn meta(Path(repo): Path<ResolveParams>) -> Result<Json<Meta>, StatusCode>
     }
     Err(StatusCode::NOT_FOUND)
 }
+
+async fn resolve(State(rs): State<Arc<ResolveState>>) -> Result<Response, StatusCode> {
+    resolve_all(rs).map_err(|_| StatusCode::NOT_FOUND)
+}
diff --git a/ee/tabby-webserver/src/repositories/resolve.rs b/ee/tabby-webserver/src/repositories/resolve.rs
@@ -1,4 +1,4 @@
-use std::{collections::HashMap, path::PathBuf, str::FromStr};
+use std::{collections::HashMap, path::PathBuf, str::FromStr, sync::Arc};
 
 use anyhow::Result;
 use axum::{
@@ -14,6 +14,8 @@ use tabby_common::{config::Config, SourceFile, Tag};
 use tower::ServiceExt;
 use tower_http::services::ServeDir;
 
+use crate::repositories::ResolveState;
+
 lazy_static! {
     static ref META: HashMap<DatasetKey, Meta> = load_meta();
 }
@@ -22,7 +24,7 @@ const DIRECTORY_MIME_TYPE: &str = "application/vnd.directory+json";
 
 #[derive(Hash, PartialEq, Eq, Debug)]
 pub struct DatasetKey {
-    local_name: String,
+    repo_name: String,
     rel_path: String,
 }
 
@@ -35,8 +37,8 @@ pub struct ResolveParams {
 impl ResolveParams {
     pub fn dataset_key(&self) -> DatasetKey {
         DatasetKey {
-            local_name: self.name.clone(),
-            rel_path: self.path_str().to_string(),
+            repo_name: self.name.clone(),
+            rel_path: self.os_path(),
         }
     }
 
@@ -47,6 +49,14 @@ impl ResolveParams {
     pub fn path_str(&self) -> &str {
         self.path.as_deref().unwrap_or("")
     }
+
+    pub fn os_path(&self) -> String {
+        if cfg!(target_os = "windows") {
+            self.path.clone().unwrap_or_default().replace('/', r"\")
+        } else {
+            self.path.clone().unwrap_or_default()
+        }
+    }
 }
 
 #[derive(Serialize)]
@@ -105,16 +115,13 @@ fn load_meta() -> HashMap<DatasetKey, Meta> {
             return dataset;
         }
     };
-    let iter = match SourceFile::all() {
-        Ok(all) => all,
-        Err(_) => {
-            return dataset;
-        }
+    let Ok(iter) = SourceFile::all() else {
+        return dataset;
     };
     for file in iter {
-        if let Some(name) = repo_conf.get(&file.git_url).map(|repo| repo.name()) {
+        if let Some(repo_name) = repo_conf.get(&file.git_url).map(|repo| repo.name()) {
             let key = DatasetKey {
-                local_name: name,
+                repo_name,
                 rel_path: file.filepath.clone(),
             };
             dataset.insert(key, file.into());
@@ -179,3 +186,21 @@ pub fn resolve_meta(key: &DatasetKey) -> Option<Meta> {
     }
     None
 }
+
+pub fn resolve_all(rs: Arc<ResolveState>) -> Result<Response> {
+    let entries: Vec<_> = rs
+        .repositories
+        .iter()
+        .map(|repo| DirEntry {
+            kind: DirEntryKind::Dir,
+            basename: repo.name(),
+        })
+        .collect();
+
+    let body = Json(ListDir { entries }).into_response();
+    let resp = Response::builder()
+        .header(header::CONTENT_TYPE, DIRECTORY_MIME_TYPE)
+        .body(body.into_body())?;
+
+    Ok(resp)
+}