Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor(torii-server): server proxy handlers #2708

Merged
merged 3 commits into from
Nov 22, 2024
Merged

refactor(torii-server): server proxy handlers #2708

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
aa3c5ac
refactor(torii-server): cleanup & handlers
Larkooo Nov 21, 2024
ff49502
refactor: handlers
Larkooo Nov 21, 2024
7d12dba
better sql error handling
Larkooo Nov 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/torii/server/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,3 +30,4 @@ tower.workspace = true
tracing.workspace = true
warp.workspace = true
form_urlencoded = "1.2.1"
async-trait = "0.1.83"
46 changes: 46 additions & 0 deletions crates/torii/server/src/handlers/graphql.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
use std::net::{IpAddr, SocketAddr};

use http::StatusCode;
use hyper::{Body, Request, Response};
use tracing::error;

use super::Handler;

pub(crate) const LOG_TARGET: &str = "torii::server::handlers::graphql";

pub struct GraphQLHandler {
client_ip: IpAddr,
graphql_addr: Option<SocketAddr>,
}

impl GraphQLHandler {
pub fn new(client_ip: IpAddr, graphql_addr: Option<SocketAddr>) -> Self {
Self { client_ip, graphql_addr }
}

Check warning on line 19 in crates/torii/server/src/handlers/graphql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/graphql.rs#L17-L19 

Added lines #L17 - L19 were not covered by tests
}

#[async_trait::async_trait]
impl Handler for GraphQLHandler {
fn should_handle(&self, req: &Request<Body>) -> bool {
req.uri().path().starts_with("/graphql")
}

Check warning on line 26 in crates/torii/server/src/handlers/graphql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/graphql.rs#L24-L26 

Added lines #L24 - L26 were not covered by tests

async fn handle(&self, req: Request<Body>) -> Response<Body> {
if let Some(addr) = self.graphql_addr {
let graphql_addr = format!("http://{}", addr);
match crate::proxy::GRAPHQL_PROXY_CLIENT.call(self.client_ip, &graphql_addr, req).await
{
Ok(response) => response,
Err(_error) => {
error!(target: LOG_TARGET, "GraphQL proxy error: {:?}", _error);
Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.body(Body::empty())
.unwrap()
}
}
Larkooo marked this conversation as resolved.
Show resolved Hide resolved
} else {
Response::builder().status(StatusCode::NOT_FOUND).body(Body::empty()).unwrap()
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Handle Response builder error case, sensei!

The unwrap() call on the Response builder could panic. Consider using unwrap_or_else for safer error handling.

-            Response::builder().status(StatusCode::NOT_FOUND).body(Body::empty()).unwrap()
+            Response::builder()
+                .status(StatusCode::NOT_FOUND)
+                .body(Body::empty())
+                .unwrap_or_else(|_| Response::new(Body::empty()))
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
Response::builder().status(StatusCode::NOT_FOUND).body(Body::empty()).unwrap()
Response::builder()
.status(StatusCode::NOT_FOUND)
.body(Body::empty())
.unwrap_or_else(|_| Response::new(Body::empty()))

}
}

Check warning on line 45 in crates/torii/server/src/handlers/graphql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/graphql.rs#L28-L45 

Added lines #L28 - L45 were not covered by tests
}
49 changes: 49 additions & 0 deletions crates/torii/server/src/handlers/grpc.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
use std::net::{IpAddr, SocketAddr};

use http::header::CONTENT_TYPE;
use hyper::{Body, Request, Response, StatusCode};
use tracing::error;

use super::Handler;

pub(crate) const LOG_TARGET: &str = "torii::server::handlers::grpc";

pub struct GrpcHandler {
client_ip: IpAddr,
grpc_addr: Option<SocketAddr>,
}

impl GrpcHandler {
pub fn new(client_ip: IpAddr, grpc_addr: Option<SocketAddr>) -> Self {
Self { client_ip, grpc_addr }
}

Check warning on line 19 in crates/torii/server/src/handlers/grpc.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/grpc.rs#L17-L19 

Added lines #L17 - L19 were not covered by tests
}

#[async_trait::async_trait]
impl Handler for GrpcHandler {
fn should_handle(&self, req: &Request<Body>) -> bool {
req.headers()
.get(CONTENT_TYPE)
.and_then(|ct| ct.to_str().ok())
.map(|ct| ct.starts_with("application/grpc"))
.unwrap_or(false)
}

Check warning on line 30 in crates/torii/server/src/handlers/grpc.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/grpc.rs#L24-L30 

Added lines #L24 - L30 were not covered by tests

async fn handle(&self, req: Request<Body>) -> Response<Body> {
if let Some(grpc_addr) = self.grpc_addr {
let grpc_addr = format!("http://{}", grpc_addr);
match crate::proxy::GRPC_PROXY_CLIENT.call(self.client_ip, &grpc_addr, req).await {
Ok(response) => response,
Err(_error) => {
error!(target: LOG_TARGET, "{:?}", _error);
Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.body(Body::empty())
.unwrap()
}
}
} else {
Response::builder().status(StatusCode::NOT_FOUND).body(Body::empty()).unwrap()
}
}

Check warning on line 48 in crates/torii/server/src/handlers/grpc.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/grpc.rs#L32-L48 

Added lines #L32 - L48 were not covered by tests
Larkooo marked this conversation as resolved.
Show resolved Hide resolved
}
15 changes: 15 additions & 0 deletions crates/torii/server/src/handlers/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
pub mod graphql;
pub mod grpc;
pub mod sql;
pub mod static_files;

use hyper::{Body, Request, Response};

#[async_trait::async_trait]
pub trait Handler: Send + Sync {
// Check if this handler should handle the given request
fn should_handle(&self, req: &Request<Body>) -> bool;

// Handle the request
async fn handle(&self, req: Request<Body>) -> Response<Body>;
}
Larkooo marked this conversation as resolved.
Show resolved Hide resolved
134 changes: 134 additions & 0 deletions crates/torii/server/src/handlers/sql.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
use std::sync::Arc;

use base64::engine::general_purpose::STANDARD;
use base64::Engine;
use http::header::CONTENT_TYPE;
use hyper::{Body, Method, Request, Response, StatusCode};
use sqlx::{Column, Row, SqlitePool, TypeInfo};

use super::Handler;

pub struct SqlHandler {
pool: Arc<SqlitePool>,
}

impl SqlHandler {
pub fn new(pool: Arc<SqlitePool>) -> Self {
Self { pool }
}

Check warning on line 18 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L16-L18 

Added lines #L16 - L18 were not covered by tests

pub async fn execute_query(&self, query: String) -> Response<Body> {
match sqlx::query(&query).fetch_all(&*self.pool).await {
Ok(rows) => {
let result: Vec<_> = rows
.iter()
.map(|row| {
let mut obj = serde_json::Map::new();
for (i, column) in row.columns().iter().enumerate() {
let value: serde_json::Value = match column.type_info().name() {
"TEXT" => row
.get::<Option<String>, _>(i)
.map_or(serde_json::Value::Null, serde_json::Value::String),
"INTEGER" | "NULL" => row
.get::<Option<i64>, _>(i)
.map_or(serde_json::Value::Null, |n| {
serde_json::Value::Number(n.into())
}),
"REAL" => row.get::<Option<f64>, _>(i).map_or(
serde_json::Value::Null,
|f| {
serde_json::Number::from_f64(f).map_or(
serde_json::Value::Null,
serde_json::Value::Number,
)
},
),
"BLOB" => row
.get::<Option<Vec<u8>>, _>(i)
.map_or(serde_json::Value::Null, |bytes| {
serde_json::Value::String(STANDARD.encode(bytes))
}),
_ => row
.get::<Option<String>, _>(i)
.map_or(serde_json::Value::Null, serde_json::Value::String),

Check warning on line 53 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L20-L53 

Added lines #L20 - L53 were not covered by tests
};
obj.insert(column.name().to_string(), value);

Check warning on line 55 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L55 

Added line #L55 was not covered by tests
}
serde_json::Value::Object(obj)
})
.collect();

Check warning on line 59 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L57-L59 

Added lines #L57 - L59 were not covered by tests

let json = match serde_json::to_string(&result) {
Ok(json) => json,
Err(e) => {
return Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.body(Body::from(format!("Failed to serialize result: {:?}", e)))
.unwrap();

Check warning on line 67 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L61-L67 

Added lines #L61 - L67 were not covered by tests
}
};
Comment on lines +61 to +69
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Ohayo sensei! Handle serialization errors more gracefully

The current error handling exposes internal error details and uses unwrap(). This could leak sensitive information and potentially crash the server.

Consider this safer approach:

                 let json = match serde_json::to_string(&result) {
                     Ok(json) => json,
                     Err(e) => {
+                        // Log the error internally
+                        tracing::error!("Serialization error: {:?}", e);
                         return Response::builder()
                             .status(StatusCode::INTERNAL_SERVER_ERROR)
-                            .body(Body::from(format!("Failed to serialize result: {:?}", e)))
+                            .body(Body::from("Internal server error"))
                             .unwrap_or_else(|_| Response::new(Body::from("Internal server error")));
                     }
                 };
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
let json = match serde_json::to_string(&result) {
Ok(json) => json,
Err(e) => {
return Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.body(Body::from(format!("Failed to serialize result: {:?}", e)))
.unwrap();
}
};
let json = match serde_json::to_string(&result) {
Ok(json) => json,
Err(e) => {
// Log the error internally
tracing::error!("Serialization error: {:?}", e);
return Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.body(Body::from("Internal server error"))
.unwrap_or_else(|_| Response::new(Body::from("Internal server error")));
}
};


Response::builder()
.status(StatusCode::OK)
.header(CONTENT_TYPE, "application/json")
.body(Body::from(json))
.unwrap()

Check warning on line 75 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L71-L75 

Added lines #L71 - L75 were not covered by tests
}
Err(e) => Response::builder()
.status(StatusCode::BAD_REQUEST)
.body(Body::from(format!("Query error: {:?}", e)))
.unwrap(),

Check warning on line 80 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L77-L80 

Added lines #L77 - L80 were not covered by tests
Larkooo marked this conversation as resolved.
Show resolved Hide resolved
}
}

Check warning on line 82 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L82 

Added line #L82 was not covered by tests
Comment on lines +20 to +82
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Ohayo sensei! Add query validation and sanitization

The execute_query method executes queries without validation or sanitization, which could lead to SQL injection or harmful operations.

Consider:

  1. Implementing a query allowlist
  2. Using parameterized queries
  3. Adding query complexity limits
  4. Restricting operations (e.g., no DROP, DELETE, etc.)

Would you like assistance implementing these security measures?


async fn extract_query(&self, req: Request<Body>) -> Result<String, Response<Body>> {
match *req.method() {

Check warning on line 85 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L84-L85 

Added lines #L84 - L85 were not covered by tests
Method::GET => {
// Get the query from the query params
let params = req.uri().query().unwrap_or_default();
form_urlencoded::parse(params.as_bytes())
.find(|(key, _)| key == "q" || key == "query")
.map(|(_, value)| value.to_string())
.ok_or(
Response::builder()
.status(StatusCode::BAD_REQUEST)
.body(Body::from("Missing 'q' or 'query' parameter."))
.unwrap(),
)

Check warning on line 97 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L88-L97 

Added lines #L88 - L97 were not covered by tests
}
Method::POST => {
// Get the query from request body
let body_bytes = hyper::body::to_bytes(req.into_body()).await.map_err(|_| {
Response::builder()
.status(StatusCode::BAD_REQUEST)
.body(Body::from("Failed to read query from request body"))
.unwrap()
})?;
String::from_utf8(body_bytes.to_vec()).map_err(|_| {
Response::builder()
.status(StatusCode::BAD_REQUEST)
.body(Body::from("Invalid query"))
.unwrap()
})

Check warning on line 112 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L101-L112 

Added lines #L101 - L112 were not covered by tests
}
_ => Err(Response::builder()
.status(StatusCode::METHOD_NOT_ALLOWED)
.body(Body::from("Only GET and POST methods are allowed"))
.unwrap()),

Check warning on line 117 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L114-L117 

Added lines #L114 - L117 were not covered by tests
}
}

Check warning on line 119 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L119 

Added line #L119 was not covered by tests
Comment on lines +84 to +119
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Ohayo sensei! Add request size limits and improve error handling

The query extraction lacks request size limits and uses unwrap() in error paths.

Consider these improvements:

  1. Add request size limits:
const MAX_BODY_SIZE: usize = 1024 * 1024; // 1MB limit

// In POST handler:
if body_bytes.len() > MAX_BODY_SIZE {
    return Err(Response::builder()
        .status(StatusCode::PAYLOAD_TOO_LARGE)
        .body(Body::from("Request body too large"))
        .unwrap_or_else(|_| Response::new(Body::from("Request body too large"))));
}
  1. Replace unwrap() with unwrap_or_else:
-                            .unwrap(),
+                            .unwrap_or_else(|_| Response::new(Body::from("Bad request"))),

}

#[async_trait::async_trait]
impl Handler for SqlHandler {
fn should_handle(&self, req: &Request<Body>) -> bool {
req.uri().path().starts_with("/sql")
}

Check warning on line 126 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L124-L126 

Added lines #L124 - L126 were not covered by tests
Comment on lines +124 to +126
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Ohayo sensei! Enhance request validation

The should_handle method only checks the path prefix. Consider adding more validation:

  1. Content-Type validation for POST requests
  2. Query parameter validation for GET requests
     fn should_handle(&self, req: &Request<Body>) -> bool {
-        req.uri().path().starts_with("/sql")
+        if !req.uri().path().starts_with("/sql") {
+            return false;
+        }
+        match *req.method() {
+            Method::POST => req
+                .headers()
+                .get(CONTENT_TYPE)
+                .map_or(false, |ct| ct == "application/json"),
+            Method::GET => true,
+            _ => false,
+        }
     }

Committable suggestion skipped: line range outside the PR's diff.


async fn handle(&self, req: Request<Body>) -> Response<Body> {
match self.extract_query(req).await {
Ok(query) => self.execute_query(query).await,
Err(response) => response,
}
}

Check warning on line 133 in crates/torii/server/src/handlers/sql.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/sql.rs#L128-L133 

Added lines #L128 - L133 were not covered by tests
}
47 changes: 47 additions & 0 deletions crates/torii/server/src/handlers/static_files.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
use std::net::{IpAddr, SocketAddr};

use hyper::{Body, Request, Response, StatusCode};
use tracing::error;

use super::Handler;

pub(crate) const LOG_TARGET: &str = "torii::server::handlers::static";

pub struct StaticHandler {
client_ip: IpAddr,
artifacts_addr: Option<SocketAddr>,
}

impl StaticHandler {
pub fn new(client_ip: IpAddr, artifacts_addr: Option<SocketAddr>) -> Self {
Self { client_ip, artifacts_addr }
}

Check warning on line 18 in crates/torii/server/src/handlers/static_files.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/static_files.rs#L16-L18 

Added lines #L16 - L18 were not covered by tests
}

#[async_trait::async_trait]
impl Handler for StaticHandler {
fn should_handle(&self, req: &Request<Body>) -> bool {
req.uri().path().starts_with("/static")
}

Check warning on line 25 in crates/torii/server/src/handlers/static_files.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/static_files.rs#L23-L25 

Added lines #L23 - L25 were not covered by tests

async fn handle(&self, req: Request<Body>) -> Response<Body> {
if let Some(artifacts_addr) = self.artifacts_addr {
let artifacts_addr = format!("http://{}", artifacts_addr);
match crate::proxy::GRAPHQL_PROXY_CLIENT
.call(self.client_ip, &artifacts_addr, req)
.await
{
Ok(response) => response,
Err(_error) => {
error!(target: LOG_TARGET, "{:?}", _error);
Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.body(Body::empty())
.unwrap()
}
}
} else {
Response::builder().status(StatusCode::NOT_FOUND).body(Body::empty()).unwrap()
}
}

Check warning on line 46 in crates/torii/server/src/handlers/static_files.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/server/src/handlers/static_files.rs#L27-L46 

Added lines #L27 - L46 were not covered by tests
}
1 change: 1 addition & 0 deletions crates/torii/server/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
pub mod artifacts;
pub(crate) mod handlers;
pub mod proxy;
Loading
Loading