Anthony Chu Contact Me

Querying Cosmos DB from Azure API Management Policies

Monday, June 12, 2017

The default user accounts in Azure API Management don't store a lot of data about the user; they mostly just store a name and an email address. But what if we want to store more information about each user and pass it along to the backend API?

In this article, we'll store additional information about our API Management users in a Cosmos DB collection. When our API is called, we'll query Cosmos DB for the user data using an API Management policy and pass this information to the backend via an HTTP header.

In addition, we'll leverage the API Management value cache to prevent calling Cosmos DB on every request.

Set up Cosmos DB

We'll start by creating a Cosmos DB collection and a document for each user in API Management. The id column contains the API Management user id.

Cosmos DB setup

Add an API Management policy

Next, we'll want to add an inbound policy to our API.

<policies>
    <inbound>
        <base />
        <set-variable name="requestDateString" value="@(DateTime.UtcNow.ToString("r"))" />
        <send-request mode="new" response-variable-name="response" timeout="10" ignore-error="false">
            <set-url>https://apim-test.documents.azure.com/dbs/3Ko0AA==/colls/3Ko0AO1cWgA=/docs</set-url>
            <set-method>POST</set-method>
            <set-header name="Authorization" exists-action="override">
                <value>@{
          var verb = "post";
          var resourceType = "docs";
          var resourceLink = "";
          var key = "";
          var keyType = "master";
          var tokenVersion = "1.0";
          var date = context.Variables.GetValueOrDefault<string>("requestDateString");

          var hmacSha256 = new System.Security.Cryptography.HMACSHA256 { Key = Convert.FromBase64String(key) };  

          verb = verb ?? "";  
          resourceType = resourceType ?? "";
          resourceLink = resourceLink ?? "";

          string payLoad = string.Format("{0}\n{1}\n{2}\n{3}\n{4}\n",  
                  verb.ToLowerInvariant(),  
                  resourceType.ToLowerInvariant(),  
                  resourceLink,  
                  date.ToLowerInvariant(),  
                  ""  
          );  

          byte[] hashPayLoad = hmacSha256.ComputeHash(System.Text.Encoding.UTF8.GetBytes(payLoad));  
          string signature = Convert.ToBase64String(hashPayLoad);  

          return System.Uri.EscapeDataString(String.Format("type={0}&ver={1}&sig={2}",  
              keyType,  
              tokenVersion,  
              signature));
        }</value>
            </set-header>
            <set-header name="Content-Type" exists-action="override">
                <value>application/query+json</value>
            </set-header>
            <set-header name="x-ms-documentdb-isquery" exists-action="override">
                <value>True</value>
            </set-header>
            <set-header name="x-ms-date" exists-action="override">
                <value>@(context.Variables.GetValueOrDefault<string>("requestDateString"))</value>
            </set-header>
            <set-header name="x-ms-version" exists-action="override">
                <value>2017-02-22</value>
            </set-header>
            <set-header name="x-ms-query-enable-crosspartition" exists-action="override">
                <value>true</value>
            </set-header>
            <set-body>
        @("{\"query\": \"SELECT * FROM c WHERE c.id = @id\", " +
          "\"parameters\": [{ \"name\": \"@id\", \"value\": \"" + context.User.Id + "\"}]}")
      </set-body>
        </send-request>
        <set-variable name="userJson" value="@(((IResponse)context.Variables["response"]).Body.As<JObject>()["Documents"][0].ToString(Newtonsoft.Json.Formatting.None))" />
        <set-header name="x-api-user" exists-action="override">
            <value>@(Convert.ToBase64String(System.Text.Encoding.UTF8.GetBytes((string)context.Variables["userJson"])))</value>
        </set-header>
    </inbound>
    <backend>
        <base />
    </backend>
    <outbound>
        <base />
    </outbound>
    <on-error>
        <base />
    </on-error>
</policies>

There's a lot going on there so let's unpack each policy statement.

Create an RFC 7231 date

We're querying Cosmos DB using its REST API. Each REST API request needs to be signed using the Cosmos DB master key. In order to limit the risk of replay attacks, each request must contain the current time in RFC 7231 format. We use DateTime.UtcNow.ToString("r") to get this value and store it in a variable called requestDateString.

Query Cosmos DB

How to query Cosmos DB using its REST API is pretty well documented here.

The <send-request> policy statement builds up the request and makes a call to Cosmos DB, storing the result in a variable called response. The large chunk of C# code generates the signature for each request. It's taken directly from here.

We're also using context.User.Id to get the current user's id to look up in Cosmos DB.

There are a couple of references (enclosed in {{ ... }}) to API Management policies. This is handy for values that change between environments and for secrets like the Cosmos DB key.

Set the backend header

The final 2 policy statements extract the first user from the Cosmos DB response and sets a base64-encoded header on the backend request. Note that some more work needs to be done if we don't want to fail on unsuccessful lookups.

Test it out

We can create a simple JavaScript HTTP triggered Azure Function to act as the backend to test this out. It decodes the header and returns it in the body.

module.exports = function (context, req) {
    const userJson = Buffer.from(req.headers['x-api-user'], 'base64').toString();
    const user = JSON.parse(userJson);

    context.res = {
        status: 200,
        body: user
    };

    context.done();
};

If everything is working properly, we should see that the Cosmos DB document belonging to the user is successfully passed to the backend. In the real world, we can use this information to authorize access to resources.

APIM test

Add caching

Now that our call to Cosmos DB works, we can further optimize our policy to cache our HTTP header by user id. Here's an example of how to cache it for up to 60 seconds:

<policies>
    <inbound>
        <base />
        <cache-lookup-value key="@(context.User.Id)" variable-name="userBase64" />
        <choose>
            <when condition="@(!context.Variables.ContainsKey("userBase64"))">
                <set-variable name="requestDateString" value="@(DateTime.UtcNow.ToString("r"))" />
                <send-request mode="new" response-variable-name="response" timeout="10" ignore-error="false">
                    <set-url>https://apim-test.documents.azure.com/dbs/3Ko0AA==/colls/3Ko0AO1cWgA=/docs</set-url>
                    <set-method>POST</set-method>
                    <set-header name="Authorization" exists-action="override">
                        <value>@{
              var verb = "post";
              var resourceType = "docs";
              var resourceLink = "";
              var key = "";
              var keyType = "master";
              var tokenVersion = "1.0";
              var date = context.Variables.GetValueOrDefault<string>("requestDateString");

              var hmacSha256 = new System.Security.Cryptography.HMACSHA256 { Key = Convert.FromBase64String(key) };  

              verb = verb ?? "";  
              resourceType = resourceType ?? "";
              resourceLink = resourceLink ?? "";

              string payLoad = string.Format("{0}\n{1}\n{2}\n{3}\n{4}\n",  
                      verb.ToLowerInvariant(),  
                      resourceType.ToLowerInvariant(),  
                      resourceLink,  
                      date.ToLowerInvariant(),  
                      ""  
              );  

              byte[] hashPayLoad = hmacSha256.ComputeHash(System.Text.Encoding.UTF8.GetBytes(payLoad));  
              string signature = Convert.ToBase64String(hashPayLoad);  

              return System.Uri.EscapeDataString(String.Format("type={0}&ver={1}&sig={2}",  
                  keyType,  
                  tokenVersion,  
                  signature));
            }</value>
                    </set-header>
                    <set-header name="Content-Type" exists-action="override">
                        <value>application/query+json</value>
                    </set-header>
                    <set-header name="x-ms-documentdb-isquery" exists-action="override">
                        <value>True</value>
                    </set-header>
                    <set-header name="x-ms-date" exists-action="override">
                        <value>@(context.Variables.GetValueOrDefault<string>("requestDateString"))</value>
                    </set-header>
                    <set-header name="x-ms-version" exists-action="override">
                        <value>2017-02-22</value>
                    </set-header>
                    <set-header name="x-ms-query-enable-crosspartition" exists-action="override">
                        <value>true</value>
                    </set-header>
                    <set-body>
            @("{\"query\": \"SELECT * FROM c WHERE c.id = @id\", " +
              "\"parameters\": [{ \"name\": \"@id\", \"value\": \"" + context.User.Id + "\"}]}")
          </set-body>
                </send-request>
                <set-variable name="userJson" value="@(((IResponse)context.Variables["response"]).Body.As<JObject>()["Documents"][0].ToString(Newtonsoft.Json.Formatting.None))" />
                <set-variable name="userBase64" value="@(Convert.ToBase64String(System.Text.Encoding.UTF8.GetBytes((string)context.Variables["userJson"])))" />
                <cache-store-value key="@(context.User.Id)" value="@((string)context.Variables["userBase64"])" duration="60" />
            </when>
        </choose>
        <set-header name="x-api-user" exists-action="override">
            <value>@((string)context.Variables["userBase64"])</value>
        </set-header>
    </inbound>
    <backend>
        <base />
    </backend>
    <outbound>
        <base />
    </outbound>
    <on-error>
        <base />
    </on-error>
</policies>

The main additions are the <cache-lookup-value> and <cache-store-value> policy statements. We're also using a <choose> statement to make sure we only call Cosmos DB if the cache doesn't contain a value for the user.

Now if we try it out, we should see in the traces that the header value for each user is cached for up to a minute at a time.