mirror of
https://github.com/openai/codex.git
synced 2026-11-12 21:20:34 +00:00
Add policy-checked TCP upstream tunneling for SOCKS (#60523)
## Why Allowing a hostname must not bypass deny or review rules on its resolved IP addresses. Forwarding DNS resolution to an upstream proxy also requires checking whether those address restrictions can still be enforced. ## What changed - Add `TcpUpstreamProxy` and `set_tcp_upstream_proxy` to require a configured HTTP `CONNECT` next hop for opaque SOCKS TCP traffic. Wait for a successful response before forwarding payloads and preserve bytes read past the response headers. - Enforce resolved IP deny and review rules at direct connection time, carry destination-scoped authorization through HTTP, SOCKS, and MITM forwarding, and report blocked resolutions as policy denials. - Recognize equivalent numeric IP spellings for restrictive rules. Permit remote TCP DNS only when local binding is allowed and no deny or review pattern can restrict IP addresses; reject upstream proxy loops. - Restrict Linux proxy bridges to managed proxy environment keys and supplied loopback ports. Bound concurrent bridge workers and release them after connection failures while allowing queued responses to drain. - Add test-only Linux DNS name mapping and SOCKS relay helpers with stable synthetic addresses and validation before opening upstream connections. ## Testing Add regression coverage for CONNECT refusal and buffered tunnel data, resolved-address policy denials, approval handling across policy refreshes, equivalent IP literals, DNS mapping stability, rejected relay requests, and bridge cleanup after resets. GitOrigin-RevId: bdd6e16de3373d0eb39995f689dad0230680b99e
This commit is contained in:
viyatb-oai
@@ -19,6 +19,10 @@ mod fd_mount;
|
||||
mod launcher;
|
||||
#[cfg(target_os = "linux")]
|
||||
mod linux_run_main;
|
||||
#[cfg(all(target_os = "linux", test))]
|
||||
mod managed_tcp;
|
||||
#[cfg(all(target_os = "linux", test))]
|
||||
mod managed_tcp_dns;
|
||||
#[cfg(target_os = "linux")]
|
||||
mod proxy_lifecycle;
|
||||
#[cfg(target_os = "linux")]
|
||||
|
||||
@@ -266,8 +266,11 @@ pub fn run_main() -> ! {
|
||||
// Outer stage: bubblewrap first, then re-enter this binary in the
|
||||
// sandboxed environment to apply seccomp. This path never falls back
|
||||
// to legacy Landlock on failure.
|
||||
let (proxy_route_spec, proxy_controls) = if allow_network_for_proxy {
|
||||
let (proxy_route_spec, controls) = prepare_host_proxy_route_spec()
|
||||
let (proxy_route_spec, proxy_controls) = if let Some(context) = managed_network.as_ref() {
|
||||
// Older explicit-proxy callers may omit the port list.
|
||||
let ports =
|
||||
(!context.loopback_ports.is_empty()).then_some(context.loopback_ports.as_slice());
|
||||
let (proxy_route_spec, controls) = prepare_host_proxy_route_spec(ports)
|
||||
.unwrap_or_else(|err| panic!("failed to prepare host proxy routing bridge: {err}"));
|
||||
(Some(proxy_route_spec), controls)
|
||||
} else {
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
//! Translate a TUN worker’s TCP destination to the original hostname or literal IP
|
||||
//! for the managed SOCKS5 bridge.
|
||||
|
||||
use crate::managed_tcp_dns::Names;
|
||||
use crate::proxy_routing::proxy_bidirectional;
|
||||
use socks5_protocol::Address;
|
||||
use socks5_protocol::AuthMethod;
|
||||
use socks5_protocol::AuthRequest;
|
||||
use socks5_protocol::AuthResponse;
|
||||
use socks5_protocol::Command as SocksCommand;
|
||||
use socks5_protocol::CommandReply;
|
||||
use socks5_protocol::CommandRequest;
|
||||
use socks5_protocol::CommandResponse;
|
||||
use socks5_protocol::Version;
|
||||
use socks5_protocol::sync::FromIO;
|
||||
use std::io;
|
||||
use std::io::Write;
|
||||
use std::net::SocketAddr;
|
||||
use std::net::TcpStream;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
|
||||
const TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
fn relay(
|
||||
mut client: TcpStream,
|
||||
proxy: SocketAddr,
|
||||
names: &Mutex<Names>,
|
||||
) -> socks5_protocol::Result<()> {
|
||||
client.set_read_timeout(Some(TIMEOUT))?;
|
||||
client.set_write_timeout(Some(TIMEOUT))?;
|
||||
Version::read_from(&mut client)?;
|
||||
let offered = AuthRequest::read_from(&mut client)?;
|
||||
if !offered.0.contains(&AuthMethod::Noauth) {
|
||||
return Err(io::Error::other("unsupported SOCKS greeting").into());
|
||||
}
|
||||
Version::V5.write_to(&mut client)?;
|
||||
AuthResponse::new(AuthMethod::Noauth).write_to(&mut client)?;
|
||||
|
||||
let upstream = match connect(&mut client, proxy, names) {
|
||||
Ok(stream) => stream,
|
||||
Err(err) => {
|
||||
let denied = CommandResponse::reply_error(CommandReply::ConnectionNotAllowedByRuleset);
|
||||
let _ = write_frame(&mut client, &denied);
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
write_frame(&mut client, &CommandResponse::success(Address::default()))?;
|
||||
client.set_read_timeout(None)?;
|
||||
client.set_write_timeout(None)?;
|
||||
upstream.set_read_timeout(None)?;
|
||||
upstream.set_write_timeout(None)?;
|
||||
Ok(proxy_bidirectional(upstream, client)?)
|
||||
}
|
||||
|
||||
fn connect(
|
||||
client: &mut impl io::Read,
|
||||
proxy: SocketAddr,
|
||||
names: &Mutex<Names>,
|
||||
) -> socks5_protocol::Result<TcpStream> {
|
||||
let request = CommandRequest::read_from(client)?;
|
||||
let (SocksCommand::Connect, Address::SocketAddr(target)) = (request.command, request.address)
|
||||
else {
|
||||
return Err(io::Error::other("expected TCP CONNECT to a TUN destination IP").into());
|
||||
};
|
||||
let destination = match names
|
||||
.lock()
|
||||
.map_err(|_| io::Error::other("DNS name table is unavailable"))?
|
||||
.lookup(target.ip())?
|
||||
{
|
||||
Some(host) => Address::Domain(host, target.port()),
|
||||
None => Address::SocketAddr(target),
|
||||
};
|
||||
|
||||
let mut upstream = TcpStream::connect_timeout(&proxy, TIMEOUT)?;
|
||||
upstream.set_read_timeout(Some(TIMEOUT))?;
|
||||
upstream.set_write_timeout(Some(TIMEOUT))?;
|
||||
Version::V5.write_to(&mut upstream)?;
|
||||
AuthRequest::new([AuthMethod::Noauth]).write_to(&mut upstream)?;
|
||||
Version::read_from(&mut upstream)?;
|
||||
if AuthResponse::read_from(&mut upstream)?.method() != AuthMethod::Noauth {
|
||||
return Err(io::Error::other("managed SOCKS5 handshake failed").into());
|
||||
}
|
||||
write_frame(&mut upstream, &CommandRequest::connect(destination))?;
|
||||
// The trusted proxy owns the approval deadline. The relay's handshake
|
||||
// timeout applies to framing and negotiation, not the human decision.
|
||||
upstream.set_read_timeout(None)?;
|
||||
let reply = CommandResponse::read_from(&mut upstream)?;
|
||||
if reply.reply != CommandReply::Succeeded
|
||||
|| matches!(reply.address, Address::Domain(ref name, _) if name.is_empty())
|
||||
{
|
||||
return Err(io::Error::other("invalid managed SOCKS5 success reply").into());
|
||||
}
|
||||
Ok(upstream)
|
||||
}
|
||||
|
||||
// Buffer the codec's field-by-field output before calling write_all. A denial
|
||||
// reply is best effort; the authorization decision does not depend on its delivery.
|
||||
fn write_frame(stream: &mut TcpStream, message: &impl FromIO) -> socks5_protocol::Result<()> {
|
||||
let mut bytes = Vec::new();
|
||||
message.write_to(&mut bytes)?;
|
||||
Ok(stream.write_all(&bytes)?)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "managed_tcp_relay_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,117 @@
|
||||
//! Assign stable per-command addresses to names. Addresses never change owners
|
||||
//! while the command runs, including after the client's DNS cache expires.
|
||||
|
||||
use hickory_proto::op::Header;
|
||||
use hickory_proto::op::Message;
|
||||
use hickory_proto::op::MessageType;
|
||||
use hickory_proto::op::OpCode;
|
||||
use hickory_proto::op::Query;
|
||||
use hickory_proto::op::ResponseCode;
|
||||
use hickory_proto::rr::DNSClass;
|
||||
use hickory_proto::rr::RData;
|
||||
use hickory_proto::rr::Record;
|
||||
use hickory_proto::rr::RecordType;
|
||||
use hickory_proto::serialize::binary::BinDecodable;
|
||||
use hickory_proto::serialize::binary::BinDecoder;
|
||||
use hickory_proto::serialize::binary::BinEncodable;
|
||||
use indexmap::IndexSet;
|
||||
use std::io;
|
||||
use std::net::IpAddr;
|
||||
use std::net::Ipv4Addr;
|
||||
use std::net::Ipv6Addr;
|
||||
|
||||
const MAX_NAMES: usize = 8192;
|
||||
const V4_BASE: u32 = u32::from_be_bytes([198, 18, 0, 0]);
|
||||
const V6_BASE: u128 = 0xfdc0_de00_0000_0000_0000_0000_0000_0000;
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct Names {
|
||||
// Append only: both synthetic address families encode the same one-based index.
|
||||
names: IndexSet<String>,
|
||||
}
|
||||
|
||||
impl Names {
|
||||
pub(super) fn lookup(&self, ip: IpAddr) -> io::Result<Option<String>> {
|
||||
let id = match ip.to_canonical() {
|
||||
IpAddr::V4(ip) if u32::from(ip) & 0xfffe_0000 == V4_BASE => u32::from(ip) - V4_BASE,
|
||||
IpAddr::V6(ip) if u128::from(ip) >> 32 == V6_BASE >> 32 => u128::from(ip) as u32,
|
||||
_ => return Ok(None),
|
||||
};
|
||||
let name = id
|
||||
.checked_sub(1)
|
||||
.and_then(|index| self.names.get_index(index as usize));
|
||||
name.cloned()
|
||||
.map(Some)
|
||||
.ok_or_else(|| io::Error::other("unassigned synthetic address"))
|
||||
}
|
||||
|
||||
fn id(&mut self, host: String) -> Option<u32> {
|
||||
let index = match self.names.get_index_of(&host) {
|
||||
Some(index) => index,
|
||||
None if self.names.len() < MAX_NAMES => self.names.insert_full(host).0,
|
||||
None => return None,
|
||||
};
|
||||
Some(index as u32 + 1)
|
||||
}
|
||||
|
||||
fn answer(&mut self, packet: &[u8]) -> Option<Vec<u8>> {
|
||||
if packet.len() > 4096 {
|
||||
return None;
|
||||
}
|
||||
// Decode one question only. EDNS and other trailing bytes never go to a resolver
|
||||
// or into the reply, and untrusted section counts never size an allocation.
|
||||
let mut decoder = BinDecoder::new(packet);
|
||||
let header = Header::read(&mut decoder).ok()?;
|
||||
if header.message_type != MessageType::Query
|
||||
|| header.op_code != OpCode::Query
|
||||
|| header.counts.queries != 1
|
||||
|| header.counts.answers != 0
|
||||
|| header.counts.authorities != 0
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let start = decoder.index();
|
||||
let query = Query::read(&mut decoder).ok()?;
|
||||
// The first question has no earlier name to reference. Reject compression
|
||||
// and labels that would require escapes when converted to a SOCKS hostname.
|
||||
if query.to_bytes().ok()? != decoder.slice_from(start).ok()?
|
||||
|| !query
|
||||
.name
|
||||
.iter()
|
||||
.flatten()
|
||||
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_'))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut reply = Message::response(header.id, OpCode::Query);
|
||||
reply.metadata.recursion_desired = header.recursion_desired;
|
||||
reply.metadata.recursion_available = true;
|
||||
if query.query_class == DNSClass::IN
|
||||
&& matches!(query.query_type, RecordType::A | RecordType::AAAA)
|
||||
&& !query.name.is_root()
|
||||
{
|
||||
let host = query
|
||||
.name
|
||||
.to_ascii()
|
||||
.trim_end_matches('.')
|
||||
.to_ascii_lowercase();
|
||||
if let Some(id) = self.id(host) {
|
||||
let ip = if query.query_type == RecordType::A {
|
||||
IpAddr::V4(Ipv4Addr::from(V4_BASE + id))
|
||||
} else {
|
||||
IpAddr::V6(Ipv6Addr::from(V6_BASE + u128::from(id)))
|
||||
};
|
||||
reply.add_answer(Record::from_rdata(query.name.clone(), 60, RData::from(ip)));
|
||||
} else {
|
||||
reply.metadata.response_code = ResponseCode::ServFail;
|
||||
}
|
||||
}
|
||||
reply.add_query(query);
|
||||
reply.to_vec().ok()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "managed_tcp_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,72 @@
|
||||
#![cfg(test)]
|
||||
|
||||
//! Verify rejected TUN requests cannot open an upstream connection.
|
||||
|
||||
use super::*;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::io::Read;
|
||||
use std::io::Write;
|
||||
use std::net::Ipv4Addr;
|
||||
use std::net::TcpListener;
|
||||
use std::thread;
|
||||
|
||||
#[test]
|
||||
fn invalid_tun_requests_never_open_upstream() {
|
||||
let proxy = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
|
||||
let addr = proxy.local_addr().unwrap();
|
||||
for request in [
|
||||
[5, 3, 0, 1, 127, 0, 0, 1, 0, 53],
|
||||
[5, 2, 0, 1, 127, 0, 0, 1, 0, 80],
|
||||
[5, 1, 1, 1, 127, 0, 0, 1, 0, 80],
|
||||
[5, 1, 0, 3, 3, b'd', b'b', b'x', 0, 80],
|
||||
[5, 1, 0, 1, 198, 18, 0, 99, 1, 187],
|
||||
] {
|
||||
assert!(connect(&mut request.as_slice(), addr, &Mutex::new(Names::default())).is_err());
|
||||
}
|
||||
proxy.set_nonblocking(true).unwrap();
|
||||
assert_eq!(
|
||||
proxy.accept().unwrap_err().kind(),
|
||||
io::ErrorKind::WouldBlock
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn denied_socks_connect_never_forwards_pipelined_payload() {
|
||||
let proxy = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
|
||||
rustix::net::sockopt::set_socket_timeout(
|
||||
&proxy,
|
||||
rustix::net::sockopt::Timeout::Recv,
|
||||
Some(TIMEOUT),
|
||||
)
|
||||
.unwrap();
|
||||
let proxy_addr = proxy.local_addr().unwrap();
|
||||
let incoming = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
|
||||
let mut client = TcpStream::connect(incoming.local_addr().unwrap()).unwrap();
|
||||
let (server, _) = incoming.accept().unwrap();
|
||||
let task = thread::spawn(move || relay(server, proxy_addr, &Mutex::new(Names::default())));
|
||||
|
||||
// A malicious client can send application data before either handshake completes.
|
||||
let connect = [5, 1, 0, 1, 192, 0, 2, 7, 0x15, 0x38];
|
||||
client.write_all(&[5, 2, 2, 0]).unwrap();
|
||||
client.write_all(&connect).unwrap();
|
||||
client.write_all(b"private marker").unwrap();
|
||||
client.shutdown(std::net::Shutdown::Write).unwrap();
|
||||
|
||||
let (mut upstream, _) = proxy.accept().unwrap();
|
||||
upstream.set_read_timeout(Some(TIMEOUT)).unwrap();
|
||||
let mut greeting = [0; 3];
|
||||
upstream.read_exact(&mut greeting).unwrap();
|
||||
assert_eq!(greeting, [5, 1, 0]);
|
||||
upstream.write_all(&[5, 0]).unwrap();
|
||||
let mut request = [0; 10];
|
||||
upstream.read_exact(&mut request).unwrap();
|
||||
assert_eq!(request, connect);
|
||||
upstream.write_all(&[5, 2, 0, 1, 0, 0, 0, 0, 0, 0]).unwrap();
|
||||
|
||||
let mut response = [0; 12]; // greeting and refusal sent to the client
|
||||
client.read_exact(&mut response).unwrap();
|
||||
assert_eq!(&response[..4], &[5, 0, 5, 2]);
|
||||
let mut byte = [0];
|
||||
assert_eq!(upstream.read(&mut byte).unwrap(), 0);
|
||||
assert!(task.join().unwrap().is_err());
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
#![cfg(test)]
|
||||
|
||||
//! Check DNS identity across address families, cache exhaustion and literal-address requests.
|
||||
|
||||
use super::*;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
fn query(name: &str, kind: RecordType) -> Vec<u8> {
|
||||
let mut packet = Message::new(/*id*/ 1, MessageType::Query, OpCode::Query);
|
||||
packet.add_query(Query::query(name.parse().unwrap(), kind));
|
||||
packet.to_vec().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn both_families_retain_the_original_name_and_literal_ips_are_separate() {
|
||||
let mut names = Names::default();
|
||||
let a = names.answer(&query("Db.Example", RecordType::A)).unwrap();
|
||||
let aaaa = names
|
||||
.answer(&query("db.example", RecordType::AAAA))
|
||||
.unwrap();
|
||||
let ip4 = Ipv4Addr::from(<[u8; 4]>::try_from(&a[a.len() - 4..]).unwrap());
|
||||
let ip6 = Ipv6Addr::from(<[u8; 16]>::try_from(&aaaa[aaaa.len() - 16..]).unwrap());
|
||||
for address in [
|
||||
IpAddr::V4(ip4),
|
||||
IpAddr::V6(ip6),
|
||||
IpAddr::V6(ip4.to_ipv6_mapped()),
|
||||
] {
|
||||
assert_eq!(
|
||||
names.lookup(address).unwrap().as_deref(),
|
||||
Some("db.example")
|
||||
);
|
||||
}
|
||||
assert_eq!(names.lookup("203.0.113.19".parse().unwrap()).unwrap(), None);
|
||||
assert_eq!(names.lookup("2001:db8::19".parse().unwrap()).unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capacity_exhaustion_never_reassigns_an_existing_address() {
|
||||
let mut names = Names::default();
|
||||
let first = names
|
||||
.answer(&query("first.example", RecordType::A))
|
||||
.unwrap();
|
||||
for n in 1..MAX_NAMES {
|
||||
names
|
||||
.answer(&query(&format!("n{n}.example"), RecordType::A))
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(
|
||||
names.answer(&query("new.example", RecordType::A)).unwrap()[3] & 15,
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
names
|
||||
.answer(&query("first.example", RecordType::A))
|
||||
.unwrap(),
|
||||
first
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_and_unsupported_questions_cannot_allocate_an_alias() {
|
||||
let mut names = Names::default();
|
||||
assert!(names.answer(&[0; 11]).is_none());
|
||||
let mut packet = query("valid.example", RecordType::A);
|
||||
packet[12] = 0xc0;
|
||||
assert!(names.answer(&packet).is_none());
|
||||
let mut counts = query("valid.example", RecordType::A);
|
||||
counts[6..12].fill(0xff);
|
||||
assert!(names.answer(&counts).is_none());
|
||||
let mut header_pointer = query("valid.example", RecordType::A);
|
||||
header_pointer.splice(12..27, [0xc0, 0x00]);
|
||||
assert!(names.answer(&header_pointer).is_none());
|
||||
assert_eq!(
|
||||
names
|
||||
.answer(&query("valid.example", RecordType::MX))
|
||||
.unwrap()[7],
|
||||
0
|
||||
);
|
||||
assert!(names.names.is_empty());
|
||||
}
|
||||
@@ -23,28 +23,12 @@ use std::os::fd::AsRawFd;
|
||||
use std::os::fd::FromRawFd;
|
||||
use std::os::fd::OwnedFd;
|
||||
use std::os::unix::net::UnixStream;
|
||||
use std::sync::Arc;
|
||||
use std::sync::mpsc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::Semaphore;
|
||||
use url::Url;
|
||||
|
||||
const PROXY_ENV_KEYS: &[&str] = &[
|
||||
"HTTP_PROXY",
|
||||
"HTTPS_PROXY",
|
||||
"WS_PROXY",
|
||||
"WSS_PROXY",
|
||||
"ALL_PROXY",
|
||||
"FTP_PROXY",
|
||||
"YARN_HTTP_PROXY",
|
||||
"YARN_HTTPS_PROXY",
|
||||
"NPM_CONFIG_HTTP_PROXY",
|
||||
"NPM_CONFIG_HTTPS_PROXY",
|
||||
"NPM_CONFIG_PROXY",
|
||||
"BUNDLE_HTTP_PROXY",
|
||||
"BUNDLE_HTTPS_PROXY",
|
||||
"PIP_PROXY",
|
||||
"DOCKER_HTTP_PROXY",
|
||||
"DOCKER_HTTPS_PROXY",
|
||||
];
|
||||
|
||||
const HOST_BRIDGE_READY: u8 = 1;
|
||||
const LOOPBACK_INTERFACE_NAME: &[u8] = b"lo";
|
||||
const HANDOFF_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 30);
|
||||
@@ -72,7 +56,9 @@ struct ProxyRoutePlan {
|
||||
has_proxy_config: bool,
|
||||
}
|
||||
|
||||
pub(crate) fn prepare_host_proxy_route_spec() -> io::Result<(String, Vec<File>)> {
|
||||
pub(crate) fn prepare_host_proxy_route_spec(
|
||||
port_allowlist: Option<&[u16]>,
|
||||
) -> io::Result<(String, Vec<File>)> {
|
||||
let (attribution_token, plan) = extract_attribution_token_and_plan(std::env::vars().collect());
|
||||
// SAFETY: the sandbox helper is single-threaded here, before it forks bridge workers or
|
||||
// executes the user command.
|
||||
@@ -89,6 +75,20 @@ pub(crate) fn prepare_host_proxy_route_spec() -> io::Result<(String, Vec<File>)>
|
||||
return Err(io::Error::new(io::ErrorKind::InvalidInput, message));
|
||||
}
|
||||
|
||||
// Validate the entire plan before forking any host-side workers or sending
|
||||
// attribution. An empty enforced list grants no routes.
|
||||
if let Some(ports) = port_allowlist
|
||||
&& plan
|
||||
.routes
|
||||
.iter()
|
||||
.any(|route| !ports.contains(&route.endpoint.port()))
|
||||
{
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::PermissionDenied,
|
||||
"proxy endpoint is outside the managed sandbox context",
|
||||
));
|
||||
}
|
||||
|
||||
let mut control_by_endpoint = BTreeMap::new();
|
||||
for route in &plan.routes {
|
||||
if control_by_endpoint.contains_key(&route.endpoint) {
|
||||
@@ -237,8 +237,12 @@ fn plan_proxy_routes(env: &HashMap<String, String>) -> ProxyRoutePlan {
|
||||
}
|
||||
|
||||
fn is_proxy_env_key(key: &str) -> bool {
|
||||
let upper = key.to_ascii_uppercase();
|
||||
PROXY_ENV_KEYS.contains(&upper.as_str())
|
||||
// Only exact key spellings overwritten by managed setup may nominate a host
|
||||
// socket. A caller-controlled mixed-case alias can otherwise create a second bridge.
|
||||
codex_network_proxy::PROXY_ENV_KEYS.contains(&key)
|
||||
&& codex_network_proxy::PROXY_URL_ENV_KEYS
|
||||
.iter()
|
||||
.any(|url_key| key.eq_ignore_ascii_case(url_key))
|
||||
}
|
||||
|
||||
fn parse_loopback_proxy_endpoint(proxy_url: &str) -> Option<SocketAddr> {
|
||||
@@ -352,6 +356,26 @@ fn spawn_host_bridge(
|
||||
Ok(File::from(OwnedFd::from(control)))
|
||||
}
|
||||
|
||||
/// Each connection worker holds a permit until it exits.
|
||||
pub(super) struct ConnectionSlots(Arc<Semaphore>);
|
||||
|
||||
impl ConnectionSlots {
|
||||
pub(super) fn new() -> Self {
|
||||
const MAX_CONNECTIONS: usize = 128;
|
||||
Self(Arc::new(Semaphore::new(MAX_CONNECTIONS)))
|
||||
}
|
||||
|
||||
pub(super) fn spawn(&self, serve: impl FnOnce() + Send + 'static) {
|
||||
let Ok(permit) = Arc::clone(&self.0).try_acquire_owned() else {
|
||||
return;
|
||||
};
|
||||
let _ = std::thread::Builder::new().spawn(move || {
|
||||
let _permit = permit;
|
||||
serve();
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn run_host_bridge(
|
||||
endpoint: SocketAddr,
|
||||
mut control: UnixStream,
|
||||
@@ -368,10 +392,11 @@ fn run_host_bridge(
|
||||
drop(control);
|
||||
|
||||
let attribution_token = attribution_token.map(str::to_owned);
|
||||
let slots = ConnectionSlots::new();
|
||||
loop {
|
||||
let (client_stream, _) = listener.accept()?;
|
||||
let attribution_token = attribution_token.clone();
|
||||
std::thread::spawn(move || {
|
||||
slots.spawn(move || {
|
||||
let mut tcp_stream = match TcpStream::connect(endpoint) {
|
||||
Ok(stream) => stream,
|
||||
Err(_) => return,
|
||||
@@ -468,22 +493,38 @@ fn ensure_loopback_interface_up() -> io::Result<()> {
|
||||
close_fd(fd)
|
||||
}
|
||||
|
||||
fn proxy_bidirectional(mut upstream: TcpStream, mut client: TcpStream) -> io::Result<()> {
|
||||
let mut upstream_reader = upstream.try_clone()?;
|
||||
let mut client_writer = client.try_clone()?;
|
||||
let upstream_to_client = std::thread::spawn(move || {
|
||||
let result = std::io::copy(&mut upstream_reader, &mut client_writer);
|
||||
let _ = client_writer.shutdown(Shutdown::Write);
|
||||
result
|
||||
});
|
||||
let client_to_upstream = std::io::copy(&mut client, &mut upstream);
|
||||
let _ = upstream.shutdown(Shutdown::Write);
|
||||
let upstream_to_client = upstream_to_client
|
||||
.join()
|
||||
.map_err(|_| io::Error::other("bridge thread panicked"))?;
|
||||
upstream_to_client?;
|
||||
client_to_upstream?;
|
||||
Ok(())
|
||||
pub(super) fn proxy_bidirectional(upstream: TcpStream, client: TcpStream) -> io::Result<()> {
|
||||
let (upload_done, upload) = mpsc::channel();
|
||||
let (download_done, download) = mpsc::channel();
|
||||
std::thread::scope(|scope| {
|
||||
let response = std::thread::Builder::new().spawn_scoped(scope, || {
|
||||
copy_and_shutdown(&upstream, &client, download_done, upload)
|
||||
})?;
|
||||
let request = copy_and_shutdown(&client, &upstream, upload_done, download);
|
||||
response
|
||||
.join()
|
||||
.map_err(|_| io::Error::other("bridge thread panicked"))??;
|
||||
request
|
||||
})
|
||||
}
|
||||
|
||||
fn copy_and_shutdown(
|
||||
mut source: &TcpStream,
|
||||
mut destination: &TcpStream,
|
||||
done: mpsc::Sender<()>,
|
||||
peer: mpsc::Receiver<()>,
|
||||
) -> io::Result<()> {
|
||||
let result = io::copy(&mut source, &mut destination);
|
||||
let _ = destination.shutdown(Shutdown::Write);
|
||||
let _ = done.send(());
|
||||
// Let the other direction flush queued bytes. A stalled peer must still release the permit.
|
||||
if result.is_err()
|
||||
&& peer.recv_timeout(Duration::from_secs(1)) == Err(mpsc::RecvTimeoutError::Timeout)
|
||||
{
|
||||
let _ = source.shutdown(Shutdown::Both);
|
||||
let _ = destination.shutdown(Shutdown::Both);
|
||||
}
|
||||
result.map(|_| ())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
//! Verify proxy route planning, attribution handling, and URL rewriting.
|
||||
//! Verify proxy route planning, attribution handling, URL rewriting, and bridge cleanup.
|
||||
|
||||
use super::ConnectionSlots;
|
||||
use super::PROXY_ATTRIBUTION_TOKEN_ENV_KEY;
|
||||
use super::ProxyRouteEntry;
|
||||
use super::ProxyRouteSpec;
|
||||
@@ -8,17 +9,81 @@ use super::extract_attribution_token_and_plan;
|
||||
use super::is_proxy_env_key;
|
||||
use super::parse_loopback_proxy_endpoint;
|
||||
use super::plan_proxy_routes;
|
||||
use super::proxy_bidirectional;
|
||||
use super::rewrite_proxy_env_value;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::collections::HashMap;
|
||||
use std::io::Read;
|
||||
use std::io::Write;
|
||||
use std::net::Ipv4Addr;
|
||||
use std::net::Shutdown;
|
||||
use std::net::SocketAddr;
|
||||
use std::net::TcpListener;
|
||||
use std::net::TcpStream;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_reset_on_either_side_releases_the_bridge_permit() {
|
||||
let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
|
||||
let connect = || {
|
||||
let peer = TcpStream::connect(listener.local_addr().unwrap()).unwrap();
|
||||
(peer, listener.accept().unwrap().0)
|
||||
};
|
||||
for reset_upstream in [false, true] {
|
||||
let (reset_peer, reset_socket) = connect();
|
||||
let (idle_peer, idle_socket) = connect();
|
||||
let slots = ConnectionSlots(Arc::new(Semaphore::new(1)));
|
||||
let (upstream, client) = if reset_upstream {
|
||||
(reset_socket, idle_socket)
|
||||
} else {
|
||||
(idle_socket, reset_socket)
|
||||
};
|
||||
let (done, result) = tokio::sync::oneshot::channel();
|
||||
slots.spawn(move || {
|
||||
let _ = done.send(proxy_bidirectional(upstream, client));
|
||||
});
|
||||
rustix::net::sockopt::set_socket_linger(&reset_peer, Some(Duration::ZERO)).unwrap();
|
||||
drop(reset_peer);
|
||||
|
||||
let _permit = tokio::time::timeout(Duration::from_secs(5), slots.0.acquire())
|
||||
.await
|
||||
.expect("bridge must release its permit while the other peer is still open")
|
||||
.unwrap();
|
||||
assert!(result.await.unwrap().is_err());
|
||||
drop(idle_peer);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recognizes_proxy_env_keys_case_insensitively() {
|
||||
fn upload_failure_drains_the_pending_response() {
|
||||
let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
|
||||
let mut server = TcpStream::connect(listener.local_addr().unwrap()).unwrap();
|
||||
let upstream = listener.accept().unwrap().0;
|
||||
let mut peer = TcpStream::connect(listener.local_addr().unwrap()).unwrap();
|
||||
let client = listener.accept().unwrap().0;
|
||||
|
||||
server.write_all(b"response").unwrap();
|
||||
assert!(upstream.peek(&mut [0]).unwrap() > 0);
|
||||
rustix::net::sockopt::set_socket_linger(&server, Some(Duration::ZERO)).unwrap();
|
||||
drop(server);
|
||||
peer.write_all(b"upload").unwrap();
|
||||
peer.shutdown(Shutdown::Write).unwrap();
|
||||
assert!(proxy_bidirectional(upstream, client).is_err());
|
||||
let mut response = Vec::new();
|
||||
peer.read_to_end(&mut response).unwrap();
|
||||
assert_eq!(response, b"response");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recognizes_only_proxy_url_keys_written_by_managed_setup() {
|
||||
assert_eq!(is_proxy_env_key("HTTP_PROXY"), true);
|
||||
assert_eq!(is_proxy_env_key("http_proxy"), true);
|
||||
assert_eq!(is_proxy_env_key("WS_PROXY"), true);
|
||||
assert_eq!(is_proxy_env_key("wss_proxy"), true);
|
||||
assert_eq!(is_proxy_env_key("HtTp_PrOxY"), false);
|
||||
assert_eq!(is_proxy_env_key("pip_proxy"), false);
|
||||
assert_eq!(is_proxy_env_key("PATH"), false);
|
||||
}
|
||||
|
||||
@@ -54,6 +119,10 @@ fn plan_proxy_routes_only_includes_valid_loopback_endpoints() {
|
||||
"HTTPS_PROXY".to_string(),
|
||||
"http://example.com:3128".to_string(),
|
||||
);
|
||||
env.insert(
|
||||
"HtTp_PrOxY".to_string(),
|
||||
"http://127.0.0.1:32102".to_string(),
|
||||
);
|
||||
env.insert("PATH".to_string(), "/usr/bin".to_string());
|
||||
|
||||
let plan = plan_proxy_routes(&env);
|
||||
|
||||
Reference in New Issue
Block a user