Team Ai
Modelpublic

Felipe97/llama-cpp-compiled

sourceHugging Faceupdated 21d agoView on Hugging Face
0likes1.2kdownloads
test_security.py246 linesDownload Raw Back to unit
1import pytest2from openai import OpenAI3from utils import *4import threading5from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer6 7server = ServerPreset.tinyllama2()8 9TEST_API_KEY = "sk-this-is-the-secret-key"10 11@pytest.fixture(autouse=True)12def create_server():13    global server14    server = ServerPreset.tinyllama2()15    server.api_key = TEST_API_KEY16 17 18@pytest.mark.parametrize("endpoint", ["/health"])19def test_access_public_endpoint(endpoint: str):20    global server21    server.start()22    res = server.make_request("GET", endpoint)23    assert res.status_code == 20024    assert "error" not in res.body25 26 27def test_access_static_assets_without_api_key():28    """Static web UI assets should not require API key authentication (issue #21229)"""29    global server30    server.start()31    for path in ["/", "/sw.js", "/manifest.webmanifest", "/_app/version.json"]:32        res = server.make_request("GET", path)33        assert res.status_code == 200, f"Expected 200 for {path}, got {res.status_code}"34 35 36@pytest.mark.parametrize("api_key", [None, "invalid-key"])37def test_incorrect_api_key(api_key: str):38    global server39    server.start()40    res = server.make_request("POST", "/completions", data={41        "prompt": "I believe the meaning of life is",42    }, headers={43        "Authorization": f"Bearer {api_key}" if api_key else None,44    })45    assert res.status_code == 40146    assert "error" in res.body47    assert res.body["error"]["type"] == "authentication_error"48 49 50def test_correct_api_key():51    global server52    server.start()53    res = server.make_request("POST", "/completions", data={54        "prompt": "I believe the meaning of life is",55    }, headers={56        "Authorization": f"Bearer {TEST_API_KEY}",57    })58    assert res.status_code == 20059    assert "error" not in res.body60    assert "content" in res.body61 62 63def test_correct_api_key_anthropic_header():64    global server65    server.start()66    res = server.make_request("POST", "/completions", data={67        "prompt": "I believe the meaning of life is",68    }, headers={69        "X-Api-Key": TEST_API_KEY,70    })71    assert res.status_code == 20072    assert "error" not in res.body73    assert "content" in res.body74 75 76def test_openai_library_correct_api_key():77    global server78    server.start()79    client = OpenAI(api_key=TEST_API_KEY, base_url=f"http://{server.server_host}:{server.server_port}")80    res = client.chat.completions.create(81        model="gpt-3.5-turbo",82        messages=[83            {"role": "system", "content": "You are a chatbot."},84            {"role": "user", "content": "What is the meaning of life?"},85        ],86    )87    assert len(res.choices) == 188 89 90@pytest.mark.parametrize("origin,cors_header,cors_header_value", [91    ("localhost", "Access-Control-Allow-Origin", "localhost"),92    ("web.mydomain.fr", "Access-Control-Allow-Origin", "web.mydomain.fr"),93    ("origin", "Access-Control-Allow-Credentials", "true"),94    ("web.mydomain.fr", "Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS"),95    ("web.mydomain.fr", "Access-Control-Allow-Headers", "*"),96])97def test_cors_options(origin: str, cors_header: str, cors_header_value: str):98    global server99    server.start()100    res = server.make_request("OPTIONS", "/completions", headers={101        "Origin": origin,102        "Access-Control-Request-Method": "POST",103        "Access-Control-Request-Headers": "Authorization",104    })105    assert res.status_code == 200106    assert cors_header in res.headers107    assert res.headers[cors_header] == cors_header_value108 109 110@pytest.mark.parametrize("origin", [111    "http://localhost",112    "http://localhost:8080",113    "http://127.0.0.1",114    "http://127.0.0.1:3000",115    "http://[::1]",116    "http://[::1]:3000",117])118def test_cors_origins_localhost_reflects(origin: str):119    global server120    server = ServerPreset.router()121    server.cors_origins = "localhost"122    server.start()123    res = server.make_request("OPTIONS", "/completions", headers={124        "Origin": origin,125        "Access-Control-Request-Method": "POST",126        "Access-Control-Request-Headers": "Authorization",127    })128    assert res.status_code == 200129    assert res.headers["Access-Control-Allow-Origin"] == origin130 131 132@pytest.mark.parametrize("origin", [133    "http://web.mydomain.fr",134    "http://evil.com",135    "http://notlocalhost",136    "http://localhost.evil.com",137])138def test_cors_origins_localhost_rejects(origin: str):139    global server140    server = ServerPreset.router()141    server.cors_origins = "localhost"142    server.start()143    res = server.make_request("OPTIONS", "/completions", headers={144        "Origin": origin,145        "Access-Control-Request-Method": "POST",146        "Access-Control-Request-Headers": "Authorization",147    })148    assert res.status_code == 200149    assert "Access-Control-Allow-Origin" not in res.headers150 151 152def test_cors_origins_defaults_to_localhost_with_tools_enabled():153    global server154    server = ServerPreset.router()155    server.server_tools = "all"156    server.start()157    res = server.make_request("OPTIONS", "/completions", headers={158        "Origin": "http://localhost:8080",159        "Access-Control-Request-Method": "POST",160        "Access-Control-Request-Headers": "Authorization",161    })162    assert res.status_code == 200163    assert res.headers["Access-Control-Allow-Origin"] == "http://localhost:8080"164 165    res = server.make_request("OPTIONS", "/completions", headers={166        "Origin": "http://evil.com",167        "Access-Control-Request-Method": "POST",168        "Access-Control-Request-Headers": "Authorization",169    })170    assert res.status_code == 200171    assert "Access-Control-Allow-Origin" not in res.headers172 173 174def test_cors_proxy_only_forwards_explicit_proxy_headers():175    class CaptureHeadersHandler(BaseHTTPRequestHandler):176        def do_GET(self):177            self.server.captured_headers = dict(self.headers)178            self.send_response(200)179            self.end_headers()180            self.wfile.write(b"ok")181 182        def log_message(self, format, *args):183            pass184 185    target = ThreadingHTTPServer(("127.0.0.1", 0), CaptureHeadersHandler)186    target.captured_headers = {}187    target_thread = threading.Thread(target=target.serve_forever, daemon=True)188    target_thread.start()189 190    try:191        server = ServerPreset.tinyllama2()192        server.api_key = TEST_API_KEY193        server.ui_mcp_proxy = True194        server.start()195 196        res = server.make_request("GET", f"/cors-proxy?url=http://127.0.0.1:{target.server_port}/capture", headers={197            "Authorization": f"Bearer {TEST_API_KEY}",198            "Proxy-Authorization": "Basic secret",199            "X-Api-Key": TEST_API_KEY,200            "Cookie": "session=secret",201            "x-llama-server-proxy-header-accept": "application/json",202            "x-llama-server-proxy-header-authorization": "Bearer explicit",203        })204 205        assert res.status_code == 200206        captured = {key.lower(): value for key, value in target.captured_headers.items()}207        assert captured["accept"] == "application/json"208        assert captured["authorization"] == "Bearer explicit"209        assert "proxy-authorization" not in captured210        assert "x-api-key" not in captured211        assert "cookie" not in captured212    finally:213        target.shutdown()214        target.server_close()215 216 217@pytest.mark.parametrize(218    "media_path, image_url, success",219    [220        (None,             "file://mtmd/test-1.jpeg",    False), # disabled media path, should fail221        ("../../../tools", "file://mtmd/test-1.jpeg",    True),222        ("../../../tools", "file:////mtmd//test-1.jpeg", True),  # should be the same file as above223        ("../../../tools", "file://mtmd/notfound.jpeg",  False), # non-existent file224        ("../../../tools", "file://../mtmd/test-1.jpeg", False), # no directory traversal225    ]226)227def test_local_media_file(media_path, image_url, success,):228    server = ServerPreset.tinygemma3()229    server.media_path = media_path230    server.start()231    res = server.make_request("POST", "/chat/completions", data={232        "max_tokens": 1,233        "messages": [234            {"role": "user", "content": [235                {"type": "text", "text": "test"},236                {"type": "image_url", "image_url": {237                    "url": image_url,238                }},239            ]},240        ],241    })242    if success:243        assert res.status_code == 200244    else:245        assert res.status_code == 400246