Code cleanup

Author: lukehinds

src/codegate/providers/copilot/mapping.py

@@ -10,6 +10,7 @@ class CopilotProxyRoute(BaseModel):
 class CoPilotMappings(BaseSettings):
 
     # Proxy routes configuration
+    # This code is to ensure that incoming proxy requests are routed to the correct target
     PROXY_ROUTES: List[tuple[str, str]] = [
         ("/github/login", "https://github.com/login"),
         ("/api/github/user", "https://api.github.com"),
@@ -30,56 +31,6 @@ class CoPilotMappings(BaseSettings):
         ("/v1", "https://copilot-proxy.githubusercontent.com/v1"),
     ]
 
-    # Headers configuration
-    PRESERVED_HEADERS: List[str] = [
-        'authorization',
-        'user-agent',
-        'content-type',
-        'accept',
-        'accept-encoding',
-        'connection',
-        'x-github-token',
-        'github-token',
-        'x-request-id',
-        'x-github-api-version',
-        'openai-organization',
-        'openai-intent',
-        'openai-model',
-        'editor-version',
-        'editor-plugin-version',
-        'vscode-sessionid',
-        'vscode-machineid',
-    ]
-
-    REMOVED_HEADERS: List[str] = [
-        'proxy-connection',
-        'proxy-authenticate',
-        'proxy-authorization',
-        'connection',
-        'keep-alive',
-        'transfer-encoding',
-        'te',
-        'trailer',
-        'proxy-authenticate',
-        'upgrade',
-        'expect',
-    ]
-
-    ENDPOINT_HEADERS: Dict[str, Dict[str, str]] = {
-        '/v1/engines/copilot-codex/completions': {
-            'Accept': 'application/json',
-            'Content-Type': 'application/json',
-            'Editor-Version': 'vscode/1.95.3',
-            'Editor-Plugin-Version': 'copilot/1.246.0',
-            'Openai-Organization': 'github-copilot',
-            'Openai-Intent': 'copilot-ghost',
-            'User-Agent': 'GithubCopilot/1.246.0',
-            'Accept-Encoding': 'gzip, deflate, br',
-            'X-Github-Api-Version': '2022-11-28',
-            'Host': 'copilot-proxy.githubusercontent.com'
-        }
-    }
-
 # Create settings instance
 mappings = CoPilotMappings()
 

src/codegate/providers/copilot/provider.py

@@ -8,7 +8,7 @@
 
 from codegate.config import Config
 from codegate.ca.codegate_ca import CertificateAuthority
-from codegate.providers.copilot.mapping import VALIDATED_ROUTES 
+from codegate.providers.copilot.mapping import VALIDATED_ROUTES
 
 logger = structlog.get_logger("codegate")
 
@@ -59,7 +59,7 @@ def extract_path(self, full_path: str) -> str:
         elif full_path.startswith('/'):
             return full_path.lstrip('/')
         return full_path
-    
+
     def get_headers(self) -> Dict[str, str]:
         """Get request headers as a dictionary"""
         logger.debug("Getting headers as dictionary fn: get_headers")
@@ -78,7 +78,7 @@ def get_headers(self) -> Dict[str, str]:
                     headers_dict[name.strip().lower()] = value.strip()
                 except ValueError:
                     continue
-                    
+
             return headers_dict
         except Exception as e:
             logger.error(f"Error getting headers: {e}")
@@ -222,37 +222,58 @@ async def handle_http_request(self):
             logger.error(f"Error handling HTTP request: {e}")
             self.send_error_response(502, str(e).encode())
 
-    def data_received(self, data: bytes):
+    def _check_buffer_size(self, new_data: bytes) -> bool:
+        """Check if adding new data would exceed the maximum buffer size"""
+        return len(self.buffer) + len(new_data) <= MAX_BUFFER_SIZE
+
+    def _handle_parsed_headers(self) -> None:
+        """Handle the request based on parsed headers"""
+        if self.method == 'CONNECT':
+            logger.debug("Handling CONNECT request")
+            self.handle_connect()
+        else:
+            logger.debug("Handling HTTP request")
+            asyncio.create_task(self.handle_http_request())
+
+    def _forward_data_to_target(self, data: bytes) -> None:
+        """Forward data to target if connection is established"""
+        if self.target_transport and not self.target_transport.is_closing():
+            self.log_decrypted_data(data, "Client to Server")
+            self.target_transport.write(data)
+
+    def data_received(self, data: bytes) -> None:
+        """Handle received data from the client"""
         logger.debug(f"Data received from {self.peername} fn: data_received")
 
         try:
-            if len(self.buffer) + len(data) > MAX_BUFFER_SIZE:
-                logger.error("Request too large")
+            # Check buffer size limit
+            if not self._check_buffer_size(data):
+                logger.error("Request exceeds maximum buffer size")
                 self.send_error_response(413, b"Request body too large")
                 return
 
+            # Append new data to buffer
             self.buffer.extend(data)
 
             if not self.headers_parsed:
+                # Try to parse headers
                 self.headers_parsed = self.parse_headers()
                 if not self.headers_parsed:
                     return
 
-                if self.method == 'CONNECT':
-                    logger.debug("Handling CONNECT request")
-                    self.handle_connect()
-                else:
-                    logger.debug("Handling HTTP request")
-                    asyncio.create_task(self.handle_http_request())
-            elif self.target_transport and not self.target_transport.is_closing():
-                self.log_decrypted_data(data, "Client to Server")
-                self.target_transport.write(data)
+                # Handle the request based on parsed headers
+                self._handle_parsed_headers()
+            else:
+                # Forward data to target if headers are already parsed
+                self._forward_data_to_target(data)
 
+        except asyncio.CancelledError:
+            logger.warning("Operation cancelled")
+            raise
         except Exception as e:
-            logger.error(f"Error in data_received: {e}")
+            logger.error(f"Error processing received data: {e}")
             self.send_error_response(502, str(e).encode())
 
- 
     def handle_connect(self):
         '''
         This where requests are sent directly via the tunnel created during
@@ -294,7 +315,7 @@ def handle_connect(self):
         except Exception as e:
             logger.error(f"Error handling CONNECT: {e}")
             self.send_error_response(502, str(e).encode())
-    
+
     def send_error_response(self, status: int, message: bytes):
         logger.debug(f"Sending error response: {status} {message} fn: send_error_response")
         response = (