diff --git a/lib/proxy/control.rs b/lib/proxy/control.rs index c535c1a..56173a5 100644 --- a/lib/proxy/control.rs +++ b/lib/proxy/control.rs @@ -208,6 +208,14 @@ impl Control { return Ok(true); } + if !self.process_input(rules)? { + return Ok(false); + } + + if !self.output.is_empty() { + return Ok(true); + } + if self.input_closed { return Ok(false); } @@ -224,7 +232,13 @@ impl Control { Ok(n) => { bytes_read += n; self.input.extend_from_slice(&buf[..n]); - self.process_input(rules)?; + if !self.process_input(rules)? { + return Ok(false); + } + + if !self.output.is_empty() { + return Ok(true); + } } Err(err) if err.kind() == ErrorKind::WouldBlock => break, Err(err) @@ -239,11 +253,7 @@ impl Control { } } - if !self.flush()? { - return Ok(false); - } - - Ok(!self.input_closed || !self.output.is_empty()) + Ok(!self.input_closed) } pub(super) fn shutdown(&self) -> Result<()> { @@ -252,7 +262,7 @@ impl Control { .context("failed to shut down the control socket") } - fn process_input(&mut self, rules: &mut PrefixMap) -> Result<()> { + fn process_input(&mut self, rules: &mut PrefixMap) -> Result { loop { if self.discarding_input { if let Some(newline) = self.input.iter().position(|byte| *byte == b'\n') { @@ -262,7 +272,7 @@ impl Control { } self.input.clear(); - return Ok(()); + return Ok(true); } let Some(newline) = self.input.iter().position(|byte| *byte == b'\n') else { @@ -274,9 +284,11 @@ impl Control { PARSE_ERROR, "request exceeds the maximum frame size", ))?; + + return self.flush(); } - return Ok(()); + return Ok(true); }; let line = self.input.drain(..=newline).collect::>(); @@ -287,6 +299,15 @@ impl Control { PARSE_ERROR, "request exceeds the maximum frame size", ))?; + + if !self.flush()? { + return Ok(false); + } + + if !self.output.is_empty() { + return Ok(true); + } + continue; } @@ -296,6 +317,14 @@ impl Control { if let Some(update) = update { *rules = self.policy.apply(update); } + + if !self.flush()? { + return Ok(false); + } + + if !self.output.is_empty() { + return Ok(true); + } } } @@ -908,7 +937,10 @@ mod tests { #[test] fn oversized_frame_is_discarded_and_following_frame_is_processed() { - let (_client, server) = UnixStream::pair().unwrap(); + let (mut client, server) = UnixStream::pair().unwrap(); + client + .set_read_timeout(Some(Duration::from_secs(1))) + .unwrap(); let mut control = control(server.as_raw_fd()).unwrap(); let mut rules = PrefixMap::new(); @@ -922,7 +954,9 @@ mod tests { control.process_input(&mut rules).unwrap(); assert!(!control.discarding_input); - let responses = std::str::from_utf8(&control.output) + let mut output = [0; 2048]; + let n = client.read(&mut output).unwrap(); + let responses = std::str::from_utf8(&output[..n]) .unwrap() .lines() .map(|line| serde_json::from_str::(line).unwrap()) @@ -988,6 +1022,38 @@ mod tests { assert!(control.output.len() - control.output_offset <= 4 * MAX_REQUEST_BYTES); } + #[test] + fn pipelined_policy_responses_stop_before_queue_overflow() { + let (_client, server) = UnixStream::pair().unwrap(); + let mut control = control(server.as_raw_fd()).unwrap(); + let mut rules = PrefixMap::new(); + let allow = (0..MAX_TARGETS) + .map(|index| format!("10.{}.{}.0/24", index / 256, index % 256)) + .collect::>(); + let mut input = serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "id": 0, + "method": "softnet.policy.set", + "params": {"allow": allow, "block": []} + })) + .unwrap(); + input.push(b'\n'); + + for id in 1..=100 { + input.extend_from_slice( + format!("{{\"jsonrpc\":\"2.0\",\"id\":{id},\"method\":\"softnet.policy.get\"}}\n") + .as_bytes(), + ); + } + + control.input = input; + assert!(control.process_input(&mut rules).unwrap()); + assert_eq!(rules.len(), MAX_TARGETS); + assert!(!control.input.is_empty()); + assert!(!control.output.is_empty()); + assert!(control.output.len() - control.output_offset <= MAX_PENDING_RESPONSE_BYTES); + } + #[test] fn response_queue_overflow_does_not_apply_a_policy_update() { let (_client, server) = UnixStream::pair().unwrap();