Improved listen_ports integration test (#3306)

Author: Razz4780

changelog.d/+listen-ports-test.internal.md

@@ -0,0 +1 @@
+Improved `listen_ports` integration test to reduce flakes.

mirrord/layer/tests/apps/listen_ports/src/main.rs

@@ -1,27 +1,33 @@
+//! Test application that listens on a sequence of localhost ports.
+//! On each port, it accepts one connection and expects to receive the string "HELLO" and a
+//! shutdown.
+//!
+//! Reads the semicolon-separated port list from the `APP_PORTS` environment variable.
+
 use std::{
     io::Read,
-    net::{SocketAddr, TcpListener},
+    net::{Ipv4Addr, SocketAddr, TcpListener},
 };
 
+const PORTS_ENV: &str = "APP_PORTS";
+
 fn main() {
-    let mut buf = [0_u8; 5];
+    let ports = std::env::var(PORTS_ENV)
+        .unwrap()
+        .split(';')
+        .map(|port| port.parse::<u16>().unwrap())
+        .collect::<Vec<_>>();
+
+    let mut buf = String::new();
 
-    let addr: SocketAddr = "127.0.0.1:80".parse().unwrap();
-    // Check `listen_ports` work
-    let listener = TcpListener::bind(addr).unwrap();
-    // verify user could connect on port 51222
-    let (mut conn, _) = listener.accept().unwrap();
-    conn.read_exact(buf.as_mut_slice()).unwrap();
-    assert_eq!(buf.as_slice(), b"HELLO");
+    for port in ports {
+        let addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), port);
+        let listener = TcpListener::bind(addr).unwrap();
 
-    // Check specific port but not in listen_ports works
-    let addr: SocketAddr = "127.0.0.1:40000".parse().unwrap();
-    let listener = TcpListener::bind(addr).unwrap();
-    // verify user could connect on port 40000
-    let (mut conn, _) = listener.accept().unwrap();
-    conn.read_exact(buf.as_mut_slice()).unwrap();
-    assert_eq!(buf.as_slice(), b"HELLO");
+        let (mut conn, _) = listener.accept().unwrap();
+        conn.read_to_string(&mut buf).unwrap();
+        assert_eq!(buf, "HELLO");
 
-    // Checking that random port works is hard because we need to make libc return an error
-    // so we miss that for now.
+        buf.clear();
+    }
 }

mirrord/layer/tests/configs/listen_ports.json

@@ -2,75 +2,109 @@
 #![feature(assert_matches)]
 #![warn(clippy::indexing_slicing)]
 
-use std::{assert_matches::assert_matches, path::Path, time::Duration};
+use std::{
+    assert_matches::assert_matches,
+    io::Write,
+    net::{Ipv4Addr, SocketAddr},
+    path::Path,
+    time::Duration,
+};
 
 use mirrord_protocol::{
     tcp::{DaemonTcp, LayerTcpSteal, StealType},
     ClientMessage, DaemonMessage,
 };
 use rstest::rstest;
-use tokio::io::AsyncWriteExt;
+use tokio::{io::AsyncWriteExt, net::TcpListener};
 
 mod common;
 
 pub use common::*;
 use tokio::net::TcpStream;
 
-/// Start an application (and load the layer into it) that listens on a port that is configured to
-/// be ignored, and verify that no messages are sent to the agent.
+/// Verifies that the layer respects `feature.network.incoming.listen_ports` mapping.
 #[rstest]
 #[tokio::test]
 #[timeout(Duration::from_secs(60))]
 async fn listen_ports(
     #[values(Application::RustListenPorts)] application: Application,
     dylib_path: &Path,
-    config_dir: &Path,
 ) {
-    let config_path = config_dir.join("listen_ports.json");
+    // We need to know ports on which the application listens,
+    // because we want to make the connections from the test code
+    // (to verify that the layer respects `feature.network.incoming.listen_ports` mapping).
+    // If we just use const ports, this test is flaky in the CI, as the app tends to fail with
+    // EADDRINUSE. Here we bind to random ports, drop the sockets, and hope that the ports
+    // remain free for the app to use.
+    let port_1 = TcpListener::bind(SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 0))
+        .await
+        .unwrap()
+        .local_addr()
+        .unwrap()
+        .port();
+    let port_2 = TcpListener::bind(SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 0))
+        .await
+        .unwrap()
+        .local_addr()
+        .unwrap()
+        .port();
+
+    let config = serde_json::json!({
+        "feature": {
+            "network": {
+                "incoming": {
+                    "mode": "steal",
+                    "listen_ports": [[80, port_1]]
+                }
+            },
+            "fs": "local"
+        }
+    });
+    // Ports for the application code to use.
+    // Due to the port mapping, the application should effectively listen on `port_1;port_2`.
+    let app_ports = format!("80;{port_2}");
+
+    let mut config_file = tempfile::NamedTempFile::with_suffix(".json").unwrap();
+    config_file
+        .as_file_mut()
+        .write_all(serde_json::to_string(&config).unwrap().as_bytes())
+        .unwrap();
+
     let (mut test_process, mut intproxy) = application
         .start_process_with_layer(
             dylib_path,
-            vec![
-                ("RUST_LOG", "mirrord=trace"),
-                ("MIRRORD_FILE_MODE", "local"),
-            ],
-            Some(&config_path),
+            vec![("RUST_LOG", "mirrord=trace"), ("APP_PORTS", &app_ports)],
+            Some(config_file.path()),
         )
         .await;
 
-    assert_matches!(
-        intproxy.recv().await,
-        ClientMessage::TcpSteal(LayerTcpSteal::PortSubscribe(StealType::All(80)))
-    );
-    intproxy
-        .send(DaemonMessage::TcpSteal(DaemonTcp::SubscribeResult(Ok(80))))
-        .await;
-    let mut stream = TcpStream::connect("127.0.0.1:51222").await.unwrap();
-    println!("connected to listener at port 51222");
-    stream.write_all(b"HELLO").await.unwrap();
-
-    assert_matches!(
-        intproxy.recv().await,
-        ClientMessage::TcpSteal(LayerTcpSteal::PortSubscribe(StealType::All(40000)))
-    );
-    intproxy
-        .send(DaemonMessage::TcpSteal(DaemonTcp::SubscribeResult(Ok(
-            40000,
-        ))))
-        .await;
-    let mut stream = TcpStream::connect("127.0.0.1:40000").await.unwrap();
-    println!("connected to listener at port 40000");
-    stream.write_all(b"HELLO").await.unwrap();
-
-    loop {
-        match intproxy.try_recv().await {
-            Some(ClientMessage::TcpSteal(LayerTcpSteal::PortUnsubscribe(40000))) => {}
-            Some(ClientMessage::TcpSteal(LayerTcpSteal::PortUnsubscribe(80))) => {}
-            None => break,
-            other => panic!("unexpected message: {:?}", other),
-        }
+    for (subscription_port, local_port) in [(80, port_1), (port_2, port_2)] {
+        assert_matches!(
+            intproxy.recv().await,
+            ClientMessage::TcpSteal(LayerTcpSteal::PortSubscribe(StealType::All(port)))
+            if port == subscription_port,
+        );
+        intproxy
+            .send(DaemonMessage::TcpSteal(DaemonTcp::SubscribeResult(Ok(
+                subscription_port,
+            ))))
+            .await;
+        let mut stream =
+            TcpStream::connect(SocketAddr::new(Ipv4Addr::LOCALHOST.into(), local_port))
+                .await
+                .unwrap();
+        println!("connected to the application on port {local_port} (application code used {subscription_port})");
+        stream.write_all(b"HELLO").await.unwrap();
+        stream.shutdown().await.unwrap();
+        assert_matches!(
+            intproxy.recv().await,
+            ClientMessage::TcpSteal(LayerTcpSteal::PortUnsubscribe(port))
+            if port == subscription_port,
+        );
     }
+
     assert_eq!(intproxy.try_recv().await, None);
+
     test_process.wait_assert_success().await;
     test_process.assert_no_error_in_stderr().await;
     test_process.assert_no_error_in_stdout().await;