[bugfix] Separate OCI Object Storage clients for ome agent#143
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @chengjieyao, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request implements a crucial security enhancement for the ome-agent's object replication by introducing distinct OCI Object Storage clients for read (download) and write (upload) operations. This separation of concerns, enforced by different authentication mechanisms, prevents unauthorized data modifications and lays the groundwork for more complex replication scenarios.
Highlights
- Enhanced Security & Maintainability: This PR refactors the OCI Object Storage client usage within the
ome-agent's object replication workflow to improve security and maintainability. It introduces separate clients for download and upload operations. - Separate Client Responsibilities: A dedicated 'source' client is now used for downloading objects, authenticated with customer-level credentials (OBO token), ensuring read-only access. A separate 'target' client is used for uploading objects, authenticated with service-level credentials, preventing unauthorized customer write access to the model store bucket.
- Dependency Injection Refactor: The underlying dependency injection (fx) framework has been updated to provide and manage multiple named OCI Object Storage client configurations (source and target) instead of a single generic one. New configuration options (
WithName,WithOboTokenEnabled) were added to support this. - Future Scalability: The new client structure is designed to facilitate future scalability, such as cross-region replication and fan-out scenarios, by clearly separating and identifying different object storage endpoints.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
op109lvb
requested review from
davidsnahm,
frankzhouhr,
heymrbox,
pallasathena92,
slin1237 and
speedmancs
as code owners
There was a problem hiding this comment.
Code Review
This PR refactors the object replication to use separate OCI Object Storage clients for source (download) and target (upload), enhancing security. My review focuses on improving test coverage and clarity, particularly for the new configuration and dependency injection logic.
op109lvb
force-pushed
the
separate-os-client
branch
from
722b1cd to
7943cb8
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request effectively refactors the object replication workflow to use separate Object Storage clients for source and target, which is a great improvement for security and maintainability. The changes are well-aligned with the PR's description. My review focuses on improving test specificity, error message clarity, and code robustness in a few areas to further enhance the quality of the changes.
op109lvb
force-pushed
the
separate-os-client
branch
from
7943cb8 to
6afdf95
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request separates object storage clients for download and upload, enhancing security and maintainability. The review focuses on improving configuration robustness and test precision, addressing potential duplicate client configurations and suggesting test refinements.
op109lvb
force-pushed
the
separate-os-client
branch
from
6afdf95 to
bff01e5
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces separate OCI Object Storage clients for the ome agent, enhancing security and maintainability. The changes involve refactoring the replica job to use distinct clients for downloading and uploading objects, authenticated with customer and service-level credentials, respectively. The review identified areas for improvement in tests, including test specificity, assertion robustness, and potential nil pointer dereference. Additionally, opportunities exist to reduce code duplication and correct a typo in the dependency injection module.
| LocalPath: "/test/path", | ||
| SourceObjectStoreURI: sourceURI, | ||
| TargetObjectStoreURI: targetURI, | ||
| SourceObjectStorageDataStore: mockDataStore.OCIOSDataStore, |
There was a problem hiding this comment.
The Config struct being created for the ReplicaAgent in this test is missing the TargetObjectStorageDataStore field. This field is required, and its absence will lead to a nil pointer dereference when agent.Start() calls uploadObject, causing the test to panic. Please initialize TargetObjectStorageDataStore to prevent this test failure.
| SourceObjectStorageDataStore: mockDataStore.OCIOSDataStore, | |
| SourceObjectStorageDataStore: mockDataStore, | |
| TargetObjectStorageDataStore: mockDataStore, |
| DownloadSizeLimitGB: 100, | ||
| EnableSizeLimitCheck: true, | ||
| NumConnections: 5, | ||
| TargetObjectStoreURI: ociobjectstore.ObjectURI{Namespace: "tgt-ns", BucketName: "tgt-bucket"}, |
There was a problem hiding this comment.
This test case, named "missing LocalPath", is also missing the SourceObjectStoreURI field, which is also a required field. While the test will still pass because an error is expected, it makes the test less specific as it's testing for more than one missing field simultaneously. To improve test clarity and focus, please ensure each "missing field" test case is only missing the single field it's named after.
| assert.Empty(t, params.ObjectStorageDataStores) | ||
| assert.Equal(t, mockLogger, params.AnotherLogger) | ||
| assert.Equal(t, mockDataStore, params.ObjectStorageDataStores) | ||
| assert.Empty(t, mockDataStores) |
There was a problem hiding this comment.
These assertions are brittle because they depend on the order of elements in the mockDataStores slice. The production logic in WithAppParams assigns datastores based on their Name property, not their position in the slice. If the order of elements in mockDataStores were to change, this test would fail even if the production code is correct.
To make the test more robust, please assert on the properties of the datastores (like their Config.Name) rather than relying on their index in the mock slice.
| assert.Empty(t, params.ObjectStorageDataStores) | |
| assert.Equal(t, mockLogger, params.AnotherLogger) | |
| assert.Equal(t, mockDataStore, params.ObjectStorageDataStores) | |
| assert.Empty(t, mockDataStores) | |
| assert.Equal(t, ociobjectstore.SourceOsConfigName, params.ObjectStorageDataStores[0].Config.Name) | |
| assert.Equal(t, ociobjectstore.TargetOsConfigName, params.ObjectStorageDataStores[1].Config.Name) |
| // ProvideSourceOCIOSDataStoreConfig ProvideSourceOCIOSDataStore constructs a Config instance for the source OCI Object Storage | ||
| // location using Viper configuration, environment context, and a logger. | ||
| // This function is intended to be used as an fx provider. The resulting Config is wrapped in | ||
| // OSDataStoreConfigWrapper and added to the "oSDataStoreConfig" value group for collective injection. |
There was a problem hiding this comment.
There's a small typo in this comment. The value group is referred to as "oSDataStoreConfig", but the struct tag in OSDataStoreConfigWrapper (and the consuming appParams) uses "casperConfigs". To avoid confusion, please correct the comment to match the implementation.
| // OSDataStoreConfigWrapper and added to the "oSDataStoreConfig" value group for collective injection. | |
| // OSDataStoreConfigWrapper and added to the "casperConfigs" value group for collective injection. |
| func ProvideSourceOCIOSDataStoreConfig(v *viper.Viper, logger logging.Interface) (OSDataStoreConfigWrapper, error) { | ||
| config, err := NewConfig(WithViper(v), WithAnotherLog(logger), WithName(SourceOsConfigName)) | ||
| if err != nil { | ||
| return OSDataStoreConfigWrapper{}, fmt.Errorf("error creating source object store config: %w", err) | ||
| } | ||
| if err = config.Validate(); err != nil { | ||
| return OSDataStoreConfigWrapper{}, fmt.Errorf("ociobjectstore config is invalid: %w", err) | ||
| } | ||
| return OSDataStoreConfigWrapper{ | ||
| OSDataStoreConfig: config, | ||
| }, nil | ||
| } | ||
|
|
||
| // ProvideTargetOCIOSDataStoreConfig constructs a Config instance for the destination (target) OCI Object Storage | ||
| // location using Viper configuration, environment context, and a logger. This function is intended to be | ||
| // used as an fx provider. The resulting Config is wrapped in OSDataStoreConfigWrapper and added to the | ||
| // "oSDataStoreConfig" value group for collective injection. | ||
| // | ||
| // Note: Destination object storage locations are currently expected to reside under service tenancies | ||
| // where customer OBO tokens are not permitted. | ||
| // | ||
| // TODO: Refactor to separate source and target configuration in OME agent config, and read specifically | ||
| // from the target config. | ||
| // | ||
| // Returns: | ||
| // - OSDataStoreConfigWrapper containing the initialized Config if successful | ||
| // - An error if configuration loading or initialization fails | ||
| func ProvideTargetOCIOSDataStoreConfig(v *viper.Viper, logger logging.Interface) (OSDataStoreConfigWrapper, error) { | ||
| config, err := NewConfig(WithViper(v), WithAnotherLog(logger), WithName(TargetOsConfigName), WithOboTokenEnabled(false)) | ||
| if err != nil { | ||
| return OSDataStoreConfigWrapper{}, fmt.Errorf("error creating target object store config: %w", err) | ||
| } | ||
| if err = config.Validate(); err != nil { | ||
| return OSDataStoreConfigWrapper{}, fmt.Errorf("ociobjectstore config is invalid: %w", err) | ||
| } | ||
| return OSDataStoreConfigWrapper{ | ||
| OSDataStoreConfig: config, | ||
| }, nil | ||
| } |
There was a problem hiding this comment.
The functions ProvideSourceOCIOSDataStoreConfig and ProvideTargetOCIOSDataStoreConfig contain very similar logic for creating and validating a Config object. This duplication can be reduced by extracting the common logic into a helper function. This would improve maintainability by centralizing the config creation and validation logic, making future changes easier.
Separate Object Storage clients Co-authored-by: Chengjie Yao <cheyao@cj-mac-m3.local>
2 participants
What type of PR is this?
/kind bug
What this PR does / why we need it:
This PR updates the object replication workflow to use separate Object Storage clients for downloading and uploading objects, improving both security and maintainability.
Background:
We want to ensure that customers do not have direct write access to the customer model store bucket. Granting write permissions would allow users operating with OBO (On-Behalf-Of) tokens to directly overwrite or modify data in the model store, which introduces significant security and data integrity risks. By clearly separating client responsibilities - one for downloading with customer-level credentials and one for uploading with service-level credentials - we prevent unauthorized or accidental data modifications by customers. This approach also provides stronger isolation of permissions and aligns with the principle of least privilege.
Changes:
Which issue(s) this PR fixes:
Fixes #
Does this PR introduce a user-facing change?