File size: 6,421 Bytes
854994d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
#![allow(clippy::expect_used)]
//! Exercise realtime TLS selection without mutating the test process's trust environment.

use std::io;
use std::net::TcpListener;
use std::process::Command;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;

use codex_api::Provider;
use codex_api::RealtimeEventParser;
use codex_api::RealtimeOutputModality;
use codex_api::RealtimeSessionConfig;
use codex_api::RealtimeSessionMode;
use codex_api::RealtimeWebsocketClient;
use codex_api::RetryConfig;
use codex_protocol::protocol::RealtimeVoice;
use http::HeaderMap;
use pretty_assertions::assert_eq;

const ADDRESS_ENV: &str = "CODEX_TEST_REALTIME_TLS_ADDRESS";
const TRUST_ENV: &str = "CODEX_TEST_REALTIME_TLS_TRUST";

#[test]
fn realtime_tls_selects_system_and_custom_trust() {
    if let Ok(address) = std::env::var(ADDRESS_ENV) {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap();
        runtime.block_on(check_connection(address));
        return;
    }

    codex_utils_rustls_provider::ensure_rustls_crypto_provider();
    let certificate = rcgen::generate_simple_self_signed(vec!["localhost".to_string()]).unwrap();
    let temp = tempfile::TempDir::new().unwrap();
    let ca = temp.path().join("ca.pem");
    std::fs::write(&ca, certificate.cert.pem()).unwrap();
    let config = Arc::new(
        rustls::ServerConfig::builder()
            .with_no_client_auth()
            .with_single_cert(
                vec![certificate.cert.der().clone()],
                certificate.signing_key.into(),
            )
            .unwrap(),
    );

    for trust in ["system", "custom"] {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let address = listener.local_addr().unwrap();
        listener.set_nonblocking(/*nonblocking*/ true).unwrap();
        let config = config.clone();
        let server = std::thread::spawn(move || {
            let deadline = Instant::now() + Duration::from_secs(/*secs*/ 30);
            let stream = loop {
                match listener.accept() {
                    Ok((stream, _)) => break stream,
                    Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
                        assert!(Instant::now() < deadline, "TLS client did not connect");
                        std::thread::sleep(Duration::from_millis(/*millis*/ 10));
                    }
                    Err(error) => panic!("TLS accept failed: {error}"),
                }
            };
            stream.set_nonblocking(/*nonblocking*/ false).unwrap();
            stream
                .set_read_timeout(Some(Duration::from_secs(/*secs*/ 30)))
                .unwrap();
            stream
                .set_write_timeout(Some(Duration::from_secs(/*secs*/ 30)))
                .unwrap();
            let tls =
                rustls::StreamOwned::new(rustls::ServerConnection::new(config).unwrap(), stream);
            match tungstenite::accept(tls) {
                Ok(mut socket) => {
                    assert_eq!(trust, "custom", "system trust must reject the generated CA");
                    let message = socket.read().unwrap().into_text().unwrap();
                    let update: serde_json::Value = serde_json::from_str(&message).unwrap();
                    assert_eq!(update["type"], "session.update");
                }
                Err(_) => assert_eq!(trust, "system", "custom CA must allow WSS"),
            }
        });
        // Re-execute this test so each connection uses its own CA environment, including on Windows.
        let mut child = Command::new(std::env::current_exe().unwrap());
        child
            .args([
                "--exact",
                "realtime_tls_selects_system_and_custom_trust",
                "--nocapture",
            ])
            .env(ADDRESS_ENV, format!("localhost:{}", address.port()))
            .env(TRUST_ENV, trust)
            .env_remove("CODEX_CA_CERTIFICATE")
            .env_remove("SSL_CERT_FILE");
        if trust == "custom" {
            child.env("CODEX_CA_CERTIFICATE", &ca).env(
                "SSL_CERT_FILE",
                temp.path().join("missing-lower-priority-ca.pem"),
            );
        }
        let output = child.output().unwrap();
        server.join().unwrap();
        assert!(
            output.status.success(),
            "{trust}: {}\n{}",
            String::from_utf8_lossy(&output.stdout),
            String::from_utf8_lossy(&output.stderr)
        );
    }
}

async fn check_connection(address: String) {
    let client = RealtimeWebsocketClient::new(Provider {
        name: "local TLS test".into(),
        base_url: format!("https://{address}"),
        query_params: None,
        headers: HeaderMap::new(),
        retry: RetryConfig {
            max_attempts: 1,
            base_delay: Duration::from_millis(/*millis*/ 1),
            retry_429: false,
            retry_5xx: false,
            retry_transport: false,
        },
        stream_idle_timeout: Duration::from_secs(/*secs*/ 5),
    });
    let result = tokio::time::timeout(
        Duration::from_secs(/*secs*/ 20),
        client.connect(
            RealtimeSessionConfig {
                instructions: "TLS test".into(),
                initial_items: Vec::new(),
                delegation_ack_filler: None,
                model: Some("realtime-test-model".into()),
                session_id: None,
                event_parser: RealtimeEventParser::V1,
                session_mode: RealtimeSessionMode::Conversational,
                output_modality: RealtimeOutputModality::Audio,
                voice: RealtimeVoice::Cove,
            },
            HeaderMap::new(),
            HeaderMap::new(),
        ),
    )
    .await
    .expect("TLS connection should finish");
    match std::env::var(TRUST_ENV)
        .expect("trust scenario should be set")
        .as_str()
    {
        "custom" => {
            result.expect("configured CA should permit realtime WSS");
        }
        "system" => {
            let error = match result {
                Ok(_) => panic!("system trust accepted an untrusted certificate"),
                Err(error) => error.to_string(),
            };
            assert!(error.to_lowercase().contains("certificate"), "{error}");
        }
        other => panic!("unknown trust scenario: {other}"),
    }
}