simple_azure_monitor_data_collector 0.1.0

use reqwest::Client as ReqwestClient;
use serde_json::Value;
use thiserror::Error;

/// A `LogClient` can be used to send records to the Data Collector API.
pub struct LogClient {
    api_version: String,
    inner_client: ReqwestClient,
    key: String,
    url: url::Url,
    workspace_id: String,
}

impl LogClient {
    /// Creates a new `LogClient`.
    pub fn new(
        api_version: String,
        workspace_id: String,
        key: String,
        timeout_in_seconds: u64,
    ) -> Result<Self, NewLogClientError> {
        use std::time::Duration;

        let inner_client = ReqwestClient::builder()
            .timeout(Duration::from_secs(timeout_in_seconds))
            .build()?;

        let url = Self::create_url(&workspace_id, &api_version)?;
        Ok(Self {
            api_version,
            inner_client,
            key,
            url,
            workspace_id,
        })
    }

    /// The url this `LogClient` will send records to.
    pub fn url(&self) -> &url::Url {
        &self.url
    }

    /// The Data Collector API version being used by this `LogClient`.
    pub fn api_version(&self) -> &str {
        &self.api_version
    }

    /// The Azure workspace ID of this `LogClient`.
    pub fn workspace_id(&self) -> &str {
        &self.workspace_id
    }

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

    fn create_url(
        workspace_id: &str,
        api_version: &str,
    ) -> Result<url::Url, url::ParseError> {
        let url = format!(
            "https://{wsid}.ods.opinsights.azure.com/api/logs?api-version={ver}",
            wsid = workspace_id,
            ver = api_version
        );
        url::Url::parse(&url)
    }

    fn build_log_request(
        &self,
        key: &str,
        workspace_id: &str,
        records: Records,
    ) -> Result<reqwest::Request, RequestBuildError> {
        let body = serde_json::to_string(&records.records)?;
        let content_length = body.as_bytes().len();

        let timestamp = rfc_1123_utc_now_timestamp();

        let auth_header_val =
            auth_header_value(key, workspace_id, content_length, &timestamp)?;

        let mut req_builder = self
            .inner_client
            .post(self.url().clone())
            .header("Authorization", auth_header_val)
            .header("Content-Type", "application/json")
            .header("Log-Type", records.log_type)
            .header("x-ms-date", timestamp)
            .body(body);

        if let Some(resource_id) = records.azure_resource_id {
            req_builder =
                req_builder.header("x-ms-AzureResourceId", resource_id);
        }

        if let Some(time_generated_field) = records.time_generated_field {
            req_builder = req_builder
                .header("time-generated-field", time_generated_field);
        }

        req_builder.build().map_err(|err| err.into())
    }

    /// Log the given records to the Data Collector API.
    pub async fn log(&self, records: Records) -> Result<(), LogError> {
        let req =
            self.build_log_request(self.key(), self.workspace_id(), records)?;
        let res = self.inner_client.execute(req).await?;
        check_response(res)
    }
}

/// If the HTTP response is a success response, return Ok. Otherwise
/// return an Err.
fn check_response(res: reqwest::Response) -> Result<(), LogError> {
    let status = res.status();

    if status.is_success() {
        Ok(())
    } else {
        use reqwest::StatusCode;
        match status {
            StatusCode::BAD_REQUEST => {
                Err(LogError::BadRequest { response: res })
            }
            StatusCode::FORBIDDEN => Err(LogError::Forbidden),
            StatusCode::NOT_FOUND => Err(LogError::NotFound),
            StatusCode::TOO_MANY_REQUESTS => Err(LogError::TooManyRequests),
            StatusCode::INTERNAL_SERVER_ERROR => {
                Err(LogError::InternalServerError)
            }
            StatusCode::SERVICE_UNAVAILABLE => {
                Err(LogError::ServiceUnavailable)
            }
            _ => Err(LogError::Other { response: res }),
        }
    }
}

/// Returns a string containing the current UTC time formatted according to
/// RFC1123.
fn rfc_1123_utc_now_timestamp() -> String {
    format!(
        "{}",
        chrono::offset::Utc::now().format("%a, %d %b %Y %H:%M:%S GMT")
    )
}

fn auth_header_value(
    key: &str,
    workspace_id: &str,
    content_length: usize,
    x_ms_date: &str,
) -> Result<String, base64::DecodeError> {
    Ok(format!(
        "SharedKey {ws_id}:{sig}",
        ws_id = workspace_id,
        sig = signature(key, content_length, x_ms_date)?
    ))
}

fn signature(
    key: &str,
    content_length: usize,
    x_ms_date: &str,
) -> Result<String, base64::DecodeError> {
    use hmac::{Hmac, Mac};
    type HmacSha256 = Hmac<sha2::Sha256>;

    let str_to_sign = string_to_sign(content_length, x_ms_date);

    let decoded_key = base64::decode(key)?;
    let mut hmac = HmacSha256::new_varkey(&decoded_key)
        .expect("HMAC can take key of any size");

    hmac.input(str_to_sign.as_bytes());

    let result = hmac.result();
    // Get the underlying array using the `code()` method. Note that
    // incorrect use of the code value may permit timing attacks which defeat
    // the security provided by the `MacResult`:
    let hmac_bytes = result.code();

    let b64_str = base64::encode(&hmac_bytes);
    Ok(b64_str)
}

fn string_to_sign(content_length: usize, x_ms_date: &str) -> String {
    format!(
        "POST\n{content_length}\napplication/json\nx-ms-date:{date}\n/api/logs",
        content_length = content_length,
        date = x_ms_date
    )
}

/// Wraps a JSON value. The JSON value should contain an array of JSON objects,
/// which represent the individual records being sent to the Data Collector API.
///
/// The records in this struct will be sent in a single request to the service.
/// The number of records in this stuct should be limited such that the
/// request's size will not exceed the limits imposed by the service.
#[derive(Clone, Debug, PartialEq)]
pub struct Records {
    log_type: String,
    records: Value,
    azure_resource_id: Option<String>,
    time_generated_field: Option<String>,
}

impl Records {
    /// Construct a new `Records` struct.
    pub fn new(
        log_type: String,
        records: Value,
        azure_resource_id: Option<String>,
        time_generated_field: Option<String>,
    ) -> Result<Self, NewRecordsError> {
        Self::check_json_value(&records).map_err(NewRecordsError::new)?;
        Ok(Self {
            log_type,
            records,
            azure_resource_id,
            time_generated_field,
        })
    }

    fn check_json_value(val: &Value) -> Result<(), String> {
        if let Some(array) = val.as_array() {
            for elem in array {
                if !elem.is_object() {
                    return Err(
                        "Each element of the JSON array must be an object"
                            .into(),
                    );
                }
            }
            Ok(())
        } else {
            Err("JSON value must be an array".into())
        }
    }
}

/// Error that occur when constructing a new `Records` struct.
#[derive(Clone, Debug, Error)]
#[error("Could not create a valid `Records` struct: {msg}")]
pub struct NewRecordsError {
    msg: String,
}

impl NewRecordsError {
    fn new(msg: String) -> Self {
        Self { msg }
    }
}

/// Errors that can occur when building a HTTP request.
#[derive(Debug, Error)]
pub enum RequestBuildError {
    /// Error occurred when decoding base64.
    #[error("Could not decode base64")]
    Base64Decode(#[from] base64::DecodeError),
    /// Error occurred in the underlying HTTP library.
    #[error("HTTP error")]
    Http(#[from] reqwest::Error),
    /// Could not convert a JSON value to a UTF-8 string.
    #[error("Could not convert JSON value to a UTF-8 string")]
    JsonToString(#[from] serde_json::Error),
}

/// Errors that can occur when creating a new `LogClient`.
#[derive(Debug, Error)]
pub enum NewLogClientError {
    #[error("HTTP Error")]
    Http(#[from] reqwest::Error),
    #[error("Could not create a valid URL for this client")]
    ParseUrl(#[from] url::ParseError),
}

/// Errors that can occur when logging data to the Data Collector API.
#[derive(Debug, Error)]
pub enum LogError {
    /// Error occurred in the underlying HTTP library.
    #[error("HTTP Error")]
    Http(#[from] reqwest::Error),
    /// Error occurred when building a HTTP request.
    #[error("Could not build a HTTP request")]
    RequestBuilder(#[from] RequestBuildError),
    /// The response indicated that the request was not valid.
    #[error("The HTTP request was not valid")]
    BadRequest { response: reqwest::Response },
    /// The service failed to authenticate the request. Verify
    /// that the workspace ID and connection key are valid.
    #[error("The service failed to authenticate the request")]
    Forbidden,
    /// Either the URL provided is incorrect, or the request is too large.
    #[error(
        "Either the URL provided is incorrect, or the request is too large"
    )]
    NotFound,
    /// The service is experiencing a high volume of data from your account.
    /// Please retry the request later.
    #[error(
        "The service is experiencing a high volume of data from your account"
    )]
    TooManyRequests,
    /// The service encountered an internal error. Please retry the request.
    #[error("The service encountered an internal error")]
    InternalServerError,
    /// The service currently is unavailable to receive requests.
    /// Please retry your request.
    #[error("The service currently is unavailable to receive requests")]
    ServiceUnavailable,
    /// The HTTP reponse had an unexpected status code.
    #[error("The HTTP reponse had an unexpected status code")]
    Other { response: reqwest::Response },
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    const API_VERSION: &str = "2016-04-01";

    #[tokio::test]
    async fn log_works_with_some_records() {
        let client = client();
        let records = records_from_value(json!(
            [
                {
                    "slot_ID": 12345,
                    "ID": "5cdad72f-c848-4df0-8aaa-ffe033e75d57",
                    "availability_Value": 100,
                    "performance_Value": 6.954,
                    "measurement_Name": "last_one_hour",
                    "duration": 3600,
                    "warning_Threshold": 0,
                    "critical_Threshold": 0,
                    "IsActive": "true"
                },
                {
                    "slot_ID": 67890,
                    "ID": "b6bee458-fb65-492e-996d-61c4d7fbb942",
                    "availability_Value": 100,
                    "performance_Value": 3.379,
                    "measurement_Name": "last_one_hour",
                    "duration": 3600,
                    "warning_Threshold": 0,
                    "critical_Threshold": 0,
                    "IsActive": "false"
                }
            ]
        ));
        client.log(records).await.unwrap();
    }

    #[tokio::test]
    async fn log_works_with_no_records() {
        let client = client();
        let records = records_from_value(json!([]));
        client.log(records).await.unwrap();
    }

    #[test]
    fn string_to_sign() {
        assert_eq!(
            super::string_to_sign(1024, "Mon, 04 Apr 2016 08:00:00 GMT"),
            "POST\n1024\napplication/json\nx-ms-date:Mon, 04 Apr 2016 08:00:00 GMT\n/api/logs"
        )
    }

    fn records_from_value(val: serde_json::Value) -> Records {
        Records::new("testlog".into(), val, None, None).unwrap()
    }

    fn client() -> LogClient {
        let api_version = API_VERSION.into();
        let timeout = 60;
        LogClient::new(api_version, workspace_id(), key(), timeout).unwrap()
    }

    fn workspace_id() -> String {
        std::env::var("AZURE_LOG_ANALYTICS_WORKSPACES_WORKSPACE_ID").unwrap()
    }

    fn key() -> String {
        std::env::var("AZURE_LOG_ANALYTICS_WORKSPACES_PRIMARY_KEY").unwrap()
    }
}