Merge pull request #5251 from Manouchehri/oidc-improvements-20240816

(oidc): Add support for loading tokens via a file, env var, and path in env var
This commit is contained in:
Krish Dholakia 2024-08-16 19:15:31 -07:00 committed by GitHub
commit 6b1be4783a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 114 additions and 0 deletions

View file

@ -19,9 +19,17 @@ LiteLLM supports the following OIDC identity providers:
| CircleCI v2 | `circleci_v2`| No | | CircleCI v2 | `circleci_v2`| No |
| GitHub Actions | `github` | Yes | | GitHub Actions | `github` | Yes |
| Azure Kubernetes Service | `azure` | No | | Azure Kubernetes Service | `azure` | No |
| File | `file` | No |
| Environment Variable | `env` | No |
| Environment Path | `env_path` | No |
If you would like to use a different OIDC provider, please open an issue on GitHub. If you would like to use a different OIDC provider, please open an issue on GitHub.
:::tip
Do not use the `file`, `env`, or `env_path` providers unless you know what you're doing, and you are sure none of the other providers will work for your use-case. Hint: they probably will.
:::
## OIDC Connect Relying Party (RP) ## OIDC Connect Relying Party (RP)
@ -46,6 +54,32 @@ For providers that do not use the `audience` parameter, you can (and should) omi
oidc/config_name_here/ oidc/config_name_here/
``` ```
#### Unofficial Providers (not recommended)
For the unofficial `file` provider, you can use the following format:
```
oidc/file/home/user/dave/this_is_a_file_with_a_token.txt
```
For the unofficial `env`, use the following format, where `SECRET_TOKEN` is the name of the environment variable that contains the token:
```
oidc/env/SECRET_TOKEN
```
For the unofficial `env_path`, use the following format, where `SECRET_TOKEN` is the name of the environment variable that contains the path to the file with the token:
```
oidc/env_path/SECRET_TOKEN
```
:::tip
If you are tempted to use oidc/env_path/AZURE_FEDERATED_TOKEN_FILE, don't do that. Instead, use `oidc/azure/`, as this will ensure continued support from LiteLLM if Azure changes their OIDC configuration and/or adds new features.
:::
## Examples ## Examples
### Google Cloud Run -> Amazon Bedrock ### Google Cloud Run -> Amazon Bedrock

View file

@ -5,6 +5,8 @@ from dotenv import load_dotenv
load_dotenv() load_dotenv()
import os import os
from uuid import uuid4
import tempfile
sys.path.insert( sys.path.insert(
0, os.path.abspath("../..") 0, os.path.abspath("../..")
@ -135,3 +137,62 @@ def test_oidc_circle_v1_with_amazon_fips():
aws_session_name="assume-v1-session-fips", aws_session_name="assume-v1-session-fips",
aws_sts_endpoint="https://sts-fips.us-west-1.amazonaws.com", aws_sts_endpoint="https://sts-fips.us-west-1.amazonaws.com",
) )
def test_oidc_env_variable():
# Create a unique environment variable name
env_var_name = "OIDC_TEST_PATH_" + uuid4().hex
os.environ[env_var_name] = "secret-" + uuid4().hex
secret_val = get_secret(
f"oidc/env/{env_var_name}"
)
print(f"secret_val: {redact_oidc_signature(secret_val)}")
assert secret_val == os.environ[env_var_name]
# now unset the environment variable
del os.environ[env_var_name]
def test_oidc_file():
# Create a temporary file
with tempfile.NamedTemporaryFile(mode='w+') as temp_file:
secret_value = "secret-" + uuid4().hex
temp_file.write(secret_value)
temp_file.flush()
temp_file_path = temp_file.name
secret_val = get_secret(
f"oidc/file/{temp_file_path}"
)
print(f"secret_val: {redact_oidc_signature(secret_val)}")
assert secret_val == secret_value
def test_oidc_env_path():
# Create a temporary file
with tempfile.NamedTemporaryFile(mode='w+') as temp_file:
secret_value = "secret-" + uuid4().hex
temp_file.write(secret_value)
temp_file.flush()
temp_file_path = temp_file.name
# Create a unique environment variable name
env_var_name = "OIDC_TEST_PATH_" + uuid4().hex
# Set the environment variable to the temporary file path
os.environ[env_var_name] = temp_file_path
# Test getting the secret using the environment variable
secret_val = get_secret(
f"oidc/env_path/{env_var_name}"
)
print(f"secret_val: {redact_oidc_signature(secret_val)}")
assert secret_val == secret_value
del os.environ[env_var_name]

View file

@ -8441,6 +8441,25 @@ def get_secret(
with open(azure_federated_token_file, "r") as f: with open(azure_federated_token_file, "r") as f:
oidc_token = f.read() oidc_token = f.read()
return oidc_token return oidc_token
elif oidc_provider == "file":
# Load token from a file
with open(oidc_aud, "r") as f:
oidc_token = f.read()
return oidc_token
elif oidc_provider == "env":
# Load token directly from an environment variable
oidc_token = os.getenv(oidc_aud)
if oidc_token is None:
raise ValueError(f"Environment variable {oidc_aud} not found")
return oidc_token
elif oidc_provider == "env_path":
# Load token from a file path specified in an environment variable
token_file_path = os.getenv(oidc_aud)
if token_file_path is None:
raise ValueError(f"Environment variable {oidc_aud} not found")
with open(token_file_path, "r") as f:
oidc_token = f.read()
return oidc_token
else: else:
raise ValueError("Unsupported OIDC provider") raise ValueError("Unsupported OIDC provider")