wgtmac commented on code in PR #892:
URL: https://github.com/apache/iceberg-cpp/pull/892#discussion_r4272755728
##########
src/iceberg/arrow/s3/arrow_s3_file_io.cc:
##########
@@ -314,6 +325,252 @@ std::string CanonicalizeS3Scheme(std::string_view
location) {
return std::string(location);
}
+// Lead time before expiry, matching Java's VendedCredentialsProvider.
+constexpr auto kRefreshLeadTime = std::chrono::minutes(5);
+
+// How long to wait after a failed refresh before asking the catalog again.
+constexpr auto kRefreshRetryBackoff = std::chrono::seconds(30);
+
+// Floor on a backoff shortened to land on the expiry.
+constexpr auto kMinRefreshRetryBackoff = std::chrono::seconds(1);
+
+// How long a caller with expired credentials waits for a refresh under way.
+constexpr auto kExpiredCredentialWait = std::chrono::seconds(10);
+
+// Whether `config` carries an access key pair.
+bool HasKeys(const std::unordered_map<std::string, std::string>& config) {
+ return FindProperty(config, S3Properties::kAccessKeyId) != nullptr &&
+ FindProperty(config, S3Properties::kSecretAccessKey) != nullptr;
+}
+
+// The properties an S3 client for `credential` is built from.
+std::unordered_map<std::string, std::string> PropertiesFor(
+ const std::unordered_map<std::string, std::string>& default_properties,
+ const StorageCredential& credential) {
+ auto properties = default_properties;
+ for (const auto& [key, value] : credential.config) {
+ properties[key] = value;
+ }
+ return properties;
+}
+
+// The client settings in `properties`: everything but the keys.
+std::unordered_map<std::string, std::string> SettingsOf(
+ std::unordered_map<std::string, std::string> properties) {
+ for (const auto key :
+ {S3Properties::kAccessKeyId, S3Properties::kSecretAccessKey,
+ S3Properties::kSessionToken, S3Properties::kSessionTokenExpiresAtMs}) {
+ properties.erase(std::string(key));
+ }
+ return properties;
+}
+
+// When the session token that `properties` sign with expires: nullopt if they
+// sign without one, or an error if the expiry is missing or invalid.
+Result<std::optional<std::chrono::system_clock::time_point>> ExpiryOf(
+ const std::unordered_map<std::string, std::string>& properties,
+ std::string_view prefix) {
+ // Only a nonempty token alongside an access key pair signs requests.
+ const auto* token = FindProperty(properties, S3Properties::kSessionToken);
+ if (!HasKeys(properties) || token == nullptr || token->empty()) {
+ return std::nullopt;
+ }
+ const auto* value = FindProperty(properties,
S3Properties::kSessionTokenExpiresAtMs);
+ if (value == nullptr) {
+ return InvalidArgument("Credential \"{}\" has a session token but no
\"{}\"", prefix,
+ S3Properties::kSessionTokenExpiresAtMs);
+ }
+ auto millis = StringUtils::ParseNumber<int64_t>(*value);
+ if (!millis.has_value()) {
+ return InvalidArgument(
+ "Credential \"{}\" has a session token but an unparseable \"{}\" value
\"{}\"",
+ prefix, S3Properties::kSessionTokenExpiresAtMs, *value);
+ }
+ // Beyond what the clock can hold, converting would overflow it.
+ constexpr auto kMaxMillis =
std::chrono::duration_cast<std::chrono::milliseconds>(
+ std::chrono::system_clock::duration::max())
+ .count();
+ constexpr auto kMinMillis =
std::chrono::duration_cast<std::chrono::milliseconds>(
+ std::chrono::system_clock::duration::min())
+ .count();
+ if (*millis > kMaxMillis || *millis < kMinMillis) {
+ return InvalidArgument("Credential \"{}\" has an out-of-range \"{}\" value
\"{}\"",
+ prefix, S3Properties::kSessionTokenExpiresAtMs,
*value);
+ }
+ return
std::chrono::system_clock::time_point(std::chrono::milliseconds(*millis));
+}
+
+// The S3 credentials in effect: the first for each canonical prefix, which is
+// the one its locations match.
+std::vector<const StorageCredential*> S3CredentialsOf(
+ const std::vector<StorageCredential>& credentials) {
+ std::vector<const StorageCredential*> s3_credentials;
+ std::unordered_set<std::string> prefixes;
+ for (const auto& credential : credentials) {
+ if (IsS3CredentialPrefix(credential.prefix) &&
+ prefixes.insert(CanonicalizeS3Scheme(credential.prefix)).second) {
+ s3_credentials.push_back(&credential);
+ }
+ }
+ return s3_credentials;
+}
+
+// When installed credentials need replacing.
+struct CredentialsExpiry {
+ // The earliest expiry known, if any.
+ std::optional<std::chrono::system_clock::time_point> expires_at;
+ // A session token lacks a usable expiry: refresh at once, as Java does, but
+ // do not treat it as expired.
+ bool unknown = false;
+};
+
+CredentialsExpiry ExpiryOf(
+ const std::unordered_map<std::string, std::string>& default_properties,
+ const std::vector<StorageCredential>& credentials) {
+ CredentialsExpiry expiry;
+ for (const auto* credential : S3CredentialsOf(credentials)) {
+ auto expires_at =
+ ExpiryOf(PropertiesFor(default_properties, *credential),
credential->prefix);
+ if (!expires_at.has_value()) {
+ ICEBERG_LOG_WARN("{}", expires_at.error().message);
+ expiry.unknown = true;
+ } else if (expires_at->has_value() &&
+ (!expiry.expires_at.has_value() || **expires_at <
*expiry.expires_at)) {
+ expiry.expires_at = **expires_at;
+ }
+ }
+ return expiry;
+}
+
+// Known expiries by canonical prefix.
+std::unordered_map<std::string, std::chrono::system_clock::time_point>
ExpiriesOf(
+ const std::unordered_map<std::string, std::string>& default_properties,
+ const std::vector<StorageCredential>& credentials) {
+ std::unordered_map<std::string, std::chrono::system_clock::time_point>
expiries;
+ for (const auto* credential : S3CredentialsOf(credentials)) {
+ // Invalid expiries are logged on install.
+ if (auto expires_at =
+ ExpiryOf(PropertiesFor(default_properties, *credential),
credential->prefix);
+ expires_at.has_value() && expires_at->has_value()) {
+ expiries.emplace(CanonicalizeS3Scheme(credential->prefix), **expires_at);
+ }
+ }
+ return expiries;
+}
+
+// The S3 prefixes `credentials` serve, each with its client settings and
+// whether it signs with keys.
+using Layout =
+ std::unordered_map<std::string,
+ std::pair<std::unordered_map<std::string, std::string>,
bool>>;
+
+Layout LayoutOf(const std::unordered_map<std::string, std::string>&
default_properties,
+ const std::vector<StorageCredential>& credentials) {
+ Layout layout;
+ for (const auto* credential : S3CredentialsOf(credentials)) {
+ auto properties = PropertiesFor(default_properties, *credential);
+ const bool keys = HasKeys(properties);
+ layout.try_emplace(CanonicalizeS3Scheme(credential->prefix),
+ SettingsOf(std::move(properties)), keys);
+ }
+ return layout;
+}
+
+// A refresh may only rotate keys, so the current clients serve it: it must
keep
+// the layout of `current`, and its tokens must be usable.
+Status ValidateRefreshed(
+ const std::unordered_map<std::string, std::string>& default_properties,
+ const std::vector<StorageCredential>& current,
+ const std::vector<StorageCredential>& credentials) {
+ for (const auto& credential : credentials) {
+ ICEBERG_RETURN_UNEXPECTED(credential.Validate());
+ }
+ if (LayoutOf(default_properties, credentials) !=
+ LayoutOf(default_properties, current)) {
+ return InvalidArgument(
+ "Refreshed credentials change the S3 prefixes, their settings or
keys");
+ }
+ const auto now = std::chrono::system_clock::now();
+ for (const auto* credential : S3CredentialsOf(credentials)) {
+ ICEBERG_ASSIGN_OR_RAISE(
+ auto expires_at,
+ ExpiryOf(PropertiesFor(default_properties, *credential),
credential->prefix));
+ if (expires_at.has_value() && *expires_at <= now) {
+ return InvalidArgument("Credential \"{}\" has already expired",
credential->prefix);
+ }
+ }
+ return {};
+}
+
+// Requires HasKeys(config).
+Aws::Auth::AWSCredentials AwsCredentialsOf(
+ const std::unordered_map<std::string, std::string>& config) {
+ Aws::Auth::AWSCredentials credentials;
+ credentials.SetAWSAccessKeyId(
+ FindProperty(config, S3Properties::kAccessKeyId)->c_str());
+ credentials.SetAWSSecretKey(
+ FindProperty(config, S3Properties::kSecretAccessKey)->c_str());
+ if (const auto* session_token = FindProperty(config,
S3Properties::kSessionToken);
+ session_token != nullptr) {
+ credentials.SetSessionToken(session_token->c_str());
+ }
+ return credentials;
+}
+
+class ArrowS3FileIO;
+
+// Gives S3 clients their FileIO's current keys. Clients may outlive the
+// FileIO, which detaches on destruction.
+class KeySource {
+ public:
+ explicit KeySource(ArrowS3FileIO* io) : io_(io) {}
+
+ void Detach() {
+ std::unique_lock lock(mutex_);
+ io_ = nullptr;
+ }
+
+ /// \brief Refreshes if due, then returns the FileIO's keys for `prefix` and
+ /// `settings`; nullopt if it has none or is gone.
+ std::optional<Aws::Auth::AWSCredentials> KeysFor(
+ std::string_view prefix,
+ const std::unordered_map<std::string, std::string>& settings);
+
+ private:
+ std::shared_mutex mutex_;
+ ArrowS3FileIO* io_;
+};
+
+// Supplies an S3 client's keys. The AWS SDK asks before each request, so open
+// handles pick up refreshed keys, as in Java. Keys vended with other settings
+// (e.g. another endpoint) are never used; the last keys stay instead.
+class VendedCredentialsProvider final : public
Aws::Auth::AWSCredentialsProvider {
+ public:
+ VendedCredentialsProvider(std::shared_ptr<KeySource> source, std::string
prefix,
+ std::unordered_map<std::string, std::string>
settings,
+ Aws::Auth::AWSCredentials credentials)
+ : source_(std::move(source)),
+ prefix_(std::move(prefix)),
+ settings_(std::move(settings)),
+ credentials_(std::move(credentials)) {}
+
+ Aws::Auth::AWSCredentials GetAWSCredentials() override {
+ auto current = source_->KeysFor(prefix_, settings_);
+ std::lock_guard lock(mutex_);
+ if (current.has_value()) {
+ credentials_ = std::move(current).value();
+ }
Review Comment:
The AWS signer calls this provider while the new client is being built,
before SetStorageCredentials installs the new keys. This replaces the initial
keys with the old ones. If the FileIO is released before the handle's first
request, that handle keeps the old keys; can we prevent this cache rollback?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
