Implementing Your Custom Policy Features in Rust

Before You Begin

Create a New Project
Define a Schema Definition

You may need to add or redefine configuration parameters while implementing your policy.
Review Policy Template, Wrapped Functions, and Flow Cancellation to learn about the policy’s source code structure before adding additional policy capabilities.

Policy Template

Omni Gateway Policy Development Kit (PDK) provides the initial src/lib.rs file in the policy project as a template to begin implementing your policy:

// Copyright 2024 Salesforce, Inc. All rights reserved.
mod generated;

use anyhow::{anyhow, Result};

use pdk::hl::*;
use pdk::logger;

use crate::generated::config::Config;

// This filter shows how to log a specific request header.
// You can extend the function and use the configurations exposed in config.rs file
async fn request_filter(request_state: RequestState, _config: &Config) {
    let headers_state = request_state.into_headers_state().await;
    let token = headers_state.handler().header("Token").unwrap_or_default();
    // Log the header value
    logger::info!("Header value: {token}");
}

#[entrypoint]
async fn configure(launcher: Launcher, Configuration(bytes): Configuration) -> Result<()> {
    let config: Config = serde_json::from_slice(&bytes).map_err(|err| {
        anyhow!(
            "Failed to parse configuration '{}'. Cause: {}",
            String::from_utf8_lossy(&bytes),
            err
        )
    })?;
    let filter = on_request(|rs| request_filter(rs, &config));
    launcher.launch(filter).await?;
    Ok(())
}

The policy performs some simple logging. For each incoming request, the policy logs the header token at the info log level.

The following elements configure the policy:

use pdk::api::hl::*;: Imports all the components of PDK into the lib.rs source code.
#[entrypoint] function: Executes when the policy is applied and calls the request_filter function. Variables defined inside this function are avaliable while the policy is applied.

The #[entrypoint] function receives the following parameters:
- Launcher: Sets the filter functions.
- Configuration: Provides the configuration parameters defined in Defining a Policy Schema Definition.
request_filter: Executes once per every request sent to the API instance to which the policy is applied. This function implements the filtering performed by the example policy. Variables defined inside this function are available for the duration of the request.

The request_filter is an example of a wrapped function.

The #[entrypoint] function executes wrapped functions:
```
let filter = on_request(|rs| request_filter(rs, &config));
launcher.launch(filter).await?;
```

Wrapped Functions

The on_response or on_request wrapper defines when the filter function executes.

You can define custom functions other than the provided requests_filter and response_filter in the #[entrypoint] function. The requests_filter and response_filter are only examples.

Instead of filtering an incoming request as shown in the previous example, you can process the outgoing response by using the on_response wrapper:

async fn response_filter(response_state: ResponseState, _config: &Config) {
    ...
}

#[entrypoint]
async fn configure(launcher: Launcher, Configuration(bytes): Configuration) -> Result<()> {
    let config: Config = serde_json::from_slice(&bytes).map_err(|err| {
        anyhow!(
            "Failed to parse configuration '{}'. Cause: {}",
            String::from_utf8_lossy(&bytes),
            err
        )
    })?;
    let filter = on_response(|rs| response_filter(rs, &config));
    launcher.launch(filter).await?;
    Ok(())
}

The request_filter function is now the response_filter function.

To filter both requests and responses, use both wrappers:

async fn request_filter(request_state: RequestState,_config: &Config) {
    ...
}

async fn response_filter(response_state: ResponseState,_config: &Config) {
    ...
}

#[entrypoint]
async fn configure(launcher: Launcher, Configuration(bytes): Configuration) -> Result<()> {
    let config: Config = serde_json::from_slice(&bytes).map_err(|err| {
        anyhow!(
            "Failed to parse configuration '{}'. Cause: {}",
            String::from_utf8_lossy(&bytes),
            err
        )
    })?;
    let filter = on_request(|rs| request_filter(rs, &config))
    .on_response(|rs| response_filter(rs, &config));
    launcher.launch(filter).await?;
    Ok(())
}

Flow Cancellation

Custom policies run on a single-threaded environment. However, every set of wrapped functions (a on_request and a on_response async function) run as a concurrent task, meaning that a single policy instance handles multiple requests concurrently.

To support this behavior, every .await is a potential task interruption point. When an .await is invoked, the underlying async runtime can set the current task to sleep for a moment, and awake another task that has not finished. The async runtime can cancel a task and never run code that comes after the last .await invocation. on_request and on_response functions assume that .await invocations are potential cancellation points.

The most common situation, as shown in the following code example, for task cancellation is when a request function waiting for a body and an upstream policy returns an early response:

// Request function for upstream policy
async fn upstream_request(state: RequestState) -> Response {
    Response::new(404)
}

// Request function for downstream policy
async fn downstream_request(state: RequestState) {

    // Request function will be cancelled after this .await point.
    let body_state = state.into_body_state().await;

    // Code here will never be executed
}

Discover PDK features with Provided Rust Source Code

PDK provides code examples for the following policy features. Review Policy Template and Wrapped Functions before adding the additional policy capabilities: