Initial swagger schemas

alisinabh · Apr 30, 2024 · e9dfebf · e9dfebf
1 parent 30234dc
commit e9dfebf
Show file tree

Hide file tree

Showing 8 changed files with 567 additions and 6 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -20,3 +20,5 @@ reqwest = { version = "0.11", features = ["stream"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 tokio = { version = "1", features = ["full"] }
+utoipa = "4"
+utoipa-swagger-ui = { version = "4", features = ["actix-web"] }
diff --git a/src/lib.rs b/src/lib.rs
@@ -8,6 +8,10 @@ pub mod services;
 use actix_web::{web, App, HttpServer};
 use maxmind_db::MaxmindDB;
 use std::io::Result;
+use utoipa::OpenApi;
+use utoipa_swagger_ui::SwaggerUi;
+
+use crate::models::{HealthCheckModel, LookupResponseModel, LookupResult};
 
 pub async fn init_db(db_path: &str, db_variant: &str) -> web::Data<MaxmindDB> {
     let maxmind_db = MaxmindDB::init(db_variant, db_path)
@@ -25,16 +29,38 @@ pub async fn start_server(
     maxmind_db_arc: web::Data<MaxmindDB>,
     host: &str,
     port: u16,
+    swagger_ui_enabled: bool,
 ) -> Result<()> {
     // Start HTTP Server
     HttpServer::new(move || {
         let reader_data = maxmind_db_arc.clone();
-        App::new()
+        let mut app = App::new()
             .app_data(reader_data)
             .service(services::lookup::handle)
-            .service(services::healthcheck::handle)
+            .service(services::healthcheck::handle);
+
+        if swagger_ui_enabled {
+            app = app.service(swagger_ui_service())
+        }
+
+        app
     })
     .bind((host, port))?
     .run()
     .await
 }
+
+fn swagger_ui_service() -> SwaggerUi {
+    SwaggerUi::new("swagger-ui/{_:.*}").url("/api-docs/openapi.json", api_doc())
+}
+
+fn api_doc() -> utoipa::openapi::OpenApi {
+    #[derive(OpenApi)]
+    #[openapi(
+        paths(services::healthcheck::handle, services::lookup::handle),
+        components(schemas(LookupResponseModel, LookupResult, HealthCheckModel))
+    )]
+    struct ApiDoc;
+
+    ApiDoc::openapi()
+}
diff --git a/src/main.rs b/src/main.rs
@@ -17,6 +17,11 @@ async fn main() -> Result<()> {
         .parse()
         .expect("Invalid PORT value");
 
+    let swagger_ui_enabled: bool = env::var("SWAGGER_UI_ENABLED")
+        .unwrap_or("true".to_string())
+        .parse()
+        .expect("Invalid SWAGGER_UI_ENABLED value. Expected `false` or `true`");
+
     let maxmind_db_arc = atlas_rs::init_db(&db_path, &db_variant).await;
 
     let subcommand = args().skip(1).next();
@@ -26,7 +31,7 @@ async fn main() -> Result<()> {
             // Start Database Updater Daemon
             atlas_rs::start_db_refresher(maxmind_db_arc.clone(), update_interval);
             // Start Server
-            atlas_rs::start_server(maxmind_db_arc, &host, port).await
+            atlas_rs::start_server(maxmind_db_arc, &host, port, swagger_ui_enabled).await
         }
         Some("init") => Ok(()),
         Some(command) => Err(Error::new(

diff --git a/src/models.rs b/src/models.rs
@@ -1,6 +1,10 @@
 use serde::Serialize;
 use std::collections::HashMap;
 use std::net::IpAddr;
+use utoipa::{
+    openapi::{schema::AdditionalProperties, ObjectBuilder},
+    ToSchema,
+};
 
 use maxminddb::geoip2::{
     AnonymousIp, Asn, City, ConnectionType, Country, DensityIncome, Enterprise, Isp,
@@ -37,8 +41,236 @@ impl<'a> Serialize for LookupResult<'a> {
     }
 }
 
-#[derive(Serialize)]
+#[derive(Serialize, ToSchema)]
 pub struct LookupResponseModel<'a> {
     pub results: LookupResult<'a>,
     pub database_build_epoch: u64,
 }
+
+pub struct HealthCheckModel;
+
+impl<'__s> utoipa::ToSchema<'__s> for HealthCheckModel {
+    fn schema() -> (
+        &'__s str,
+        utoipa::openapi::RefOr<utoipa::openapi::schema::Schema>,
+    ) {
+        (
+            "HealthCheckModel",
+            ObjectBuilder::new()
+                .schema_type(utoipa::openapi::SchemaType::String)
+                .example(Some("Ok".into()))
+                .into(),
+        )
+    }
+}
+
+impl<'__s> utoipa::ToSchema<'__s> for LookupResult<'__s> {
+    fn schema() -> (
+        &'__s str,
+        utoipa::openapi::RefOr<utoipa::openapi::schema::Schema>,
+    ) {
+        (
+            "LookupResult",
+            utoipa::openapi::ObjectBuilder::new()
+                .property(
+                    "{ip_address}",
+                    utoipa::openapi::Schema::OneOf(
+                        utoipa::openapi::schema::OneOfBuilder::new()
+                            .item(LookupResult::city_schema())
+                            .into(),
+                    ),
+                )
+                .into(),
+        )
+    }
+}
+
+impl<'a> LookupResult<'a> {
+    fn localized_string() -> ObjectBuilder {
+        ObjectBuilder::new()
+            .schema_type(utoipa::openapi::SchemaType::Object)
+            .title(Some("LocalizedString"))
+            .property(
+                "{language_code}",
+                ObjectBuilder::new().schema_type(utoipa::openapi::SchemaType::String),
+            )
+            .additional_properties(Some(AdditionalProperties::FreeForm(true)))
+    }
+
+    fn geoname_id() -> ObjectBuilder {
+        ObjectBuilder::new()
+            .schema_type(utoipa::openapi::SchemaType::Integer)
+            .example(Some(serde_json::Value::Number(1234.into())))
+    }
+
+    fn city_schema() -> ObjectBuilder {
+        ObjectBuilder::new()
+            .title(Some("CityLookupResult"))
+            .schema_type(utoipa::openapi::SchemaType::Object)
+            .property(
+                "city",
+                ObjectBuilder::new()
+                    .property("geoname_id", Self::geoname_id())
+                    .property(
+                        "names",
+                        Self::localized_string()
+                            .example(serde_json::json!({"en": "San Diego"}).into()),
+                    ),
+            )
+            .property(
+                "continent",
+                ObjectBuilder::new()
+                    .property(
+                        "code",
+                        ObjectBuilder::new()
+                            .schema_type(utoipa::openapi::SchemaType::String)
+                            .example(Some("NA".into())),
+                    )
+                    .property("geoname_id", Self::geoname_id())
+                    .property(
+                        "names",
+                        Self::localized_string()
+                            .example(serde_json::json!({"en": "North America"}).into()),
+                    ),
+            )
+            .property(
+                "country",
+                ObjectBuilder::new()
+                    .property("geoname_id", Self::geoname_id())
+                    .property(
+                        "is_in_european_union",
+                        ObjectBuilder::new().schema_type(utoipa::openapi::SchemaType::Boolean),
+                    )
+                    .property(
+                        "iso_code",
+                        ObjectBuilder::new()
+                            .schema_type(utoipa::openapi::SchemaType::String)
+                            .example(Some("US".into())),
+                    )
+                    .property(
+                        "names",
+                        Self::localized_string()
+                            .example(serde_json::json!({"en": "United States"}).into()),
+                    ),
+            )
+            .property(
+                "location",
+                ObjectBuilder::new()
+                    .property(
+                        "accuracy_radius",
+                        ObjectBuilder::new()
+                            .schema_type(utoipa::openapi::SchemaType::Number)
+                            .example(Some(200.into())),
+                    )
+                    .property(
+                        "latitude",
+                        ObjectBuilder::new()
+                            .schema_type(utoipa::openapi::SchemaType::Number)
+                            .example(Some((30.0406).into())),
+                    )
+                    .property(
+                        "longitude",
+                        ObjectBuilder::new()
+                            .schema_type(utoipa::openapi::SchemaType::Number)
+                            .example(Some((-80.26).into())),
+                    )
+                    .property(
+                        "metro_code",
+                        ObjectBuilder::new()
+                            .schema_type(utoipa::openapi::SchemaType::Number)
+                            .example(Some(518.into())),
+                    )
+                    .property(
+                        "time_zone",
+                        ObjectBuilder::new()
+                            .schema_type(utoipa::openapi::SchemaType::String)
+                            .example(Some("America/New_York".into())),
+                    ),
+            )
+            .property(
+                "postal",
+                ObjectBuilder::new().property(
+                    "code",
+                    ObjectBuilder::new()
+                        .schema_type(utoipa::openapi::SchemaType::String)
+                        .example(Some("27127".into())),
+                ),
+            )
+            .property(
+                "registered_country",
+                ObjectBuilder::new()
+                    .property("geoname_id", Self::geoname_id())
+                    .property(
+                        "is_in_european_union",
+                        ObjectBuilder::new().schema_type(utoipa::openapi::SchemaType::Boolean),
+                    )
+                    .property(
+                        "iso_code",
+                        ObjectBuilder::new()
+                            .schema_type(utoipa::openapi::SchemaType::String)
+                            .example(Some("US".into())),
+                    )
+                    .property(
+                        "names",
+                        Self::localized_string()
+                            .example(serde_json::json!({"en": "United States"}).into()),
+                    ),
+            )
+            .property(
+                "represented_country",
+                ObjectBuilder::new()
+                    .property("geoname_id", Self::geoname_id())
+                    .property(
+                        "is_in_european_union",
+                        ObjectBuilder::new().schema_type(utoipa::openapi::SchemaType::Boolean),
+                    )
+                    .property(
+                        "iso_code",
+                        ObjectBuilder::new()
+                            .schema_type(utoipa::openapi::SchemaType::String)
+                            .example(Some("US".into())),
+                    )
+                    .property(
+                        "names",
+                        Self::localized_string()
+                            .example(serde_json::json!({"en": "United States"}).into()),
+                    )
+                    .property(
+                        "type",
+                        ObjectBuilder::new().schema_type(utoipa::openapi::SchemaType::String),
+                    ),
+            )
+            .property(
+                "subdivisions",
+                ObjectBuilder::new()
+                    .property("geoname_id", Self::geoname_id())
+                    .property(
+                        "iso_code",
+                        ObjectBuilder::new()
+                            .schema_type(utoipa::openapi::SchemaType::String)
+                            .example(Some("NC".into())),
+                    )
+                    .property(
+                        "names",
+                        Self::localized_string()
+                            .example(serde_json::json!({"en": "North Carolina"}).into()),
+                    ),
+            )
+            .property(
+                "traits",
+                ObjectBuilder::new()
+                    .property(
+                        "is_anonymous_proxy",
+                        ObjectBuilder::new().schema_type(utoipa::openapi::SchemaType::Boolean),
+                    )
+                    .property(
+                        "is_anycast",
+                        ObjectBuilder::new().schema_type(utoipa::openapi::SchemaType::Boolean),
+                    )
+                    .property(
+                        "is_satellite_provider",
+                        ObjectBuilder::new().schema_type(utoipa::openapi::SchemaType::Boolean),
+                    ),
+            )
+    }
+}
diff --git a/src/services/healthcheck.rs b/src/services/healthcheck.rs
@@ -1,6 +1,13 @@
 use actix_web::{get, HttpResponse, Responder};
 
+#[utoipa::path(
+    get,
+    path = "/health",
+    responses(
+        (status = 200, description = "Ok", body = HealthCheckModel, content_type = "text/plain")
+    ),
+)]
 #[get("/health")]
 pub async fn handle() -> impl Responder {
-    HttpResponse::Ok().body("Ok")
+    HttpResponse::Ok().content_type("text/plain").body("Ok")
 }
diff --git a/src/services/lookup.rs b/src/services/lookup.rs
@@ -6,6 +6,17 @@ use crate::network_utils::SpecialIPCheck;
 use actix_web::{get, web, HttpResponse, Responder};
 use std::net::IpAddr;
 
+#[utoipa::path(
+    get,
+    path = "/geoip/lookup/{lookup_type}/{ip_addresses}",
+    responses(
+        (status = 200, description = "Ok", body = LookupResponseModel)
+    ),
+    params(
+        ("lookup_type" = String, Path, description = "Type of the lookup", example = "city"),
+        ("ip_addresses" = String, Path, description = "List of ip addresses separated by comma", example = "4.2.2.4")
+    )
+)]
 #[get("/geoip/lookup/{lookup_type}/{ip_addresses}")]
 async fn handle(data: web::Data<MaxmindDB>, path: web::Path<(String, String)>) -> impl Responder {
     let (lookup_type, ip_addresses) = path.into_inner();

diff --git a/tests/healthcheck_service.rs b/tests/healthcheck_service.rs
@@ -1,9 +1,11 @@
-use actix_web::{test, App};
+use actix_web::{body::MessageBody, test, web, App};
 
 #[actix_web::test]
 async fn test_healthcheck_endpoint() {
     let app = test::init_service(App::new().service(atlas_rs::services::healthcheck::handle)).await;
     let req = test::TestRequest::get().uri("/health").to_request();
     let resp = test::call_service(&app, req).await;
     assert!(resp.status().is_success());
+    let body = resp.into_body().try_into_bytes().unwrap();
+    assert_eq!(body, web::Bytes::from_static(b"Ok"));
 }