ve-tos-rust-sdk 2.3.0

/*
 * Copyright (2024) Volcengine
 *
 * Licensed 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.
 */

use crate::config::ConfigHolder;
use crate::constant::{HEADER_AUTHORIZATION, HEADER_CONTENT_SHA256, HEADER_CONTENT_TYPE_LOWER, HEADER_HOST, HEADER_HOST_LOWER, HEADER_PREFIX, HEADER_REQUEST_DATE, HEADER_SECURITY_TOKEN, LONG_DATE_FORMAT, QUERY_ALGORITHM, QUERY_CREDENTIAL, QUERY_DATE, QUERY_EXPIRES, QUERY_SECURITY_TOKEN, QUERY_SIGNATURE, QUERY_SIGNED_HEADERS};
use crate::credential::{Credentials, CredentialsProvider};
use crate::enumeration::HttpMethodType;
use crate::error::TosError;
use crate::http::HttpRequest;
use crate::internal::{check_bucket, hex, hex_sha256, hmac_sha256, url_encode, url_encode_with_safe, InputDescriptor};
use arc_swap::ArcSwap;
use chrono::Utc;
use std::collections::HashMap;
use tracing::log::debug;

pub trait SignerAPI {
    fn pre_signed_url(&self, input: &PreSignedURLInput) -> Result<PreSignedURLOutput, TosError>;
}

#[derive(Debug, Clone, PartialEq, Default)]
pub struct PreSignedURLInput {
    pub(crate) http_method: HttpMethodType,
    pub(crate) bucket: String,
    pub(crate) key: String,
    pub(crate) expires: i64,
    pub(crate) header: HashMap<String, String>,
    pub(crate) query: HashMap<String, String>,
    pub(crate) alternative_endpoint: String,
    pub(crate) is_custom_domain: Option<bool>,
    pub(crate) is_signed_all_headers: bool,
}

impl PreSignedURLInput {
    pub fn new(bucket: impl Into<String>) -> Self {
        let mut input = Self::default();
        input.bucket = bucket.into();
        input
    }

    pub fn new_with_key(bucket: impl Into<String>, key: impl Into<String>) -> Self {
        let mut input = Self::default();
        input.bucket = bucket.into();
        input.key = key.into();
        input
    }

    pub fn http_method(&self) -> &HttpMethodType {
        &self.http_method
    }

    pub fn bucket(&self) -> &str {
        &self.bucket
    }

    pub fn key(&self) -> &str {
        &self.key
    }

    pub fn expires(&self) -> i64 {
        self.expires
    }

    pub fn header(&self) -> &HashMap<String, String> {
        &self.header
    }

    pub fn query(&self) -> &HashMap<String, String> {
        &self.query
    }

    pub fn alternative_endpoint(&self) -> &str {
        &self.alternative_endpoint
    }

    pub fn is_custom_domain(&self) -> Option<bool> {
        self.is_custom_domain
    }

    pub fn is_signed_all_headers(&self) -> bool {
        self.is_signed_all_headers
    }

    pub fn set_http_method(&mut self, http_method: impl Into<HttpMethodType>) {
        self.http_method = http_method.into();
    }

    pub fn set_bucket(&mut self, bucket: impl Into<String>) {
        self.bucket = bucket.into();
    }

    pub fn set_key(&mut self, key: impl Into<String>) {
        self.key = key.into();
    }

    pub fn set_expires(&mut self, expires: i64) {
        self.expires = expires;
    }

    pub fn set_header(&mut self, header: impl Into<HashMap<String, String>>) {
        self.header = header.into();
    }

    pub fn set_query(&mut self, query: impl Into<HashMap<String, String>>) {
        self.query = query.into();
    }

    pub fn set_alternative_endpoint(&mut self, alternative_endpoint: impl Into<String>) {
        self.alternative_endpoint = alternative_endpoint.into();
    }

    pub fn set_is_custom_domain(&mut self, is_custom_domain: impl Into<Option<bool>>) {
        self.is_custom_domain = is_custom_domain.into();
    }

    pub fn set_is_signed_all_headers(&mut self, is_signed_all_headers: bool) {
        self.is_signed_all_headers = is_signed_all_headers;
    }
}

impl InputDescriptor for PreSignedURLInput {
    fn operation(&self) -> &str {
        "PreSignedURL"
    }

    fn bucket(&self) -> Result<&str, TosError> {
        Ok(&self.bucket)
    }

    fn key(&self) -> Result<&str, TosError> {
        Ok(&self.key)
    }
}

pub struct PreSignedURLOutput {
    pub(crate) signed_url: String,
    pub(crate) signed_header: HashMap<String, String>,
}

impl PreSignedURLOutput {
    pub fn signed_url(&self) -> &str {
        &self.signed_url
    }

    pub fn signed_header(&self) -> &HashMap<String, String> {
        &self.signed_header
    }
}

pub(crate) fn pre_signed_url(config_holder: &ArcSwap<ConfigHolder>, ak: &str, sk: &str, security_token: &str, input: &PreSignedURLInput) -> Result<PreSignedURLOutput, TosError> {
    let config_holder = config_holder.load();
    let mut is_custom_domain = config_holder.is_custom_domain;
    if let Some(x) = input.is_custom_domain {
        is_custom_domain = x;
    }
    let bucket = input.bucket();
    if !is_custom_domain {
        check_bucket(bucket.trim())?;
    }
    let schema;
    let domain;
    if input.alternative_endpoint != "" {
        (schema, domain) = config_holder.split_endpoint(input.alternative_endpoint.as_str())?;
    } else {
        schema = "".to_string();
        domain = "".to_string();
    }

    let is_anonymous = ak == "" || sk == "";
    if is_anonymous {
        let signed_url = config_holder.get_endpoint_with_domain(bucket, input.key(), &schema, &domain, true, is_custom_domain);
        return Ok(PreSignedURLOutput {
            signed_url,
            signed_header: HashMap::default(),
        });
    }

    let (long_date, short_date, credential_scope) = calc_date_and_credential_scope(&config_holder.region);

    let mut signed_header = HashMap::with_capacity(input.header.len() + 1);
    for (key, value) in &input.header {
        signed_header.insert(key.clone(), value.to_string());
    }
    signed_header.insert(HEADER_HOST.to_string(), config_holder.get_host_with_domain(bucket, &domain, is_custom_domain));
    let (canonical_headers, signed_headers, mut content_sha256) = calc_canonical_headers(&signed_header, &None, input.is_signed_all_headers);

    let mut query = HashMap::with_capacity(input.query.len() + 7);
    for (key, value) in &input.query {
        query.insert(key.as_str(), value.to_string());
    }
    let mut expires = input.expires;
    if expires <= 0 {
        expires = 3600;
    }
    query.insert(QUERY_ALGORITHM, ALGORITHM.to_string());
    query.insert(QUERY_CREDENTIAL, format!("{}/{}", ak, credential_scope));
    query.insert(QUERY_DATE, long_date.clone());
    query.insert(QUERY_EXPIRES, expires.to_string());
    query.insert(QUERY_SIGNED_HEADERS, signed_headers.clone());
    if security_token != "" {
        query.insert(QUERY_SECURITY_TOKEN, security_token.to_string());
    }
    if content_sha256 == "" {
        content_sha256 = UNSIGNED_PAYLOAD.to_string();
    }

    let mut query = Some(query);
    let canonical_request = calc_canonical_request(input.http_method.as_str(), &mut query, input.key(), &canonical_headers, &signed_headers, content_sha256.as_str());
    let string_to_sign = calc_string_to_sign(&long_date, &credential_scope, &canonical_request);
    let signature = calc_signature(&string_to_sign, &short_date, &config_holder.region, sk)?;
    query.as_mut().unwrap().insert(QUERY_SIGNATURE, signature);
    let mut signed_url = config_holder.get_endpoint_with_domain(input.bucket(), input.key(), &schema, &domain, true, is_custom_domain);

    if let Some(query) = query.as_ref() {
        if query.len() > 0 {
            signed_url.push('?');
            for (idx, kv) in query.iter().enumerate() {
                signed_url.push_str(url_encode(*kv.0).as_str());
                signed_url.push('=');
                signed_url.push_str(kv.1);
                if idx != query.len() - 1 {
                    signed_url.push('&');
                }
            }
        }
    }
    Ok(PreSignedURLOutput {
        signed_url,
        signed_header,
    })
}

pub(crate) const SERVICE_TAG: &str = "tos";
pub(crate) const REQUEST_TAG: &str = "request";
pub(crate) const EMPTY_HASH_PAYLOAD: &str = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";
pub(crate) const UNSIGNED_PAYLOAD: &str = "UNSIGNED-PAYLOAD";
pub(crate) const ALGORITHM: &str = "TOS4-HMAC-SHA256";

pub(crate) fn sign_header<P, C, B>(request: &mut HttpRequest<B>, credentials_provider: &P, config_holder: &ConfigHolder) -> Result<(), TosError>
where
    P: CredentialsProvider<C>,
    C: Credentials,
{
    let cred = credentials_provider.credentials();
    let is_anonymous = cred.ak() == "" || cred.sk() == "";
    if !is_anonymous && cred.security_token() != "" {
        request.header.insert(HEADER_SECURITY_TOKEN, cred.security_token().to_string());
    }
    let (long_date, short_date, credential_scope) = calc_date_and_credential_scope(&config_holder.region);
    let bucket = request.bucket;
    request.header.insert(HEADER_HOST, config_holder.get_host(bucket));
    request.header.insert(HEADER_REQUEST_DATE, long_date.clone());

    let (canonical_headers, signed_headers, _) = calc_canonical_headers(&request.header, &request.meta, false);
    let canonical_request = calc_canonical_request(request.method.as_str(), &mut request.query, request.key, &canonical_headers, &signed_headers, EMPTY_HASH_PAYLOAD);
    debug!("canonical_request: {}", canonical_request);
    if is_anonymous {
        return Ok(());
    }

    let string_to_sign = calc_string_to_sign(&long_date, &credential_scope, &canonical_request);
    // println!("{}", string_to_sign);
    debug!("string_to_sign: {}", string_to_sign);
    let signature = calc_signature(&string_to_sign, &short_date, &config_holder.region, cred.sk())?;
    // println!("{}", signature);

    let mut authorization = String::with_capacity(ALGORITHM.len() + cred.ak().len() +
        credential_scope.len() + signed_headers.len() + signature.len() + 64);
    authorization = authorization + ALGORITHM +
        " Credential=" + cred.ak() + "/" + credential_scope.as_str() +
        ", SignedHeaders=" + signed_headers.as_str() +
        ", Signature=" + signature.as_str();
    // println!("{}", authorization);
    request.header.insert(HEADER_AUTHORIZATION, authorization);
    Ok(())
}

pub(crate) fn calc_string_to_sign(long_date: &str, credential_scope: &str, canonical_request: &str) -> String {
    let mut string_to_sign = String::with_capacity(ALGORITHM.len() + long_date.len() + credential_scope.len() + 3);
    string_to_sign = string_to_sign + ALGORITHM + "\n" + long_date + "\n" + credential_scope + "\n" + hex_sha256(canonical_request).as_str();
    string_to_sign
}

pub(crate) fn calc_signature(string_to_sign: &str, short_date: &str, region: &str, sk: &str) -> Result<String, TosError> {
    let result = hmac_sha256(string_to_sign, hmac_sha256(REQUEST_TAG,
                                                         hmac_sha256(SERVICE_TAG,
                                                                     hmac_sha256(region,
                                                                                 hmac_sha256(short_date, sk)?)?)?)?)?;
    Ok(hex(result))
}

pub(crate) fn calc_canonical_request(method: &str, query: &mut Option<HashMap<&str, String>>, key: &str, canonical_headers: &str, signed_headers: &str, content_sha256: &str) -> String {
    let mut canonical_request = String::with_capacity(key.len() * 2 + 64);
    canonical_request.push_str(method);
    canonical_request.push('\n');
    canonical_request.push('/');
    if key != "" {
        canonical_request.push_str(url_encode_with_safe(key, "/").as_str());
    }
    canonical_request.push('\n');
    if let Some(query) = query.as_mut() {
        let mut keys = Vec::<&str>::with_capacity(query.len());
        for key in query.keys() {
            keys.push(key);
        }
        keys.sort();
        let mut value;
        let mut encoded_value;
        for (idx, key) in keys.iter().enumerate() {
            value = query.get(*key).unwrap();
            canonical_request.push_str(url_encode(*key).as_str());
            canonical_request.push('=');
            if *value != "" {
                encoded_value = url_encode(value);
            } else {
                encoded_value = String::new();
            }
            canonical_request.push_str(encoded_value.as_str());
            if idx != keys.len() - 1 {
                canonical_request.push('&');
            }
            query.insert(*key, encoded_value);
        }
    }
    canonical_request.push('\n');

    canonical_request.push_str(canonical_headers);
    canonical_request.push('\n');

    canonical_request.push_str(signed_headers);
    canonical_request.push('\n');
    canonical_request.push_str(content_sha256);
    // println!("{}", canonical_request);
    canonical_request
}

pub(crate) fn calc_date_and_credential_scope(region: &str) -> (String, String, String) {
    let long_date = Utc::now().format(LONG_DATE_FORMAT).to_string();
    let short_date = &long_date[0..8];
    let mut credential_scope = String::with_capacity(long_date.len() + short_date.len() + SERVICE_TAG.len() + REQUEST_TAG.len() + 3);
    credential_scope.push_str(short_date);
    credential_scope.push('/');
    credential_scope.push_str(region);
    credential_scope.push('/');
    credential_scope.push_str(SERVICE_TAG);
    credential_scope.push('/');
    credential_scope.push_str(REQUEST_TAG);
    let short_date = short_date.to_string();
    (long_date, short_date, credential_scope)
}

pub(crate) fn calc_canonical_headers(header: &HashMap<impl AsRef<str>, String>, meta: &Option<HashMap<String, String>>, is_signed_all_headers: bool) -> (String, String, String) {
    let mut all_header: HashMap<String, &str>;
    let mut keys;
    let mut total_len = 0;
    if let Some(m) = meta {
        keys = Vec::<String>::with_capacity(header.len() + m.len());
        all_header = HashMap::with_capacity(header.len() + m.len());
        for (key, value) in m {
            let key = key.to_lowercase();
            all_header.insert(key.clone(), value);
            total_len += key.len();
            keys.push(key);
        }
    } else {
        keys = Vec::<String>::with_capacity(header.len());
        all_header = HashMap::with_capacity(header.len());
    }

    for (key, value) in header {
        let key = key.as_ref().to_lowercase();
        all_header.insert(key.clone(), value);
        total_len += key.len();
        keys.push(key);
    }

    keys.sort();
    let mut content_sha256 = "".to_string();
    let mut signed_headers = String::with_capacity(total_len + keys.len());
    let mut canonical_headers = String::with_capacity(total_len * 2 + keys.len() * 2);
    for key in keys.iter() {
        if !is_signed_all_headers && key != HEADER_HOST_LOWER && key != HEADER_CONTENT_TYPE_LOWER && !key.starts_with(HEADER_PREFIX) {
            continue;
        }

        let val = all_header.get(key.as_str()).unwrap().trim();
        if key == HEADER_CONTENT_SHA256 {
            content_sha256 = val.to_string();
        }

        signed_headers.push_str(key);
        signed_headers.push(';');
        canonical_headers.push_str(key);
        canonical_headers.push(':');
        canonical_headers.push_str(val);
        canonical_headers.push('\n');
    }

    signed_headers = (&signed_headers[0..signed_headers.len() - 1]).to_string();
    (canonical_headers, signed_headers, content_sha256)
}