-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add JWT API key authentication to YOLOv4 (#65)
* Add JWT API key authentication * Add authentication header to locust testing and add API demonstration test * Add invalid API key test
- Loading branch information
Showing
8 changed files
with
208 additions
and
18 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,2 +1,3 @@ | ||
pillow==10.3.0 | ||
locust==2.25.0 | ||
pytest==7.2.2 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,6 +1,7 @@ | ||
# inference server requirements | ||
fastapi==0.85.1 | ||
uvicorn==0.19.0 | ||
pyjwt==2.7.0 | ||
python-multipart==0.0.5 | ||
|
||
-f https://download.pytorch.org/whl/cpu/torch_stable.html |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,52 @@ | ||
# SPDX-License-Identifier: Apache-2.0 | ||
# | ||
# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC | ||
|
||
#!/usr/bin/env python3 | ||
import argparse | ||
import json | ||
import os | ||
|
||
import sys | ||
|
||
import jwt | ||
|
||
|
||
def main(): | ||
parser = argparse.ArgumentParser(description="Generate signed JWT payload") | ||
parser.add_argument("mode", type=str, help="'encode' or 'decode'") | ||
parser.add_argument( | ||
"payload", type=str, help="JSON string if 'encode', token if 'decode'" | ||
) | ||
parser.add_argument( | ||
"--secret", | ||
type=str, | ||
dest="secret", | ||
help="JWT secret if not provided as environment variable JWT_SECRET", | ||
) | ||
args = parser.parse_args() | ||
|
||
try: | ||
jwt_secret = os.environ.get("JWT_SECRET", args.secret) | ||
except KeyError: | ||
print("ERROR: Expected JWT_SECRET environment variable to be provided") | ||
sys.exit(1) | ||
|
||
try: | ||
if args.mode == "encode": | ||
json_payload = json.loads(args.payload) | ||
encoded_jwt = jwt.encode(json_payload, jwt_secret, algorithm="HS256") | ||
print(encoded_jwt) | ||
elif args.mode == "decode": | ||
decoded_jwt = jwt.decode(args.payload, jwt_secret, algorithms="HS256") | ||
print(decoded_jwt) | ||
else: | ||
print("ERROR: Expected mode to be 'encode' or 'decode'") | ||
sys.exit(1) | ||
except json.decoder.JSONDecodeError: | ||
print("ERROR: Expected payload to be a valid JSON string") | ||
sys.exit(1) | ||
|
||
|
||
if __name__ == "__main__": | ||
main() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,47 @@ | ||
# SPDX-License-Identifier: Apache-2.0 | ||
# | ||
# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC | ||
|
||
from http import HTTPStatus | ||
import os | ||
import pytest | ||
import requests | ||
from utils import get_auth_header, sample_file | ||
|
||
|
||
DEPLOY_URL = "http://127.0.0.1" | ||
SERVICE_PORT = int(os.getenv("SERVICE_PORT", 7000)) | ||
API_BASE_URL = f"{DEPLOY_URL}:{SERVICE_PORT}" | ||
API_URL = f"{API_BASE_URL}/objdetection_v2" | ||
HEALTH_URL = f"{API_BASE_URL}/health" | ||
|
||
|
||
def test_valid_api_call(): | ||
# get sample image file | ||
file = sample_file() | ||
# make request with auth headers | ||
headers = get_auth_header() | ||
response = requests.post(API_URL, files=file, headers=headers) | ||
# perform status and value checking | ||
assert response.status_code == HTTPStatus.OK | ||
assert isinstance(response.json(), list) | ||
|
||
|
||
def test_invalid_api_call(): | ||
# get sample image file | ||
file = sample_file() | ||
# make request with INVALID auth header | ||
headers = get_auth_header() | ||
headers.update(Authorization="INVALID API KEY") | ||
response = requests.post(API_URL, files=file, headers=headers) | ||
# assert request was unauthorized | ||
assert response.status_code == HTTPStatus.UNAUTHORIZED | ||
|
||
|
||
@pytest.mark.skip( | ||
reason="Not implemented, see https://github.com/tenstorrent/tt-inference-server/issues/63" | ||
) | ||
def test_get_health(): | ||
headers = {} | ||
response = requests.get(HEALTH_URL, headers=headers, timeout=35) | ||
assert response.status_code == 200 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,34 @@ | ||
# SPDX-License-Identifier: Apache-2.0 | ||
# | ||
# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC | ||
|
||
import io | ||
import os | ||
from PIL import Image | ||
import requests | ||
|
||
|
||
def get_auth_header(): | ||
if authorization_header := os.getenv("AUTHORIZATION", None): | ||
headers = {"Authorization": authorization_header} | ||
return headers | ||
else: | ||
raise RuntimeError("AUTHORIZATION environment variable is undefined.") | ||
|
||
|
||
# save image as JPEG in-memory | ||
def sample_file(): | ||
# load sample image | ||
url = "http://images.cocodataset.org/val2017/000000039769.jpg" | ||
pil_image = Image.open(requests.get(url, stream=True).raw) | ||
pil_image = pil_image.resize((320, 320)) # Resize to target dimensions | ||
# convert to bytes | ||
buf = io.BytesIO() | ||
# format as JPEG | ||
pil_image.save( | ||
buf, | ||
format="JPEG", | ||
) | ||
byte_im = buf.getvalue() | ||
file = {"file": byte_im} | ||
return file |