Add support for ASGI pathsend (#171)

src/asgi/io.rs

@@ -7,13 +7,17 @@ use http_body_util::BodyExt;
 use hyper::{
     body,
     header::{HeaderMap, HeaderName, HeaderValue, SERVER as HK_SERVER},
-    Response,
+    Response, StatusCode,
 };
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyDict};
 use std::{borrow::Cow, sync::Arc};
-use tokio::sync::{mpsc, oneshot, Mutex};
+use tokio::{
+    fs::File,
+    sync::{mpsc, oneshot, Mutex},
+};
 use tokio_tungstenite::WebSocketStream;
+use tokio_util::io::ReaderStream;
 use tungstenite::Message;
 
 use super::{
@@ -22,7 +26,7 @@ use super::{
 };
 use crate::{
     conversion::BytesToPy,
-    http::{HTTPRequest, HTTPResponse, HTTPResponseBody, HV_SERVER},
+    http::{response_404, HTTPRequest, HTTPResponse, HTTPResponseBody, HV_SERVER},
     runtime::{empty_future_into_py, future_into_py_futlike, future_into_py_iter, RuntimeRef},
     ws::{HyperWebsocket, UpgradeData},
 };
@@ -169,6 +173,29 @@ impl ASGIHTTPProtocol {
                     _ => error_flow!(),
                 }
             }
+            Ok(ASGIMessageType::HTTPFile) => match (self.response_started, adapt_file(py, data), self.tx.take()) {
+                (true, Ok(file_path), Some(tx)) => {
+                    let status = self.response_status.unwrap();
+                    let headers = self.response_headers.take().unwrap();
+                    future_into_py_iter(self.rt.clone(), py, async move {
+                        let res = match File::open(file_path).await {
+                            Ok(file) => {
+                                let stream = ReaderStream::new(file);
+                                let stream_body = http_body_util::StreamBody::new(stream.map_ok(body::Frame::data));
+                                let mut res =
+                                    Response::new(BodyExt::map_err(stream_body, std::convert::Into::into).boxed());
+                                *res.status_mut() = StatusCode::from_u16(status as u16).unwrap();
+                                *res.headers_mut() = headers;
+                                res
+                            }
+                            Err(_) => response_404(),
+                        };
+                        let _ = tx.send(res);
+                        Ok(())
+                    })
+                }
+                _ => error_flow!(),
+            },
             Err(err) => Err(err.into()),
             _ => error_message!(),
         }
@@ -315,6 +342,7 @@ fn adapt_message_type(message: &PyDict) -> Result<ASGIMessageType, UnsupportedAS
             match message_type {
                 "http.response.start" => Ok(ASGIMessageType::HTTPStart),
                 "http.response.body" => Ok(ASGIMessageType::HTTPBody),
+                "http.response.pathsend" => Ok(ASGIMessageType::HTTPFile),
                 "websocket.accept" => Ok(ASGIMessageType::WSAccept),
                 "websocket.close" => Ok(ASGIMessageType::WSClose),
                 "websocket.send" => Ok(ASGIMessageType::WSMessage),
@@ -364,6 +392,14 @@ fn adapt_body(py: Python, message: &PyDict) -> (Box<[u8]>, bool) {
     (body.into(), more)
 }
 
+#[inline(always)]
+fn adapt_file(py: Python, message: &PyDict) -> PyResult<String> {
+    match message.get_item(pyo3::intern!(py, "path"))? {
+        Some(item) => item.extract(),
+        _ => error_flow!(),
+    }
+}
+
 #[inline(always)]
 fn ws_message_into_rs(py: Python, message: &PyDict) -> PyResult<Message> {
     match (

src/asgi/types.rs

@@ -15,6 +15,7 @@ static ASGI_EXTENSIONS: OnceCell<PyObject> = OnceCell::new();
 pub(crate) enum ASGIMessageType {
     HTTPStart,
     HTTPBody,
+    HTTPFile,
     WSAccept,
     WSClose,
     WSMessage,
@@ -142,6 +143,7 @@ impl ASGIScope {
             ASGI_EXTENSIONS
                 .get_or_try_init(|| {
                     let rv = PyDict::new(py);
+                    rv.set_item("http.response.pathsend", PyDict::new(py))?;
                     Ok::<PyObject, PyErr>(rv.into())
                 })?
                 .as_ref(py),

tests/apps/asgi.py

@@ -1,4 +1,5 @@
 import json
+import pathlib
 
 
 PLAINTEXT_RESPONSE = {
@@ -7,6 +8,11 @@
     'headers': [[b'content-type', b'text/plain; charset=utf-8']],
 }
 JSON_RESPONSE = {'type': 'http.response.start', 'status': 200, 'headers': [[b'content-type', b'application/json']]}
+MEDIA_RESPONSE = {
+    'type': 'http.response.start',
+    'status': 200,
+    'headers': [[b'content-type', b'image/png'], [b'content-length', b'95']],
+}
 
 
 async def info(scope, receive, send):
@@ -24,6 +30,7 @@ async def info(scope, receive, send):
                     'path': scope['path'],
                     'query_string': scope['query_string'].decode('latin-1'),
                     'headers': {k.decode('utf8'): v.decode('utf8') for k, v in scope['headers']},
+                    'extensions': scope['extensions'],
                     'state': scope['state'],
                 }
             ).encode('utf8'),
@@ -43,6 +50,12 @@ async def echo(scope, receive, send):
     await send({'type': 'http.response.body', 'body': body, 'more_body': False})
 
 
+async def pathsend(scope, receive, send):
+    path = pathlib.Path.cwd() / 'tests' / 'fixtures' / 'media.png'
+    await send(MEDIA_RESPONSE)
+    await send({'type': 'http.response.pathsend', 'path': str(path)})
+
+
 async def ws_reject(scope, receive, send):
     return
 
@@ -116,6 +129,7 @@ def app(scope, receive, send):
     return {
         '/info': info,
         '/echo': echo,
+        '/file': pathsend,
         '/ws_reject': ws_reject,
         '/ws_info': ws_info,
         '/ws_echo': ws_echo,

tests/test_asgi.py

@@ -23,6 +23,7 @@ async def test_scope(asgi_server, threading_mode):
     assert data['path'] == '/info'
     assert data['query_string'] == 'test=true'
     assert data['headers']['host'] == f'localhost:{port}'
+    assert 'http.response.pathsend' in data['extensions']
     assert data['state']['global'] == 'test'
 
 
@@ -65,3 +66,14 @@ async def test_protocol_error(asgi_server, threading_mode):
         res = httpx.get(f'http://localhost:{port}/err_proto')
 
     assert res.status_code == 500
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('threading_mode', ['runtime', 'workers'])
+async def test_file(asgi_server, threading_mode):
+    async with asgi_server(threading_mode) as port:
+        res = httpx.get(f'http://localhost:{port}/file')
+
+    assert res.status_code == 200
+    assert res.headers['content-type'] == 'image/png'
+    assert res.headers['content-length'] == '95'