be/src/runtime/aws_msk_iam_auth.h

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#pragma once

#include <aws/core/auth/AWSCredentialsProvider.h>
#include <librdkafka/rdkafkacpp.h>

#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>

#include "common/status.h"

namespace doris {
/**
 * AWS MSK IAM authentication token generator.
 * 
 * This class generates SASL/OAUTHBEARER tokens for AWS MSK IAM authentication.
 * It uses AWS SDK for C++ to obtain IAM credentials and generates signed tokens
 * that can be used with librdkafka's OAUTHBEARER mechanism.
 */
class AwsMskIamAuth {
public:
    struct Config {
        std::string region;       // AWS region (e.g., "us-east-1"), required
        std::string access_key;   // AWS Access Key ID (optional, for explicit credentials)
        std::string secret_key;   // AWS Secret Access Key (optional, for explicit credentials)
        std::string role_arn;     // IAM role ARN (optional, for assume role)
        std::string external_id;  // Optional external ID for STS AssumeRole
        std::string profile_name; // AWS profile name (optional, reads from ~/.aws/credentials)
        std::string
                credentials_provider; // Credentials provider type (optional, e.g., "ENV", "INSTANCE_PROFILE")
        int token_refresh_margin_ms = 60000; // Refresh token 60s before expiry
    };

    explicit AwsMskIamAuth(Config config);
    ~AwsMskIamAuth() = default;

    /**
     * Generate AWS MSK IAM authentication token.
     * 
     * The token is a base64url-encoded presigned URL following AWS SigV4 format:
     * https://kafka.<region>.amazonaws.com/?Action=kafka-cluster:Connect
     *   &X-Amz-Algorithm=AWS4-HMAC-SHA256
     *   &X-Amz-Credential=<access-key>/<date>/<region>/kafka-cluster/aws4_request
     *   &X-Amz-Date=<timestamp>
     *   &X-Amz-Expires=900
     *   &X-Amz-SignedHeaders=host
     *   &X-Amz-Signature=<signature>
     *   &X-Amz-Security-Token=<session-token>  // if using temporary credentials
     *   &User-Agent=doris-msk-iam-auth/1.0
     * 
     * Reference: https://github.com/aws/aws-msk-iam-sasl-signer-python
     * 
     * @param broker_hostname The MSK broker hostname (used for logging, not in token)
     * @param token Output: base64url-encoded signed URL token
     * @param token_lifetime_ms Output: token lifetime in milliseconds (3600000ms = 1 hour)
     * @return Status indicating success or failure
     */
    Status generate_token(const std::string& broker_hostname, std::string* token,
                          int64_t* token_lifetime_ms);

    /**
     * Get current AWS credentials.
     * This will refresh credentials if they are expired or about to expire.
     */
    Status get_credentials(Aws::Auth::AWSCredentials* credentials);

private:
    // Create AWS credentials provider based on configuration
    std::shared_ptr<Aws::Auth::AWSCredentialsProvider> _create_credentials_provider();
    std::shared_ptr<Aws::Auth::AWSCredentialsProvider> _create_provider_from_type(
            const std::string& provider_type);
    std::shared_ptr<Aws::Auth::AWSCredentialsProvider> _create_assume_role_base_provider();

    // HMAC-SHA256 returning hex string
    std::string _hmac_sha256_hex(const std::string& key, const std::string& data);

    std::string _url_encode(const std::string& value);

    std::string _base64url_encode(const std::string& input);

    // Calculate AWS SigV4 signing key
    std::string _calculate_signing_key(const std::string& secret_key, const std::string& date_stamp,
                                       const std::string& region, const std::string& service);

    std::string _hmac_sha256(const std::string& key, const std::string& data);

    std::string _sha256(const std::string& data);

    std::string _get_timestamp();

    std::string _get_date_stamp(const std::string& timestamp);

    bool _should_refresh_credentials();

    Config _config;
    std::shared_ptr<Aws::Auth::AWSCredentialsProvider> _credentials_provider;
    std::mutex _mutex;
    Aws::Auth::AWSCredentials _cached_credentials;
    std::chrono::time_point<std::chrono::system_clock> _credentials_expiry;
};

/**
 * librdkafka OAUTHBEARER callback for AWS MSK IAM authentication.
 * 
 * This callback is invoked by librdkafka when it needs to refresh the
 * OAUTHBEARER token. It uses AwsMskIamAuth to generate the token.
 */
class AwsMskIamOAuthCallback : public RdKafka::OAuthBearerTokenRefreshCb {
public:
    /**
     * Create an OAuth callback from Kafka custom properties.
     * 
     * This factory method checks if AWS MSK IAM authentication is configured
     * (security.protocol=SASL_SSL and sasl.mechanism=OAUTHBEARER) and creates
     * the callback with proper configuration.
     * 
     * @param custom_properties Kafka custom properties map
     * @param brokers Kafka broker list (comma-separated)
     * @return unique_ptr to callback if IAM auth is configured, nullptr otherwise
     */
    static std::unique_ptr<AwsMskIamOAuthCallback> create_from_properties(
            const std::unordered_map<std::string, std::string>& custom_properties,
            const std::string& brokers);

    explicit AwsMskIamOAuthCallback(std::shared_ptr<AwsMskIamAuth> auth,
                                    std::string broker_hostname);

    /**
     * Synchronously refresh and set OAuth token.
     * Can be called directly during initialization or by the callback.
     *
     * @param handle The Kafka handle (consumer or producer)
     * @return Status indicating success or failure
     */
    Status refresh_now(RdKafka::Handle* handle);

    /**
     * Callback invoked by librdkafka to refresh OAuth token.
     *
     * @param handle The Kafka handle (consumer or producer)
     * @param oauthbearer_config Configuration string from 'sasl.oauthbearer.config'
     */
    void oauthbearer_token_refresh_cb(RdKafka::Handle* handle,
                                      const std::string& oauthbearer_config) override;

private:
    std::shared_ptr<AwsMskIamAuth> _auth;
    std::string _broker_hostname;
};

} // namespace doris