fix(realtime): stop passing the library's 1002 upstream as the client's close code

crates/aisix-proxy/src/realtime.rs

@@ -1011,13 +1011,26 @@ async fn run_session(
                     }
                 }
                 Some(Ok(AxMessage::Close(frame))) => {
-                    // The transport has already queued the reply echoing
-                    // the client's code; it goes out on the flush, and
-                    // without one the client sees 1006.
+                    // The transport has already queued its reply (the
+                    // client's code, or the library's 1002 for one that may
+                    // not be sent); it goes out on the flush, and without
+                    // one the client sees 1006.
                     let _ = client_tx.close().await;
-                    let frame = frame.map(|f| TgCloseFrame {
-                        code: f.code.into(),
-                        reason: f.reason.to_string().into(),
+                    // A client code that may not be sent arrives as the
+                    // library's substitute; the upstream gets 1011, not a
+                    // 1002 the client never sent.
+                    let frame = frame.map(|f| {
+                        if (f.code, f.reason.as_ref()) == LIBRARY_REJECTED_CLOSE {
+                            TgCloseFrame {
+                                code: 1011.into(),
+                                reason: "".into(),
+                            }
+                        } else {
+                            TgCloseFrame {
+                                code: f.code.into(),
+                                reason: f.reason.to_string().into(),
+                            }
+                        }
                     });
                     close_upstream(&mut up_tx, &mut up_rx, frame, &mut usage).await;
                     break;
@@ -1297,9 +1310,10 @@ enum UpstreamEnd<'a> {
     Dropped,
 }
 
-/// tungstenite replaces a received Close whose code may not be sent
-/// (1004, 1016–2999, …) with exactly this frame and drops the original
-/// code, so it is indistinguishable from an upstream sending it itself.
+/// tungstenite (under both the upstream client and axum's server side)
+/// replaces a received Close whose code may not be sent (1004, 1016–2999,
+/// …) with exactly this frame and drops the original code, so it is
+/// indistinguishable from the peer sending it itself.
 const LIBRARY_REJECTED_CLOSE: (u16, &str) = (1002, "Protocol violation");
 
 /// Relay the upstream's end of the session to the client: its code and
@@ -1487,25 +1501,30 @@ mod tests {
     async fn library_rejected_close_matches_what_tungstenite_substitutes() {
         use futures::StreamExt;
         use tokio::io::AsyncWriteExt;
-        let (local, mut remote) = tokio::io::duplex(64);
-        let mut ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
-            local,
-            tokio_tungstenite::tungstenite::protocol::Role::Client,
-            None,
-        )
-        .await;
-        // Unmasked server Close carrying 1004, a code that may not be sent.
-        remote
-            .write_all(&[0x88, 0x04, 0x03, 0xEC, b'r', b'1'])
-            .await
-            .unwrap();
-        let Some(Ok(TgMessage::Close(Some(frame)))) = ws.next().await else {
-            panic!("expected a Close frame");
-        };
-        assert_eq!(
-            (u16::from(frame.code), frame.reason.as_ref()),
-            super::LIBRARY_REJECTED_CLOSE
-        );
+        use tokio_tungstenite::tungstenite::protocol::Role;
+        // A Close carrying 1004, a code that may not be sent: unmasked as
+        // the upstream sends it to the gateway's client role, and masked
+        // (all-zero key) as a client sends it to axum's server role.
+        for (role, frame) in [
+            (Role::Client, &[0x88, 0x04, 0x03, 0xEC, b'r', b'1'][..]),
+            (
+                Role::Server,
+                &[0x88, 0x84, 0, 0, 0, 0, 0x03, 0xEC, b'r', b'1'][..],
+            ),
+        ] {
+            let (local, mut remote) = tokio::io::duplex(64);
+            let mut ws =
+                tokio_tungstenite::WebSocketStream::from_raw_socket(local, role, None).await;
+            remote.write_all(frame).await.unwrap();
+            let Some(Ok(TgMessage::Close(Some(close)))) = ws.next().await else {
+                panic!("expected a Close frame ({role:?})");
+            };
+            assert_eq!(
+                (u16::from(close.code), close.reason.as_ref()),
+                super::LIBRARY_REJECTED_CLOSE,
+                "{role:?}"
+            );
+        }
     }
 
     /// JSON lets both a key and a string value be spelled with escapes, and

tests/e2e/src/cases/realtime-close-e2e.test.ts

@@ -154,6 +154,22 @@ describe("/v1/realtime close handshake", () => {
     expect(upstream!.closesSeen[before]).toBe(1000);
   });
 
+  test("a client close with a code that may not be sent reaches the upstream as 1011", async (ctx) => {
+    if (!ready(ctx)) return;
+    const before = upstream!.closesSeen.length;
+    const s = await openSession(app!);
+    // `ws` refuses to send 1004, so write the masked Close frame by hand
+    // (an all-zero mask leaves the payload as is).
+    const payload = Buffer.concat([Buffer.from([1004 >> 8, 1004 & 0xff]), Buffer.from("r1004")]);
+    (s.ws as unknown as { _socket: NodeJS.WritableStream })._socket.write(
+      Buffer.concat([Buffer.from([0x88, 0x80 | payload.length, 0, 0, 0, 0]), payload]),
+    );
+    // The client's own close is still answered rather than reset.
+    expect((await s.closed).code).not.toBe(1006);
+    await expect.poll(() => upstream!.closesSeen.length).toBe(before + 1);
+    expect(upstream!.closesSeen[before]).toBe(1011);
+  });
+
   test("a normal upstream close reaches the client as is, with no error event", async (ctx) => {
     if (!ready(ctx)) return;
     const before = upstream!.closesSeen.length;