ve-tos-rust-sdk 2.0.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 std::collections::HashMap;

use chrono::Utc;

use crate::config::ConfigHolder;
use crate::constant::{HEADER_AUTHORIZATION, HEADER_CONTENT_TYPE_LOWER, HEADER_HOST, HEADER_HOST_LOWER, HEADER_PREFIX, HEADER_REQUEST_DATE, HEADER_SECURITY_TOKEN, LONG_DATE_FORMAT};
use crate::credential::{Credentials, CredentialsProvider};
use crate::error::TosError;
use crate::http::HttpRequest;
use crate::internal::{hex, hex_sha256, hmac_sha256, url_encode, url_encode_with_safe};

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 mut bucket = request.bucket;
    if config_holder.is_custom_domain {
        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);
    let canonical_request = calc_canonical_request(request.method.as_str(), &mut request.query, request.key, &canonical_headers, &signed_headers, false);
    if is_anonymous {
        return Ok(());
    }

    let string_to_sign = calc_string_to_sign(&long_date, &credential_scope, &canonical_request);
    // println!("{}", 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, is_query: bool) -> 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');
    if is_query {
        canonical_request.push_str(UNSIGNED_PAYLOAD);
    } else {
        canonical_request.push_str(EMPTY_HASH_PAYLOAD);
    }

    // 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<&str, String>, meta: &Option<HashMap<String, String>>) -> (String, String) {
    let mut all_header: HashMap<&str, &str>;
    let mut keys;
    let mut total_len = 0;
    if let Some(m) = meta {
        all_header = HashMap::with_capacity(header.len() + m.len());
        keys = Vec::<&str>::with_capacity(header.len() + m.len());
        for (key, value) in m {
            keys.push(key);
            total_len += key.len();
            all_header.insert(key, value);
        }
    } else {
        all_header = HashMap::with_capacity(header.len());
        keys = Vec::<&str>::with_capacity(header.len());
    }

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

    keys.sort();
    let mut signed_headers = String::with_capacity(total_len + keys.len());
    let mut canonical_headers = String::with_capacity(total_len * 2 + keys.len() * 2);
    let mut lower_key;
    if let Some(m) = meta {
        for (idx, key) in keys.iter().enumerate() {
            lower_key = (*key).to_lowercase();
            if lower_key != HEADER_HOST_LOWER && lower_key != HEADER_CONTENT_TYPE_LOWER && !lower_key.starts_with(HEADER_PREFIX) {
                continue;
            }
            signed_headers.push_str(lower_key.as_str());
            if idx != keys.len() - 1 {
                signed_headers.push(';');
            }
            canonical_headers.push_str(lower_key.as_str());
            canonical_headers.push(':');
            if let Some(value) = header.get(*key) {
                canonical_headers.push_str(value.trim());
            } else {
                canonical_headers.push_str(m.get(*key).unwrap().trim());
            }
            canonical_headers.push('\n');
        }
    } else {
        for (idx, key) in keys.iter().enumerate() {
            lower_key = (*key).to_lowercase();
            if lower_key != HEADER_HOST_LOWER && lower_key != HEADER_CONTENT_TYPE_LOWER && !lower_key.starts_with(HEADER_PREFIX) {
                continue;
            }
            signed_headers.push_str(lower_key.as_str());
            if idx != keys.len() - 1 {
                signed_headers.push(';');
            }
            canonical_headers.push_str(lower_key.as_str());
            canonical_headers.push(':');
            canonical_headers.push_str(header.get(*key).unwrap().trim());
            canonical_headers.push('\n');
        }
    }

    (canonical_headers, signed_headers)
}