fix(node): close the connection when the abort signal fires during the body read

impit-node/src/abortable_stream.rs

@@ -6,7 +6,7 @@ use std::task::Poll;
 
 /// A Tokio stream wrapper that can be aborted via a channel
 pub(crate) struct AbortableStream<S> {
-  inner: Pin<Box<S>>,
+  inner: Option<Pin<Box<S>>>,
   abort_rx: Arc<tokio::sync::Mutex<Option<tokio::sync::mpsc::Receiver<()>>>>,
 }
 
@@ -16,7 +16,7 @@ impl<S> AbortableStream<S> {
     abort_rx: Arc<tokio::sync::Mutex<Option<tokio::sync::mpsc::Receiver<()>>>>,
   ) -> Self {
     Self {
-      inner: Box::pin(inner),
+      inner: Some(Box::pin(inner)),
       abort_rx,
     }
   }
@@ -29,10 +29,14 @@ where
   type Item = Result<Vec<u8>>;
 
   fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-    if let Ok(mut guard) = self.abort_rx.try_lock() {
+    let this = &mut *self;
+    if let Ok(mut guard) = this.abort_rx.try_lock() {
       if let Some(ref mut rx) = *guard {
         match rx.poll_recv(cx) {
           Poll::Ready(Some(_)) => {
+            // napi only drops the stream on close or cancel, not on error, so drop the body here
+            // to close the connection.
+            this.inner = None;
             return Poll::Ready(Some(Err(napi::Error::new(
               napi::Status::GenericFailure,
               "impit: Request was aborted through AbortSignal",
@@ -43,6 +47,9 @@ where
       }
     }
 
-    self.inner.as_mut().poll_next(cx)
+    match this.inner.as_mut() {
+      Some(inner) => inner.as_mut().poll_next(cx),
+      None => Poll::Ready(None),
+    }
   }
 }

impit-node/test/basics.test.ts

@@ -767,6 +767,33 @@ describe.each([
 
             expect(getEventListeners(signal, 'abort').length).toBe(0);
         });
+
+        // https://github.com/apify/impit/issues/597
+        test('closes the connection when aborted during the body read', async () => {
+            let socketClosed: () => void = () => {};
+            const closed = new Promise<boolean>((resolve) => { socketClosed = () => resolve(true); });
+
+            const server = http.createServer((req, res) => {
+                const interval = setInterval(() => res.write('x'.repeat(1024)), 10);
+                res.on('close', () => { clearInterval(interval); socketClosed(); });
+            }).listen(0);
+
+            try {
+                const controller = new AbortController();
+                const response = await impit.fetch(`http://127.0.0.1:${(server.address() as AddressInfo).port}/`, { signal: controller.signal });
+                const reading = response.text();
+
+                await new Promise((resolve) => setTimeout(resolve, 100));
+                controller.abort();
+                await expect(reading).rejects.toThrow(/abort/i);
+
+                const timeout = new Promise<boolean>((resolve) => setTimeout(() => resolve(false), 500));
+                expect(await Promise.race([closed, timeout])).toBe(true);
+            } finally {
+                server.closeAllConnections();
+                server.close();
+            }
+        });
     });
 
     describe('Response.clone()', () => {