yuxh1996 commited on
Commit
37495c1
Β·
1 Parent(s): 98727e9

Fix MCP API key bypass vulnerability

Browse files

CRITICAL SECURITY FIX:
- Add strict MCP request path detection (/mcp/, /gradio_api/mcp)
- Enforce mandatory API key for all MCP endpoint requests
- Prevent bypass via --transport sse-only parameter
- Add detailed request debugging and path analysis
- Improve error messages with clear configuration examples

Changes:
- Detect MCP requests by URL path analysis
- Block Space API key usage for MCP endpoints
- Add comprehensive request logging for debugging
- Ensure MCP clients cannot access without proper API key

This prevents the security issue where MCP clients could
access the service without API keys using sse-only transport.

Files changed (2) hide show
  1. app.py +42 -4
  2. test_no_key.py +43 -0
app.py CHANGED
@@ -17,6 +17,7 @@ def get_api_client():
17
  # Try to get API key from multiple sources
18
  api_key = None
19
  user_agent = ""
 
20
 
21
  # 1. Try from request headers (for MCP clients)
22
  try:
@@ -25,7 +26,10 @@ def get_api_client():
25
  headers = dict(request.headers)
26
  api_key = get_api_key_from_headers(headers)
27
  user_agent = headers.get('user-agent', '')
 
 
28
  print(f"πŸ” Request headers found - User-Agent: {user_agent}")
 
29
  print(
30
  f"πŸ” API key from headers: {'Found' if api_key else 'Not found'}")
31
  except Exception as e:
@@ -39,6 +43,13 @@ def get_api_client():
39
 
40
  # 3. Determine if this is a web browser request or MCP client request
41
  is_web_request = False
 
 
 
 
 
 
 
42
  if user_agent:
43
  user_agent_lower = user_agent.lower()
44
  # Web browsers typically have 'mozilla' in user agent
@@ -52,15 +63,42 @@ def get_api_client():
52
  is_web_request = False
53
  print("πŸ” No User-Agent found - assuming MCP client request")
54
 
55
- # 4. Use Space API key ONLY for web browser requests on Hugging Face Space
56
- if not api_key and is_space and space_api_key and is_web_request:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
57
  print("πŸ“‘ Using API key from Space environment variable (web demo)")
58
  return A1DAPIClient(space_api_key)
59
 
60
- # 5. For MCP clients or when no Space API key available, user API key is mandatory
61
  if not api_key:
62
  error_msg = (
63
- "πŸ”‘ API key is required for MCP clients!\n\n"
64
  "Please provide API_KEY in request headers.\n"
65
  "Get your API key at https://a1d.ai\n\n"
66
  "Configuration example:\n"
 
17
  # Try to get API key from multiple sources
18
  api_key = None
19
  user_agent = ""
20
+ request_path = ""
21
 
22
  # 1. Try from request headers (for MCP clients)
23
  try:
 
26
  headers = dict(request.headers)
27
  api_key = get_api_key_from_headers(headers)
28
  user_agent = headers.get('user-agent', '')
29
+ request_path = getattr(request, 'url', {}).path if hasattr(
30
+ request, 'url') else ""
31
  print(f"πŸ” Request headers found - User-Agent: {user_agent}")
32
+ print(f"πŸ” Request path: {request_path}")
33
  print(
34
  f"πŸ” API key from headers: {'Found' if api_key else 'Not found'}")
35
  except Exception as e:
 
43
 
44
  # 3. Determine if this is a web browser request or MCP client request
45
  is_web_request = False
46
+ is_mcp_request = False
47
+
48
+ # Check if this is an MCP request
49
+ if request_path and ('/mcp/' in request_path or '/gradio_api/mcp' in request_path):
50
+ is_mcp_request = True
51
+ print("πŸ” Detected MCP API request")
52
+
53
  if user_agent:
54
  user_agent_lower = user_agent.lower()
55
  # Web browsers typically have 'mozilla' in user agent
 
63
  is_web_request = False
64
  print("πŸ” No User-Agent found - assuming MCP client request")
65
 
66
+ # 4. STRICT RULE: MCP requests MUST have API key
67
+ if is_mcp_request and not api_key:
68
+ error_msg = (
69
+ "πŸ”‘ API key is REQUIRED for MCP requests!\n\n"
70
+ "This is an MCP API endpoint. You must provide your API key.\n"
71
+ "Get your API key at https://a1d.ai\n\n"
72
+ "Configuration example:\n"
73
+ '{\n'
74
+ ' "mcpServers": {\n'
75
+ ' "a1d": {\n'
76
+ ' "command": "npx",\n'
77
+ ' "args": [\n'
78
+ ' "mcp-remote@latest",\n'
79
+ ' "https://aigchacker-a1d-mcp-server.hf.space/gradio_api/mcp/sse",\n'
80
+ ' "--header",\n'
81
+ ' "API_KEY:${MCP_API_KEY}"\n'
82
+ ' ],\n'
83
+ ' "env": {\n'
84
+ ' "MCP_API_KEY": "your_a1d_api_key_here"\n'
85
+ ' }\n'
86
+ ' }\n'
87
+ ' }\n'
88
+ '}'
89
+ )
90
+ print(f"❌ MCP API key validation failed: {error_msg}")
91
+ raise ValueError(error_msg)
92
+
93
+ # 5. Use Space API key ONLY for web browser requests on Hugging Face Space
94
+ if not api_key and is_space and space_api_key and is_web_request and not is_mcp_request:
95
  print("πŸ“‘ Using API key from Space environment variable (web demo)")
96
  return A1DAPIClient(space_api_key)
97
 
98
+ # 6. For all other cases, user API key is mandatory
99
  if not api_key:
100
  error_msg = (
101
+ "πŸ”‘ API key is required!\n\n"
102
  "Please provide API_KEY in request headers.\n"
103
  "Get your API key at https://a1d.ai\n\n"
104
  "Configuration example:\n"
test_no_key.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Test script to verify API key enforcement without environment variables
4
+ """
5
+
6
+ import os
7
+ import sys
8
+
9
+ # Remove any existing API key environment variable
10
+ if 'A1D_API_KEY' in os.environ:
11
+ del os.environ['A1D_API_KEY']
12
+
13
+ # Import after removing environment variable
14
+ from app import remove_bg_wrapper
15
+
16
+ def test_no_api_key():
17
+ """Test that API key is required when not provided"""
18
+ try:
19
+ result = remove_bg_wrapper("https://example.com/test.jpg")
20
+ print(f"❌ FAILED: Function should have failed but returned: {result}")
21
+ return False
22
+ except Exception as e:
23
+ print(f"βœ… SUCCESS: Function correctly failed with error: {str(e)}")
24
+ return True
25
+
26
+ if __name__ == "__main__":
27
+ print("πŸ§ͺ Testing API key enforcement...")
28
+ print("=" * 50)
29
+
30
+ # Check environment
31
+ print(f"A1D_API_KEY in environment: {'A1D_API_KEY' in os.environ}")
32
+ print(f"SPACE_ID in environment: {'SPACE_ID' in os.environ}")
33
+
34
+ # Run test
35
+ success = test_no_api_key()
36
+
37
+ print("=" * 50)
38
+ if success:
39
+ print("βœ… Test PASSED: API key enforcement is working")
40
+ sys.exit(0)
41
+ else:
42
+ print("❌ Test FAILED: API key enforcement is NOT working")
43
+ sys.exit(1)