Enhance encryption security by BenWilson2 · Pull Request #19253 · mlflow/mlflow

BenWilson2 · 2025-12-06T02:25:21Z

🛠 DevTools 🛠

Install mlflow from this PR

# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/19253/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/19253/merge#subdirectory=libs/skinny

For Databricks, use the following command:

%sh curl -LsSf https://raw.githubusercontent.com/mlflow/mlflow/HEAD/dev/install-skinny.sh | sh -s pull/19253/merge

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Changes a few things that I noticed while red-teaming the implementation:

Force the KEK version to be part of the salt. This ensures that if an admin reuses a passphrase in N >=2 rotations, the KEK will be unique even in this case, protecting against an attack vector of a leaked passphrase that wasn't graveyarded.
Convert the test-verification nonce arg to be private with a clear warnings about how dangerous it is to not use a random nonce. We DO need this argument in order to test the encryption / decryption stack, but it does NOT need to be public.
Add NB comments wherever someone might change the secret name in the backend. It is CRITICAL that this name never be made mutable (as it would break decryption). I felt that it was important to protect against foot-gunning in the future.

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

No. You can skip the rest of this section.
Yes. Give a description of this change to be included in the release notes for MLflow users.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

How should the PR be classified in the release notes? Choose one:

rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
rn/feature - A new user-facing feature worth mentioning in the release notes
rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?

Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
Bug fixes, doc updates and new features usually go into minor releases.
Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
Bug fixes and doc updates usually go into patch releases.

Yes (this PR will be cherry-picked and included in the next patch release)
No (this PR will be included in the next minor release)

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

BenWilson2 · 2025-12-06T02:27:03Z

@TomeHirata @B-Step62 I was running through some attack scenarios this evening and I managed to find a few ways that security could be degraded / this feature could break if someone didn't read up on these algorithms.
I think it's pretty important to make these changes to harden the security implementation.

Copilot

Pull request overview

This PR enhances the security of MLflow's encryption implementation through three key defensive measures identified during red-team testing:

Forces KEK version into the PBKDF2 salt to ensure unique KEKs even if an admin accidentally reuses a passphrase across rotations, protecting against leaked passphrase attack vectors
Converts the test-verification nonce parameter to private (_nonce) with keyword-only access and prominent warnings to prevent dangerous misuse in production code while preserving testability
Adds comprehensive NB (nota bene) comments documenting the critical immutability requirements of secret_id and secret_name fields used in AAD (Additional Authenticated Data) for AES-GCM encryption

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
`mlflow/utils/crypto.py`	Implements versioned salt by appending KEK version as 4 big-endian bytes; refactors `encrypt_with_aes_gcm` to private `_encrypt_with_aes_gcm` with `_nonce` as keyword-only parameter; adds detailed security documentation
`tests/utils/test_crypto.py`	Adds test verifying same passphrase with different versions produces different KEKs; updates all function calls to use `_encrypt_with_aes_gcm` with keyword argument `_nonce`
`mlflow/store/tracking/dbmodels/models.py`	Documents immutability constraints for `secret_id` and `secret_name` database columns with cross-references to AAD implementation
`mlflow/entities/gateway_secrets.py`	Documents immutability requirements for `secret_id` and `secret_name` entity fields with explanation of AAD binding

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

github-actions · 2025-12-06T02:32:54Z

Documentation preview for bc10b82 is available at:

https://pr-19253--mlflow-docs-preview.netlify.app/docs/latest/

More info

Ignore this comment if this PR does not change the documentation.
The preview is updated when a new commit is pushed to this PR.
This comment was created by this workflow run.
The documentation was built by this workflow run.

serena-ruan · 2025-12-08T05:28:55Z

mlflow/utils/crypto.py

-def encrypt_with_aes_gcm(
-    plaintext: bytes, key: bytes, nonce: bytes | None = None, aad: bytes | None = None
+def _encrypt_with_aes_gcm(
+    plaintext: bytes, key: bytes, *, aad: bytes | None = None, _nonce: bytes | None = None


Why do we use _nonce instead of nonce? This is already a private function

Even though it's end user private, this interface is callable internally by developers. I wanted to make this particularly field extra clear to devs that this is somewhat dangerous to provide a value for outside of testing (as it degrades the security of encryption to a large degree if a static value is passed in). It's intentionally meant to look abnormal.

serena-ruan · 2025-12-08T05:30:24Z

mlflow/utils/crypto.py

        aad: Optional Additional Authenticated Data. If provided, this data is
             authenticated but not encrypted. Useful for binding encryption to
             metadata (e.g., secret_id + secret_name) to prevent substitution attacks.
+        _nonce: TESTING ONLY. 12-byte nonce. If None (default), generates random nonce.


Is this testing only?

I'll change the argument name here to make it even more clear :)

serena-ruan · 2025-12-08T05:33:36Z

mlflow/entities/gateway_secrets.py

+        secret_id: Unique identifier for this secret. IMMUTABLE - used in AAD for encryption.
+        secret_name: User-friendly name for the secret. IMMUTABLE - used in AAD for encryption.


I think this docstring makes no difference, we need to update code to make it really immutable

The DB API layer already enforces immutability. But, what it doesn't do is actually block updates to this field via a DB trigger that enforces this. In order to prevent any future manipulation, I'll add that constraint.
Thanks for the perspective!

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

serena-ruan · 2025-12-11T01:51:50Z

tests/store/tracking/test_gateway_sql_store.py

    assert all(s.provider == "openai" for s in openai_secrets)


+def test_secret_id_and_name_are_immutable_at_database_level(store: SqlAlchemyStore):


Could we run this test file in database job

mlflow/.github/workflows/master.yml

Lines 132 to 135 in 312edd5

./tests/db/compose.sh run --rm --no-TTY $service pytest \

tests/store/tracking/test_sqlalchemy_store.py \

tests/store/model_registry/test_sqlalchemy_store.py \

tests/db

?

Great idea!! Updated CI config :D

serena-ruan · 2025-12-11T02:03:50Z

mlflow/entities/gateway_secrets.py



-@dataclass
+@dataclass(frozen=True)


This means last_updated_at can't be updated as well, if it's never modified then it's fine

Yep - totally fine, we return a new instance when the API is called, but on the update route for a key, the DB merges the record and updates the last_update_time.

TomeHirata · 2025-12-11T05:13:44Z

mlflow/utils/crypto.py

        nonce = os.urandom(GCM_NONCE_LENGTH)
-    elif len(nonce) != GCM_NONCE_LENGTH:
-        raise ValueError(f"Nonce must be {GCM_NONCE_LENGTH} bytes (96 bits), got {len(nonce)}")
+    elif len(_nonce_for_testing) != GCM_NONCE_LENGTH:


nit: do we still need this validation if this argument is test only?

Thanks for the catch! This is totally not needed :)

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

serena-ruan

LGTM! Do we need to add v3.8.0 label?

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

harden encryption further

7e4719e

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

Copilot AI review requested due to automatic review settings December 6, 2025 02:25

github-actions bot added area/tracking Tracking service, tracking client APIs, autologging rn/feature Mention under Features in Changelogs. labels Dec 6, 2025

BenWilson2 added the team-review Trigger a team review request label Dec 6, 2025

Copilot started reviewing on behalf of BenWilson2 December 6, 2025 02:25 View session

github-actions bot requested review from B-Step62, TomeHirata, WeichenXu123, daniellok-db, harupy, kevin-lyn, serena-ruan and xq-yin December 6, 2025 02:25

Copilot AI reviewed Dec 6, 2025

View reviewed changes

serena-ruan reviewed Dec 8, 2025

View reviewed changes

feedback

787bea9

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

BenWilson2 requested a review from serena-ruan December 8, 2025 16:22

serena-ruan reviewed Dec 11, 2025

View reviewed changes

github-actions bot assigned serena-ruan Dec 11, 2025

TomeHirata reviewed Dec 11, 2025

View reviewed changes

github-actions bot assigned TomeHirata Dec 11, 2025

BenWilson2 added 3 commits December 11, 2025 16:36

feedback

ca73732

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

Merge branch 'master' into kek-reuse-hardening

4b9dcfd

merge fix

aa606b3

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

serena-ruan approved these changes Dec 12, 2025

View reviewed changes

BenWilson2 added the v3.8.0 label Dec 12, 2025

Merge branch 'master' into kek-reuse-hardening

bc10b82

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

BenWilson2 added this pull request to the merge queue Dec 12, 2025

Merged via the queue into mlflow:master with commit b024559 Dec 12, 2025
47 checks passed

BenWilson2 deleted the kek-reuse-hardening branch December 12, 2025 22:02

		secret_id: Unique identifier for this secret. IMMUTABLE - used in AAD for encryption.
		secret_name: User-friendly name for the secret. IMMUTABLE - used in AAD for encryption.

		assert all(s.provider == "openai" for s in openai_secrets)


		def test_secret_id_and_name_are_immutable_at_database_level(store: SqlAlchemyStore):

	./tests/db/compose.sh run --rm --no-TTY $service pytest \
	tests/store/tracking/test_sqlalchemy_store.py \
	tests/store/model_registry/test_sqlalchemy_store.py \
	tests/db



		@dataclass
		@dataclass(frozen=True)

Conversation

BenWilson2 commented Dec 6, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Install mlflow from this PR

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

How should the PR be classified in the release notes? Choose one:

Should this PR be included in the next patch release?

Uh oh!

BenWilson2 commented Dec 6, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

github-actions bot commented Dec 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

serena-ruan Dec 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

serena-ruan Dec 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

serena-ruan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

BenWilson2 commented Dec 6, 2025 •

edited by github-actions bot

Loading

github-actions bot commented Dec 6, 2025 •

edited

Loading

serena-ruan Dec 8, 2025 •

edited

Loading

serena-ruan Dec 11, 2025 •

edited

Loading