-
Notifications
You must be signed in to change notification settings - Fork 188
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.
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
feat(torii): sql proxy endpoint for querying #2706
Changes from 3 commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -3,6 +3,8 @@ use std::net::{IpAddr, SocketAddr}; | |
use std::sync::Arc; | ||
use std::time::Duration; | ||
|
||
use base64::engine::general_purpose::STANDARD; | ||
use base64::Engine; | ||
use http::header::CONTENT_TYPE; | ||
use http::{HeaderName, Method}; | ||
use hyper::client::connect::dns::GaiResolver; | ||
|
@@ -12,6 +14,7 @@ use hyper::service::make_service_fn; | |
use hyper::{Body, Client, Request, Response, Server, StatusCode}; | ||
use hyper_reverse_proxy::ReverseProxy; | ||
use serde_json::json; | ||
use sqlx::{Column, Row, SqlitePool, TypeInfo}; | ||
use tokio::sync::RwLock; | ||
use tower::ServiceBuilder; | ||
use tower_http::cors::{AllowOrigin, CorsLayer}; | ||
|
@@ -60,6 +63,7 @@ pub struct Proxy { | |
grpc_addr: Option<SocketAddr>, | ||
artifacts_addr: Option<SocketAddr>, | ||
graphql_addr: Arc<RwLock<Option<SocketAddr>>>, | ||
pool: Arc<SqlitePool>, | ||
} | ||
|
||
impl Proxy { | ||
|
@@ -69,13 +73,15 @@ impl Proxy { | |
grpc_addr: Option<SocketAddr>, | ||
graphql_addr: Option<SocketAddr>, | ||
artifacts_addr: Option<SocketAddr>, | ||
pool: Arc<SqlitePool>, | ||
) -> Self { | ||
Self { | ||
addr, | ||
allowed_origins, | ||
grpc_addr, | ||
graphql_addr: Arc::new(RwLock::new(graphql_addr)), | ||
artifacts_addr, | ||
pool, | ||
} | ||
} | ||
|
||
|
@@ -93,6 +99,7 @@ impl Proxy { | |
let grpc_addr = self.grpc_addr; | ||
let graphql_addr = self.graphql_addr.clone(); | ||
let artifacts_addr = self.artifacts_addr; | ||
let pool = self.pool.clone(); | ||
|
||
let make_svc = make_service_fn(move |conn: &AddrStream| { | ||
let remote_addr = conn.remote_addr().ip(); | ||
|
@@ -129,12 +136,14 @@ impl Proxy { | |
), | ||
}); | ||
|
||
let pool_clone = pool.clone(); | ||
let graphql_addr_clone = graphql_addr.clone(); | ||
let service = ServiceBuilder::new().option_layer(cors).service_fn(move |req| { | ||
let pool = pool_clone.clone(); | ||
let graphql_addr = graphql_addr_clone.clone(); | ||
async move { | ||
let graphql_addr = graphql_addr.read().await; | ||
handle(remote_addr, grpc_addr, artifacts_addr, *graphql_addr, req).await | ||
handle(remote_addr, grpc_addr, artifacts_addr, *graphql_addr, pool, req).await | ||
} | ||
}); | ||
|
||
|
@@ -156,6 +165,7 @@ async fn handle( | |
grpc_addr: Option<SocketAddr>, | ||
artifacts_addr: Option<SocketAddr>, | ||
graphql_addr: Option<SocketAddr>, | ||
pool: Arc<SqlitePool>, | ||
req: Request<Body>, | ||
) -> Result<Response<Body>, Infallible> { | ||
if req.uri().path().starts_with("/static") { | ||
|
@@ -224,6 +234,84 @@ async fn handle( | |
} | ||
} | ||
|
||
if req.uri().path().starts_with("/sql") { | ||
let query = if req.method() == Method::GET { | ||
// Get the query from URL parameters | ||
let params = req.uri().query().unwrap_or_default(); | ||
form_urlencoded::parse(params.as_bytes()) | ||
.find(|(key, _)| key == "q") | ||
.map(|(_, value)| value.to_string()) | ||
.unwrap_or_default() | ||
} else if req.method() == Method::POST { | ||
// Get the query from request body | ||
let body_bytes = hyper::body::to_bytes(req.into_body()).await.unwrap_or_default(); | ||
String::from_utf8(body_bytes.to_vec()).unwrap_or_default() | ||
} else { | ||
return Ok(Response::builder() | ||
.status(StatusCode::METHOD_NOT_ALLOWED) | ||
.body(Body::from("Only GET and POST methods are allowed")) | ||
.unwrap()); | ||
}; | ||
|
||
// Execute the query in a read-only transaction | ||
return match sqlx::query(&query).fetch_all(&*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" => 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), | ||
}; | ||
obj.insert(column.name().to_string(), value); | ||
} | ||
serde_json::Value::Object(obj) | ||
}) | ||
.collect(); | ||
|
||
let json = serde_json::to_string(&result).unwrap(); | ||
|
||
Ok(Response::builder() | ||
.status(StatusCode::OK) | ||
.header(CONTENT_TYPE, "application/json") | ||
.body(Body::from(json)) | ||
.unwrap()) | ||
} | ||
Err(e) => { | ||
error!("SQL query error: {:?}", e); | ||
Ok(Response::builder() | ||
.status(StatusCode::BAD_REQUEST) | ||
.body(Body::from("Query error")) | ||
.unwrap()) | ||
} | ||
}; | ||
} | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Critical Security Risk: Executing Arbitrary SQL Queries Ohayo, sensei! The I recommend removing this endpoint or implementing a safe abstraction that only allows predefined queries or parameterized inputs. Would you like assistance in refactoring this to enhance security? |
||
|
||
let json = json!({ | ||
"service": "torii", | ||
"success": true | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ Refactor suggestion
Ohayo, sensei! Avoid silently ignoring errors when reading request body
Using
unwrap_or_default()
when reading the request body may mask errors that should be handled appropriately. Consider properly handling errors when converting the request body to bytes and when converting bytes to a UTF-8 string.Apply this diff to handle errors explicitly: