From 1d778b487d384514db0ddb7c829e2605dc1fe042 Mon Sep 17 00:00:00 2001
From: yylt <yang8518296@163.com>
Date: Mon, 12 Jan 2026 22:12:21 +0800
Subject: [PATCH] update cookie when html

---
 Makefile         |  3 +++
 src/proxy/api.rs |  5 ++---
 src/proxy/mod.rs | 43 +++++++++++--------------------------------
 3 files changed, 16 insertions(+), 35 deletions(-)

diff --git a/Makefile b/Makefile
index dff2745..48f959d 100644
--- a/Makefile
+++ b/Makefile
@@ -7,3 +7,6 @@ deploy: ## deploy to cf workers
 .PHONY: dev
 dev: ## run the project locally
 	@ npx wrangler dev -c .wrangler.dev.toml
+
+lint:
+	cargo clippy --all-targets --all-features -- -D warnings
\ No newline at end of file
diff --git a/src/proxy/api.rs b/src/proxy/api.rs
index ab62751..d5248ce 100644
--- a/src/proxy/api.rs
+++ b/src/proxy/api.rs
@@ -6,7 +6,6 @@ use regex::Regex;
 
 static REGISTRY: &str = "registry-1.docker.io";
 
-
 fn replace_host(content: &mut str, src: &str, dest: &str) -> Result<String> {
 
     let re = Regex::new(r#"(?P<attr>src|href)(?P<eq>=)(?P<quote>['"]?)(?P<url>(//|https://))"#)
@@ -112,7 +111,6 @@ pub async fn handler(mut req: Request, uri: Url, dst_host: &str) -> Result<Respo
                 }         
             }
             (401, "www-authenticate") => value.replace("https://", &format!("https://{}/", my_host)),
-            (_, "set-cookie") => value.replace(dst_host, &my_host),
             _ => value,
         };
         resp_header.set(&key, &new_value)?;
@@ -121,9 +119,10 @@ pub async fn handler(mut req: Request, uri: Url, dst_host: &str) -> Result<Respo
     let _ = resp_header.set("access-control-allow-origin", "*");
     if let Some(s) = resp_header.get("content-type")? {
         if s.contains("text/html")  {
+            let _ = resp_header.delete("content-encoding");
+            let _ = resp_header.set("set-cookie", format!("{}={}; Path=/; Max-Age=3600", COOKIE_HOST_KEY, dst_host).as_str());
             let mut body = response.text().await?;
             let newbody = replace_host(&mut body, dst_host, &my_host)?;
-            let _ = resp_header.delete("content-encoding");
             let resp = Response::builder()
                 .with_headers(resp_header)
                 .with_status(status)
diff --git a/src/proxy/mod.rs b/src/proxy/mod.rs
index ea5961d..661d655 100644
--- a/src/proxy/mod.rs
+++ b/src/proxy/mod.rs
@@ -204,31 +204,21 @@ pub async fn handler(req: Request, cx: RouteContext<()>) -> Result<Response> {
                 .and_then(|cookie| get_cookie_by_name(&cookie, COOKIE_HOST_KEY));
 
             let (mut domain, port, mut path) = parse_path(&origin_path);
-            // when not resolve, will try find domain by cookie.
-            let mut notresolve= true;
-            // when only domain, will update cookie in response.
-            let mut onlydomain = false;
             let scheme = "https";
-            
-            match domain {
-                Some(d) if d.contains('.') => {
-                    if (dns::is_cf_address(&Address::Domain(d)).await).is_ok() {
-                        notresolve = false;
-                        if path.is_none() || path.as_ref().unwrap().len()<2 {
-                            onlydomain = true;
-                        }
-                    }
-                },
-                _ => {},
-            }
 
-            match (notresolve, &cookie_host) {
-                (true, Some(host)) => {
+            // when not resolve, will try find domain by cookie.
+            let resolve = match domain {
+                Some(d) => d.contains('.') && dns::is_cf_address(&Address::Domain(d)).await.is_ok(),
+                _ => false,
+            };
+
+            match (resolve, &cookie_host) {
+                (false, Some(host)) => {
                     domain = Some(host.as_ref());  
                     path = Some(origin_path.as_str());
                 }
-                (true, None) => return Response::error("Not Found", 404),
-                (false, _) => {},
+                (false, None) => return Response::error("Not Found", 404),
+                (true, _) => {},
             }
 
             let host = domain.unwrap();
@@ -250,18 +240,7 @@ pub async fn handler(req: Request, cx: RouteContext<()>) -> Result<Response> {
                     .join("&")
                     .as_str());
             }
-            let mut resp = api::handler(req,  Url::parse(&url)?, host).await?;
-            match resp.headers().get("content-type")? {
-                Some(s) if s.contains("text/html") => {
-                    if onlydomain {
-                        console_debug!("set cookie domain: {:?}", host);
-                        let _ = resp.headers_mut()
-                            .set("set-cookie", format!("{}={}; Path=/; Max-Age=3600", COOKIE_HOST_KEY, host).as_str());
-                    }
-                }
-                _ => {}
-            }
-            Ok(resp)    
+            api::handler(req,  Url::parse(&url)?, host).await   
         }
     }   
 }