Lab 4: Securing APIs at runtime using Flex Gateway Policy Development Kit (Optional)

// Copyright 2023 Salesforce, Inc. All rights reserved.

use hmac::Hmac;
use pdk::logger;

use hmac::Mac;

use md5::Md5;
use sha2::Sha256;

type HmacMd5 = Hmac<Md5>;
type HmacSha256 = Hmac<Sha256>;


/***************************
 * Verify the message against a key provided
 */
pub fn verify(message: &Vec<u8>, secret_key: String, hash: String, alg: String) -> Result<(), String> {

    // Decode the provided hash from hex
    let decoded_hash = match hex::decode(hash.clone().into_bytes()) {
        Ok(bytes) => bytes,
        Err(_) => {
            // Log and return an error if the hash is invalid
            logger::info!("====== Policy Violation, invalid signature hash");
            return Err("HMAC validation failed, invalid signature hash".to_owned());
        },
    };

    logger::info!("======= Validating for {alg}");
    let original_body = String::from_utf8(message.clone()).expect("Our bytes should be valid utf8");
    logger::info!("======= Body: {original_body} ");

    // Create an HMAC context using the provided secret key
    match alg.as_str() {
        "SHA-256" => match HmacSha256::new_from_slice(&secret_key.as_bytes()) {
            Ok(mut mac) => {
                // Update the HMAC context with the message
                mac.update(message.as_slice());

                // Verify the HMAC against the decoded hash
                if mac.verify_slice(&decoded_hash).is_err() {
                    // Log and return an error if the validation fails
                    logger::debug!("==== Policy Violation, HMAC validation failed");
                    return Err("SHA-256 HMAC validation failed".to_owned());
                } else {
                    logger::info!("==== Verified!!");
                }
            },
            Err(_error) => {
                // Return an error if the secret key is invalid
                return Err("Invalid configured secret key".to_string())
            }
        },
        "MD-5" => match HmacMd5::new_from_slice(&secret_key.as_bytes()) {
            Ok(mut mac) => {
                // Update the HMAC context with the message
                mac.update(message.as_slice());

                // Verify the HMAC against the decoded hash
                if mac.verify_slice(&decoded_hash).is_err() {
                    // Log and return an error if the validation fails
                    logger::debug!("==== Policy Violation, HMAC validation failed");
                    return Err("MD-5 HMAC validation failed".to_owned());
                } else {
                    logger::info!("==== Verified!!");
                }
            },
            Err(_error) => {
                // Return an error if the secret key is invalid
                return Err("Invalid configured secret key".to_string())
            }
        },
        _e =>  {
                // Return an error if the algorithm is invalid
                return Err("======Unsupported algorithm {alg}".to_string())
            }
    };

    // Return success if the HMAC validation is successful
    Ok(())
}

use serde_json::json;
mod hmac;

async fn request_filter(request_state: RequestState, _config: &Config) -> Flow<()> {
    logger::info!("============ START");
    let headers_state = request_state.into_headers_state().await;
    let headers_handler = headers_state.handler();

    let hmac_algorithm: String = _config.hmac_algorithm.clone();
    let secret_key: String = _config.secret_key.clone();
    let signature_header_name: String = _config.signature_header_name.clone();
    let signature_header_value: String = headers_handler.header(&signature_header_name).unwrap_or_default();
    let content_length = headers_handler.header(&"Content-Length");

    let body_state = headers_state.into_body_state().await;
    let body_handler = body_state.handler();


    // Log the header value
    logger::info!("======= Header name: {signature_header_name} with header value {signature_header_value}");


    //Get the body
    let body = body_handler.body();

    // Validate the hash if the request has no body
    match content_length {
        Some(cl) => {
            logger::info!("==== Content-Length: {cl}");
            match cl.parse::<u64>() {
                Ok(len) => {

                    // Nothing to verify
                    if len == 0 {
                        logger::info!("==== Empty request body accepted");
                        return Flow::Continue(());
                    }
                    logger::info!("==== Verifying request body");
                    // Verify the HMAC signature hash with the empty body.
                    // If the verification fails, reject the request and log the error.
                    match hmac::verify(&body, secret_key, signature_header_value, hmac_algorithm) {
                        Ok(_v) => {
                            let original_body =  String::from_utf8(body).expect("Our bytes should be valid utf8");
                            logger::info!("==== Request body accepted: {original_body}");
                            return Flow::Continue(()) ;
                        },
                        Err(e) => {
                            logger::info!("==== Error during verification: {e}");
                            let response_body = json!({
                                "error": e
                            }).to_string();
                            return Flow::Break(Response::new(400).with_body(response_body))
                        }
                    }
                },
                Err(_e) => {
                    logger::info!("==== Content-Length: not parseable: {cl}");
                    let response_body = json!({
                        "error": "========= Unable to parse content-length {cl} "
                    }).to_string();
                    return Flow::Break(Response::new(400).with_body(response_body))

                }
            }
        },
        None => {
            logger::info!("==== Request body accepted by content-length error");
            return Flow::Continue(()) ;
        }
    }
}

cd <project home>
make build