heliosdb_proxy/
http_gateway.rs

1//! HTTP SQL gateway — `@neondatabase/serverless`-compatible.
2//!
3//! When `[http_gateway] enabled = true`, the proxy exposes a `POST /sql`
4//! endpoint that runs one SQL statement over the backend PG-wire client and
5//! returns a Neon-style JSON result (`{ command, rowCount, rows, fields }`).
6//! This lets edge/serverless runtimes (Cloudflare Workers, Vercel Edge) that
7//! cannot hold a TCP socket talk to vanilla PostgreSQL or HeliosDB-Nano over
8//! plain HTTP.
9//!
10//! Parameterised queries (`$1`,`$2`) are supported via a JSON `params` array.
11//! `Neon-Array-Mode: true` returns each row as an array instead of an object.
12//! A WebSocket session/transaction mode is the planned follow-on; this is the
13//! one-shot HTTP path the serverless driver uses for non-transactional queries.
14
15use std::sync::Arc;
16use std::time::Duration;
17
18use serde_json::{json, Value};
19use tokio::io::{AsyncWriteExt, BufReader};
20use tokio::net::TcpListener;
21
22use crate::backend::client::QueryResult;
23use crate::backend::types::TextValue;
24use crate::backend::{
25    tls::default_client_config, BackendClient, BackendConfig, ParamValue, TlsMode,
26};
27use crate::config::HttpGatewayConfig;
28use crate::{ProxyError, Result};
29
30pub struct HttpGateway {
31    config: HttpGatewayConfig,
32}
33
34impl HttpGateway {
35    pub fn new(config: HttpGatewayConfig) -> Self {
36        Self { config }
37    }
38
39    pub async fn run(self) -> Result<()> {
40        let listener = TcpListener::bind(&self.config.listen_address)
41            .await
42            .map_err(|e| {
43                ProxyError::Network(format!(
44                    "HTTP gateway bind {}: {}",
45                    self.config.listen_address, e
46                ))
47            })?;
48        tracing::info!(addr = %self.config.listen_address, "HTTP SQL gateway listening");
49        let cfg = Arc::new(self.config);
50        loop {
51            let (stream, peer) = match listener.accept().await {
52                Ok(x) => x,
53                Err(e) => {
54                    tracing::warn!("HTTP gateway accept error: {}", e);
55                    continue;
56                }
57            };
58            let cfg = cfg.clone();
59            tokio::spawn(async move {
60                if let Err(e) = Self::handle(stream, cfg).await {
61                    tracing::debug!(%peer, "HTTP gateway error: {}", e);
62                }
63            });
64        }
65    }
66
67    async fn handle(mut stream: tokio::net::TcpStream, cfg: Arc<HttpGatewayConfig>) -> Result<()> {
68        use crate::http_util;
69        let (reader, mut writer) = stream.split();
70        let mut reader = BufReader::new(reader);
71
72        // Bounded request read: overall deadline + header count/byte caps.
73        let deadline = tokio::time::Instant::now() + http_util::HTTP_READ_TIMEOUT;
74        let head = match http_util::read_head(&mut reader, deadline).await {
75            Ok(h) => h,
76            Err(_) => return Ok(()), // timeout / oversized headers / early close
77        };
78        let method = head.method.as_str();
79        let path = head.path.as_str();
80
81        // Liveness probe.
82        if method == "GET" && (path == "/health" || path == "/") {
83            return Self::respond(&mut writer, 200, &json!({"status":"ok"})).await;
84        }
85        // Constant-time Bearer check (no `==` short-circuit oracle; no per-request
86        // format! of the expected token).
87        let authorized = match cfg.auth_token.as_ref() {
88            None => true,
89            Some(tok) => head
90                .header("authorization")
91                .and_then(|v| v.strip_prefix("Bearer "))
92                .map(|got| http_util::constant_time_eq_str(got, tok))
93                .unwrap_or(false),
94        };
95        if !authorized {
96            return Self::respond(&mut writer, 401, &json!({"error":"unauthorized"})).await;
97        }
98        if method != "POST" {
99            return Self::respond(&mut writer, 405, &json!({"error":"use POST /sql"})).await;
100        }
101        // Reject an oversized declared body BEFORE allocating for it.
102        if head.content_length > http_util::MAX_HTTP_BODY_BYTES {
103            return Self::respond(&mut writer, 413, &json!({"error":"request body too large"}))
104                .await;
105        }
106        let array_mode = head
107            .header("neon-array-mode")
108            .map(|v| v.eq_ignore_ascii_case("true"))
109            .unwrap_or(false);
110
111        let body_buf = match http_util::read_body(&mut reader, head.content_length, deadline).await
112        {
113            Ok(b) => b,
114            Err(_) => return Ok(()),
115        };
116        let req: Value = match serde_json::from_slice(&body_buf) {
117            Ok(v) => v,
118            Err(e) => {
119                return Self::respond(
120                    &mut writer,
121                    400,
122                    &json!({"error": format!("invalid JSON: {}", e)}),
123                )
124                .await
125            }
126        };
127        let sql = req
128            .get("query")
129            .and_then(|q| q.as_str())
130            .unwrap_or("")
131            .trim();
132        if sql.is_empty() {
133            return Self::respond(&mut writer, 400, &json!({"error":"missing 'query'"})).await;
134        }
135        let params = parse_params(req.get("params"));
136
137        match Self::run_sql(&cfg, sql, &params).await {
138            Ok(qr) => {
139                let body = neon_result(&qr, array_mode);
140                Self::respond(&mut writer, 200, &body).await
141            }
142            Err(e) => Self::respond(&mut writer, 400, &json!({ "error": e })).await,
143        }
144    }
145
146    async fn run_sql(
147        cfg: &HttpGatewayConfig,
148        sql: &str,
149        params: &[ParamValue],
150    ) -> std::result::Result<QueryResult, String> {
151        let bcfg = BackendConfig {
152            host: cfg.backend_host.clone(),
153            port: cfg.backend_port,
154            user: cfg.backend_user.clone(),
155            password: cfg.backend_password.clone(),
156            database: cfg.backend_database.clone(),
157            application_name: Some("heliosproxy-http".to_string()),
158            tls_mode: TlsMode::Disable,
159            connect_timeout: Duration::from_secs(5),
160            query_timeout: Duration::from_secs(30),
161            tls_config: default_client_config(),
162        };
163        let mut client = BackendClient::connect(&bcfg)
164            .await
165            .map_err(|e| format!("backend connect: {}", e))?;
166        let res = if params.is_empty() {
167            client.simple_query(sql).await
168        } else {
169            client.query_with_params(sql, params).await
170        };
171        client.close().await;
172        res.map_err(|e| format!("{}", e))
173    }
174
175    async fn respond(
176        writer: &mut tokio::net::tcp::WriteHalf<'_>,
177        status: u16,
178        body: &Value,
179    ) -> Result<()> {
180        let payload = serde_json::to_vec(body).unwrap_or_default();
181        let status_text = match status {
182            200 => "OK",
183            400 => "Bad Request",
184            401 => "Unauthorized",
185            405 => "Method Not Allowed",
186            _ => "Error",
187        };
188        let head = format!(
189            "HTTP/1.1 {} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
190            status, status_text, payload.len()
191        );
192        writer
193            .write_all(head.as_bytes())
194            .await
195            .map_err(|e| ProxyError::Network(format!("write: {}", e)))?;
196        writer
197            .write_all(&payload)
198            .await
199            .map_err(|e| ProxyError::Network(format!("write: {}", e)))?;
200        Ok(())
201    }
202}
203
204/// Map a JSON `params` array to text-format ParamValues.
205fn parse_params(v: Option<&Value>) -> Vec<ParamValue> {
206    match v.and_then(|v| v.as_array()) {
207        None => Vec::new(),
208        Some(arr) => arr
209            .iter()
210            .map(|p| match p {
211                Value::Null => ParamValue::Null,
212                Value::Bool(b) => ParamValue::Bool(*b),
213                Value::Number(n) => {
214                    if let Some(i) = n.as_i64() {
215                        ParamValue::Int(i)
216                    } else {
217                        ParamValue::Float(n.as_f64().unwrap_or(0.0))
218                    }
219                }
220                Value::String(s) => ParamValue::Text(s.clone()),
221                other => ParamValue::Text(other.to_string()),
222            })
223            .collect(),
224    }
225}
226
227/// Build the Neon-serverless-style result body.
228fn neon_result(qr: &QueryResult, array_mode: bool) -> Value {
229    let command = qr
230        .command_tag
231        .split_whitespace()
232        .next()
233        .unwrap_or("")
234        .to_string();
235    let fields: Vec<Value> = qr
236        .columns
237        .iter()
238        .map(|c| json!({ "name": c.name, "dataTypeID": c.type_oid }))
239        .collect();
240    let rows: Vec<Value> = qr
241        .rows
242        .iter()
243        .map(|row| {
244            if array_mode {
245                Value::Array(row.iter().map(cell_to_json).collect())
246            } else {
247                let mut obj = serde_json::Map::new();
248                for (i, c) in qr.columns.iter().enumerate() {
249                    let v = row.get(i).map(cell_to_json).unwrap_or(Value::Null);
250                    obj.insert(c.name.clone(), v);
251                }
252                Value::Object(obj)
253            }
254        })
255        .collect();
256    // rowCount mirrors libpq: affected rows for writes, else row count.
257    let row_count = qr.rows_affected().unwrap_or(qr.rows.len() as u64);
258    json!({
259        "command": command,
260        "rowCount": row_count,
261        "rows": rows,
262        "fields": fields,
263        "rowAsArray": array_mode,
264    })
265}
266
267fn cell_to_json(v: &TextValue) -> Value {
268    match v {
269        TextValue::Null => Value::Null,
270        TextValue::Text(s) => Value::String(s.clone()),
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277    use crate::backend::client::ColumnMeta;
278
279    fn qr() -> QueryResult {
280        QueryResult {
281            columns: vec![
282                ColumnMeta {
283                    name: "id".into(),
284                    type_oid: 23,
285                },
286                ColumnMeta {
287                    name: "name".into(),
288                    type_oid: 25,
289                },
290            ],
291            rows: vec![
292                vec![TextValue::Text("1".into()), TextValue::Text("alice".into())],
293                vec![TextValue::Text("2".into()), TextValue::Null],
294            ],
295            command_tag: "SELECT 2".into(),
296        }
297    }
298
299    #[test]
300    fn neon_object_mode() {
301        let v = neon_result(&qr(), false);
302        assert_eq!(v["command"], "SELECT");
303        assert_eq!(v["rowCount"], 2);
304        assert_eq!(v["rows"][0]["id"], "1");
305        assert_eq!(v["rows"][0]["name"], "alice");
306        assert_eq!(v["rows"][1]["name"], Value::Null);
307        assert_eq!(v["fields"][0]["name"], "id");
308        assert_eq!(v["fields"][0]["dataTypeID"], 23);
309    }
310
311    #[test]
312    fn neon_array_mode() {
313        let v = neon_result(&qr(), true);
314        assert_eq!(v["rowAsArray"], true);
315        assert_eq!(v["rows"][0][0], "1");
316        assert_eq!(v["rows"][0][1], "alice");
317    }
318
319    #[test]
320    fn params_mapping() {
321        let p = parse_params(Some(&json!([1, "x", true, null, 2.5])));
322        assert!(matches!(p[0], ParamValue::Int(1)));
323        assert!(matches!(p[1], ParamValue::Text(ref s) if s == "x"));
324        assert!(matches!(p[2], ParamValue::Bool(true)));
325        assert!(matches!(p[3], ParamValue::Null));
326        assert!(matches!(p[4], ParamValue::Float(_)));
327    }
328}
heliosdb_proxy/http_gateway.rs

heliosdb_proxy/
http_gateway.rs
