Update proxy_handler.py

proxy_handler.py

@@ -22,7 +22,6 @@ class ProxyHandler:
             limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
             http2=True,
         )
-        # The primary secret key from the reference code.
         self.primary_secret = "junjie".encode('utf-8')
 
     async def aclose(self):
@@ -33,24 +32,21 @@ class ProxyHandler:
         return int(time.time() * 1000)
 
     def _parse_jwt_token(self, token: str) -> Dict[str, str]:
-        """A simple JWT payload decoder to get user ID ('sub' claim)."""
         try:
             parts = token.split('.')
             if len(parts) != 3: return {"user_id": ""}
             payload_b64 = parts[1]
-            payload_b64 += '=' * (-len(payload_b64) % 4) # Add padding if needed
+            payload_b64 += '=' * (-len(payload_b64) % 4)
             payload_json = base64.urlsafe_b64decode(payload_b64).decode('utf-8')
             payload = json.loads(payload_json)
             return {"user_id": payload.get("sub", "")}
         except Exception:
-            # It's okay if this fails; we'll proceed with an empty user_id.
             return {"user_id": ""}
 
-    # MODIFICATION START: The function now accepts the timestamp to ensure consistency.
     def _generate_signature(self, e_payload: str, t_payload: str, timestamp_ms: int) -> Dict[str, Any]:
         """
         Generates the signature based on the logic from the reference JS code.
-        A consistent timestamp is now passed into this function.
+        This version corrects the HMAC chaining issue by using .digest() for the intermediate key.
 
         Args:
             e_payload (str): The simplified payload string (e.g., "requestId,...,timestamp,...").
@@ -61,20 +57,23 @@ class ProxyHandler:
             A dictionary with 'signature' and 'timestamp'.
         """
         b64_encoded_t = base64.b64encode(t_payload.encode("utf-8")).decode("utf-8")
-        
-        # Use the passed-in timestamp instead of generating a new one.
         message_string = f"{e_payload}|{b64_encoded_t}|{timestamp_ms}"
-        
         n = timestamp_ms // (5 * 60 * 1000) 
         
+        # --- MODIFICATION START: Correct HMAC Chaining ---
+        
+        # 1. First HMAC: Calculate intermediate key as RAW BYTES using .digest()
         msg1 = str(n).encode("utf-8")
-        intermediate_key = hmac.new(self.primary_secret, msg1, hashlib.sha256).hexdigest()
+        intermediate_key_bytes = hmac.new(self.primary_secret, msg1, hashlib.sha256).digest()
 
+        # 2. Second HMAC: Use the raw bytes of the intermediate key directly.
+        #    The final result is converted to a hex string for the header.
         msg2 = message_string.encode("utf-8")
-        final_signature = hmac.new(intermediate_key.encode("utf-8"), msg2, hashlib.sha256).hexdigest()
+        final_signature = hmac.new(intermediate_key_bytes, msg2, hashlib.sha256).hexdigest()
+        
+        # --- MODIFICATION END ---
 
         return {"signature": final_signature, "timestamp": timestamp_ms}
-    # MODIFICATION END
 
     def _clean_thinking_content(self, text: str) -> str:
         if not text: return ""
@@ -88,7 +87,7 @@ class ProxyHandler:
     def _clean_answer_content(self, text: str) -> str:
         if not text: return ""
         cleaned_text = re.sub(r'<details[^>]*>.*?</details>', '', text, flags=re.DOTALL)
-        cleaned_text = re.sub(r'<glm_block.*?</glm_block>|<summary>.*?</summary>', '', cleaned_text, flags=re.DOTALL)
+        cleaned_text = re.sub(r'<glm_block.*?</glm_block>|<summary>.*?</summary>', '', text, flags=re.DOTALL)
         cleaned_text = re.sub(r'\s*duration="\d+"[^>]*>', '', cleaned_text)
         return cleaned_text
 
@@ -102,36 +101,28 @@ class ProxyHandler:
         return out
 
     async def _prep_upstream(self, req: ChatCompletionRequest) -> Tuple[Dict[str, Any], Dict[str, str], str, str]:
-        """Prepares the request body, headers, cookie, and URL for the upstream API."""
         ck = await cookie_manager.get_next_cookie()
         if not ck: raise HTTPException(503, "No available cookies")
 
         model = settings.UPSTREAM_MODEL if req.model == settings.MODEL_NAME else req.model
         
-        # MODIFICATION START: Generate timestamp ONCE and reuse it everywhere.
         timestamp_ms = self._get_timestamp_millis()
         payload_timestamp = str(timestamp_ms)
-        # MODIFICATION END
 
         payload_user_id = str(uuid.uuid4())
         payload_request_id = str(uuid.uuid4())
         
-        # e: The simplified payload for the signature, now using the single consistent timestamp.
         e_payload = f"requestId,{payload_request_id},timestamp,{payload_timestamp},user_id,{payload_user_id}"
 
-        # t: The last message content
         t_payload = ""
         if req.messages:
             last_message = req.messages[-1]
             if isinstance(last_message.content, str):
                 t_payload = last_message.content
 
-        # MODIFICATION START: Pass the consistent timestamp into the signature generator.
         signature_data = self._generate_signature(e_payload, t_payload, timestamp_ms)
-        # MODIFICATION END
         
         signature = signature_data["signature"]
-        # This timestamp now correctly matches the one in the URL params.
         signature_timestamp = signature_data["timestamp"]
         
         url_params = {
@@ -288,7 +279,6 @@ class ProxyHandler:
             logger.exception("Non-stream processing failed"); raise
 
     async def handle_chat_completion(self, req: ChatCompletionRequest):
-        """Determines whether to stream or not and handles the request."""
         stream = bool(req.stream) if req.stream is not None else settings.DEFAULT_STREAM
         if stream:
             return StreamingResponse(self.stream_proxy_response(req), media_type="text/event-stream",