Skip to content

Added bucket-ownership checks for Amazon S3#18542

Merged
serena-ruan merged 8 commits intomlflow:masterfrom
kingroryg:bucket-check
Nov 10, 2025
Merged

Added bucket-ownership checks for Amazon S3#18542
serena-ruan merged 8 commits intomlflow:masterfrom
kingroryg:bucket-check

Conversation

@kingroryg
Copy link
Contributor

@kingroryg kingroryg commented Oct 27, 2025

🛠 DevTools 🛠

Open in GitHub Codespaces

Install mlflow from this PR

# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/18542/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/18542/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/18542/merge

Related Issues/PRs

Fix #18541

What changes are proposed in this pull request?

  1. Environment Variable (mlflow/environment_variables.py)
    Added MLFLOW_S3_BUCKET_OWNER environment variable to specify the expected AWS account ID that owns the S3 bucket. This is optional and backward compatible - when not set, the system behaves as before.
  2. S3 Artifact Repository (mlflow/store/artifact/s3_artifact_repo.py)
    • Added _bucket_owner attribute initialized from the environment variable
    • Updated all S3 API calls to include ExpectedBucketOwner parameter when configured:
      • upload_file (single file uploads)
      • download_file (file downloads)
      • list_objects_v2 (listing artifacts)
      • delete_objects (deleting artifacts)
      • create_multipart_upload (multipart upload initiation)
      • complete_multipart_upload (multipart upload completion)
      • abort_multipart_upload (multipart upload abortion)
      • generate_presigned_url (presigned URL generation)
  3. Optimized S3 Artifact Repository (mlflow/store/artifact/optimized_s3_artifact_repo.py)
    • Applied the same bucket ownership checks to all S3 operations
    • Updated head_bucket, upload_file, download_file, list_objects_v2, delete_objects, multipart upload operations, and presigned URL generation
  4. Tests (tests/store/artifact/test_s3_artifact_repo.py)
    • Bucket ownership verification with environment variable set
    • Backward compatibility when environment variable is not set
    • Bucket takeover scenario simulation
    • List artifacts with bucket owner
    • Multipart upload with bucket owner
    • Delete artifacts with bucket owner

How is this PR tested?

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

Does this PR require documentation update?

  • [] No. You can skip the rest of this section.
  • Yes. I've updated:
    • Examples
    • API references
    • Instructions

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

  • area/tracking: Tracking Service, tracking client APIs, autologging
  • area/models: MLmodel format, model serialization/deserialization, flavors
  • area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
  • area/scoring: MLflow Model server, model deployment tools, Spark UDFs
  • area/evaluation: MLflow model evaluation features, evaluation metrics, and evaluation workflows
  • area/gateway: MLflow AI Gateway client APIs, server, and third-party integrations
  • area/prompts: MLflow prompt engineering features, prompt templates, and prompt management
  • area/tracing: MLflow Tracing features, tracing APIs, and LLM tracing functionality
  • area/projects: MLproject format, project running backends
  • area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
  • area/build: Build and test infrastructure for MLflow
  • area/docs: MLflow documentation pages

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)
Screenshot 2025-10-27 at 3 32 52 PM

@github-actions
Copy link
Contributor

@kingroryg Thank you for the contribution! Could you fix the following issue(s)?

⚠ DCO check

The DCO check failed. Please sign off your commit(s) by following the instructions here. See https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.md#sign-your-work for more details.

@github-actions github-actions bot added v3.5.2 area/build Build and test infrastructure for MLflow rn/bug-fix Mention under Bug Fixes in Changelogs. labels Oct 27, 2025
@kingroryg kingroryg force-pushed the bucket-check branch 2 times, most recently from eb360f5 to c5a9ba6 Compare October 27, 2025 22:46
Copy link
Collaborator

@serena-ruan serena-ruan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for this fix, looks solid! Could you address comments?

@kingroryg
Copy link
Contributor Author

@serena-ruan The requested changes have been pushed. Thanks for the prompt review!

@kingroryg
Copy link
Contributor Author

@serena-ruan Pushed the requested changes. Apologies for being remiss about the imports. Thanks.

Copy link
Collaborator

@serena-ruan serena-ruan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Thanks for the contribution!! Let's fix the doc format (https://github.com/mlflow/mlflow/actions/runs/18916054315/job/54040568892?pr=18542) and we're good to go!

@kingroryg
Copy link
Contributor Author

Done, thx! @serena-ruan

@harupy
Copy link
Member

harupy commented Oct 31, 2025

/review

@github-actions
Copy link
Contributor

github-actions bot commented Oct 31, 2025

Documentation preview for 5c5b2bb is available at:

Changed Pages (1)
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.

@kingroryg
Copy link
Contributor Author

Thanks @harupy. I've pushed all the changes except this: #18542 (comment)

Signed-off-by: Sarthak Munshi <smunshii@amazon.com>
harupy and others added 4 commits November 4, 2025 18:47
Replace _add_bucket_owner_if_present with _get_bucket_owner_params to return
bucket owner parameters as a dict instead of mutating kwargs in place. This
improves code clarity by making the return value explicit and eliminates
side effects. Also adds type hints and uses modern Python patterns including
the walrus operator for cleaner conditional assignments.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
Replace _get_bucket_owner_params() method with _bucket_owner_params instance
variable computed once during initialization. This eliminates redundant dict
creation on every S3 operation and simplifies the code by removing the method
and the separate _bucket_owner variable.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
Update test assertions to check _bucket_owner_params dict instead of
the removed _bucket_owner attribute. Tests now verify the dict structure
matches expected values.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
Replace unnecessary file writing with file_path.touch() in bucket owner
tests. The tests only need files to exist, not specific content.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
@harupy harupy requested a review from Copilot November 4, 2025 11:00
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR adds S3 bucket ownership verification to protect against bucket takeover attacks. It introduces a new environment variable MLFLOW_S3_EXPECTED_BUCKET_OWNER that, when set, includes the ExpectedBucketOwner parameter in all S3 API calls to verify that the bucket is owned by the expected AWS account.

Key changes:

  • Added MLFLOW_S3_EXPECTED_BUCKET_OWNER environment variable
  • Updated S3 artifact repository classes to include bucket owner verification in all S3 operations
  • Added comprehensive test coverage for the bucket ownership verification feature

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
mlflow/environment_variables.py Defines the new MLFLOW_S3_EXPECTED_BUCKET_OWNER environment variable
mlflow/store/artifact/s3_artifact_repo.py Implements bucket owner verification in S3ArtifactRepository by adding _bucket_owner_params to all S3 API calls
mlflow/store/artifact/optimized_s3_artifact_repo.py Implements bucket owner verification in OptimizedS3ArtifactRepository by adding _bucket_owner_params to all S3 API calls
tests/store/artifact/test_s3_artifact_repo.py Adds comprehensive tests for bucket ownership verification including upload, download, list, delete, and multipart upload operations
docs/docs/self-hosting/architecture/artifact-store.mdx Documents the new bucket ownership verification feature with usage examples and security context

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

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
@serena-ruan serena-ruan added this pull request to the merge queue Nov 10, 2025
Merged via the queue into mlflow:master with commit 9a3591e Nov 10, 2025
47 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/build Build and test infrastructure for MLflow rn/bug-fix Mention under Bug Fixes in Changelogs. v3.5.2

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] MLflow S3 Integration prone to bucket sniping

4 participants