From 0675cfc45120747ec91f9def4d5fe9e1a241f069 Mon Sep 17 00:00:00 2001
From: zonyitoo <zonyitoo@gmail.com>
Date: Thu, 20 Jun 2024 23:00:30 +0800
Subject: [PATCH] feat(local): PingBalancer check firefox portal allow 200
 status (#1560)

---
 Cargo.lock                                    |  1 +
 crates/shadowsocks-service/Cargo.toml         |  3 +-
 crates/shadowsocks-service/src/config.rs      |  9 ++-
 .../src/local/loadbalancing/ping_balancer.rs  | 58 ++++++++++---------
 4 files changed, 41 insertions(+), 30 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index af7ba9ca..4a11c129 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3273,6 +3273,7 @@ dependencies = [
  "hickory-resolver",
  "http 1.1.0",
  "http-body-util",
+ "httparse",
  "hyper",
  "idna 1.0.0",
  "ipnet",
diff --git a/crates/shadowsocks-service/Cargo.toml b/crates/shadowsocks-service/Cargo.toml
index a24ee93e..77d94e7c 100644
--- a/crates/shadowsocks-service/Cargo.toml
+++ b/crates/shadowsocks-service/Cargo.toml
@@ -27,7 +27,7 @@ full = [
 ]
 
 # Enable local server
-local = []
+local = ["httparse"]
 # Enable remote server
 server = []
 # Enable manager server
@@ -170,6 +170,7 @@ libc = "0.2.141"
 hyper = { version = "1.3", optional = true, features = ["full"] }
 http-body-util = { version = "0.1", optional = true }
 http = { version = "1.1", optional = true }
+httparse = { version = "1.9", optional = true }
 
 hickory-resolver = { version = "0.24", optional = true, features = [
     "serde-config",
diff --git a/crates/shadowsocks-service/src/config.rs b/crates/shadowsocks-service/src/config.rs
index d4440647..9f37ad3c 100644
--- a/crates/shadowsocks-service/src/config.rs
+++ b/crates/shadowsocks-service/src/config.rs
@@ -70,7 +70,14 @@ use serde::{Deserialize, Serialize};
 use shadowsocks::relay::socks5::Address;
 use shadowsocks::{
     config::{
-        ManagerAddr, Mode, ReplayAttackPolicy, ServerAddr, ServerConfig, ServerSource, ServerUser, ServerUserManager,
+        ManagerAddr,
+        Mode,
+        ReplayAttackPolicy,
+        ServerAddr,
+        ServerConfig,
+        ServerSource,
+        ServerUser,
+        ServerUserManager,
         ServerWeight,
     },
     crypto::CipherKind,
diff --git a/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs b/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs
index 0c715f69..33f3fbe6 100644
--- a/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs
+++ b/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs
@@ -841,6 +841,8 @@ impl PingChecker {
     /// Detect TCP connectivity with Chromium [Network Portal Detection](https://www.chromium.org/chromium-os/chromiumos-design-docs/network-portal-detection)
     #[allow(dead_code)]
     async fn check_request_tcp_chromium(&self) -> io::Result<()> {
+        use std::io::{Error, ErrorKind};
+
         static GET_BODY: &[u8] =
             b"GET /generate_204 HTTP/1.1\r\nHost: clients3.google.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";
 
@@ -860,27 +862,28 @@ impl PingChecker {
         let mut buf = Vec::new();
         reader.read_until(b'\n', &mut buf).await?;
 
-        static EXPECTED_HTTP_STATUS_LINE: &[u8] = b"HTTP/1.1 204 No Content\r\n";
-        if buf != EXPECTED_HTTP_STATUS_LINE {
-            use std::io::{Error, ErrorKind};
+        let mut headers = [httparse::EMPTY_HEADER; 1];
+        let mut response = httparse::Response::new(&mut headers);
 
-            debug!(
-                "unexpected response from http://clients3.google.com/generate_204, {:?}",
-                ByteStr::new(&buf)
-            );
-
-            let err = Error::new(
-                ErrorKind::InvalidData,
-                "unexpected response from http://clients3.google.com/generate_204",
-            );
-            return Err(err);
+        if let Ok(..) = response.parse(&buf) {
+            if matches!(response.code, Some(204)) {
+                return Ok(());
+            }
         }
 
-        Ok(())
+        Err(Error::new(
+            ErrorKind::InvalidData,
+            format!(
+                "unexpected response from http://clients3.google.com/generate_204, {:?}",
+                ByteStr::new(&buf)
+            ),
+        ))
     }
 
     /// Detect TCP connectivity with Firefox's http://detectportal.firefox.com/success.txt
     async fn check_request_tcp_firefox(&self) -> io::Result<()> {
+        use std::io::{Error, ErrorKind};
+
         static GET_BODY: &[u8] =
             b"GET /success.txt HTTP/1.1\r\nHost: detectportal.firefox.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";
 
@@ -900,23 +903,22 @@ impl PingChecker {
         let mut buf = Vec::new();
         reader.read_until(b'\n', &mut buf).await?;
 
-        static EXPECTED_HTTP_STATUS_LINE: &[u8] = b"HTTP/1.1 200 OK\r\n";
-        if buf != EXPECTED_HTTP_STATUS_LINE {
-            use std::io::{Error, ErrorKind};
+        let mut headers = [httparse::EMPTY_HEADER; 1];
+        let mut response = httparse::Response::new(&mut headers);
 
-            debug!(
-                "unexpected response from http://detectportal.firefox.com/success.txt, {:?}",
-                ByteStr::new(&buf)
-            );
-
-            let err = Error::new(
-                ErrorKind::InvalidData,
-                "unexpected response from http://detectportal.firefox.com/success.txt",
-            );
-            return Err(err);
+        if let Ok(..) = response.parse(&buf) {
+            if matches!(response.code, Some(200) | Some(204)) {
+                return Ok(());
+            }
         }
 
-        Ok(())
+        Err(Error::new(
+            ErrorKind::InvalidData,
+            format!(
+                "unexpected response from http://detectportal.firefox.com/success.txt, {:?}",
+                ByteStr::new(&buf)
+            ),
+        ))
     }
 
     async fn check_request_udp(&self) -> io::Result<()> {