PEP and PDP for Secure Authorization with AVP and ABAC

In the first part, PEP and PDP for Secure Authorization with Cognito, I introduced the concept of Policy Enforcement Points (PEPs) and Policy Decision Points (PDPs) using Lambda Authorizer in API Gateway as PEP and a Lambda based service as PDP. The authorization decision was based on the roles users held, represented by Cognito groups, with the permissions mapped in a DynamoDB table.

In the second part, PEP and PDP for Secure Authorization with AVP, we replaced the DynamoDB permission mapping with Amazon Verified Permissions (AVP) and implemented a Role-Based Access Control (RBAC) system using Cedar policies. This gave us a much more powerful way to handle authorization, especially for multi-tenant SaaS applications.

Now, in this third part, we'll take our authorization system to the next level by adding Attribute-Based Access Control (ABAC) alongside RBAC using AVP. This combination provides even more dynamic and context-aware authorization decisions, which is exactly what we need for complex multi-tenant applications.

Understanding Attribute-Based Access Control (ABAC)

So, what exactly is ABAC? At its core, ABAC is an authorization strategy that makes decisions based on attributes associated with users, resources, actions, and the environment. Unlike RBAC, which grants permissions solely based on roles, ABAC evaluates various attributes to make more fine grained authorization decisions.

Think of it this way - RBAC is like having different keys (roles) that open different doors (resources). ABAC, on the other hand, is like having a smart lock that considers not just who you are, but also where you're coming from, what time it is, what you're trying to do, and even what's going on around you before deciding whether to let you in.

As an example some common attributes used in ABAC could be:

User attributes: Department, location, clearance level, seniority
Resource attributes: Classification, owner, sensitivity level
Action attributes: Time of day, encryption status, access method
Environmental attributes: Device type, network location, security level

For example, an ABAC policy might state that Data scientists can access sensitive data, but only from corporate networks, during business hours, and if they have completed security training in the last 6 months.

Why Combine RBAC and ABAC?

But if ABAC is so powerful, why don't we just use it instead of RBAC? I would say that RBAC is really simple to understand and implement, which makes it great for basic authorization needs. However, when permissions need to be more granular or context-dependent, it can become very hard and complex to manage.

On the other hand, ABAC offers amazing flexibility but it can be very complex to implement from scratch and harder to reason about. By combining both approaches, we get the best of both worlds!

In our unicorn racing example, we can use roles to define broad access patterns (Admin, Trainer, Rider), then we'll use an attribute dataAccess to restrict access to specific datasets within those roles.

Implementing ABAC in Amazon Verified Permissions

AVP with its Cedar policy language is perfectly suited for implementing ABAC. Cedar policies can evaluate conditions based on attributes from the principal (user), resource, action, and context.

Token Mapping in AVP - A Critical Concept

An important part of implementing ABAC is understanding how token claims map to your Cedar policy schema. AVP handles this differently depending on whether you're using ID tokens or Access tokens.

ID Token: Claims are mapped to attributes of the principal entity representing the user
Access Token: Claims are mapped to the context.token object in policy evaluation

This difference is super important when writing Cedar policies because it affects how you access attributes in your policy conditions!

For example, with an ID token a policy would look something like this

permit (principal, action, resource)
when { principal.department == "Engineering" };

But with an access token, we need to use the context object and not the principal

permit (principal, action, resource)
when { context.token["department"] == "Engineering" };

In our implementation, we'll use the access token approach since I think it provides a cleaner separation between identity (who the user is) and access permissions and attributes (what they can do and what properties they have).

Why Access Tokens for Authorization

While ID tokens are primarily meant for authentication (proving identity), access tokens are designed specifically for authorization (determining permissions). I see several advantages using access tokens:

Separation of concerns: Authentication details remain in ID tokens, while authorization details live in access tokens
Reduced token size: ID tokens won't be bloated with authorization attributes
Independent lifetimes: Access token expiration can be shorter than ID token expiration
Easier token revocation: You can revoke access without affecting authentication

In our AVP-based PDP, we'll use access tokens and map custom claims to the context for better authorization decisions.

Implementing AVP with RBAC and ABAC

Now that we understand the concepts, let's get our hands dirty and implement our enhanced authorization system!

Architecture

Our architecture remains similar but with a few changes. Where we add a new Lambda function to be able to enrich our access token with custom claims.

This will give us an updated call flow like this.

The changes we need to do would be.

First of all we need to add dataAccess as a custom claim in our access token, to do that we must configure a Pre-token Generation Lambda function, that will be invoked by Cognito.

Second we need to update our schema and add the dataAccess attribute on each action we like the policy to be able to evaluate according to ABAC.

Last we need to update our Cedar policies to be context aware and evaluate both role (RBAC) and attributes (ABAC)

Adding Custom Attributes to Access Tokens

To implement ABAC, we need a way to add custom attributes to our access tokens. We'll use Cognito's Pre-Token Generation Lambda trigger for this purpose, and we need to use version 2 or 3 of the event. With version 1 we can only modify the ID Token.

Below is a implementation that adds the dataAccess attribute from user attributes to the access token:

def handler(event, context):
    user_attributes = event["request"]["userAttributes"]
    claims_to_add_to_access_token = {}

    if "custom:dataAccess" in user_attributes:
        data_access_values = [
            value.strip() for value in user_attributes["custom:dataAccess"].split(",")
        ]
        claims_to_add_to_access_token["custom:dataAccess"] = data_access_values

    response = {
        "claimsAndScopeOverrideDetails": {
            "accessTokenGeneration": {
                "claimsToAddOrOverride": claims_to_add_to_access_token,
            }
        }
    }

    event["response"] = response
    return event

This function checks if the user has a custom:dataAccess attribute. If found, it splits the comma-separated values into an array and adds them to the access token. This way, our AVP policies can check if the user has access to specific data categories.

To have Cognito call our function we need to configure the Lambda Triggers of the User Pool.

  PreTokenGenerationFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: Lambda/PreTokenGeneration
      Handler: index.handler
      Runtime: python3.12
      Architectures:
        - x86_64
      MemorySize: 128
      Description: A Lambda function that adds custom attributes to the JWT access token
      Policies:
        - AWSLambdaBasicExecutionRole

  UserPool:
    Type: AWS::Cognito::UserPool
    Properties:
      UsernameConfiguration:
        CaseSensitive: false
      AutoVerifiedAttributes:
        - email
      UserPoolName: !Sub ${ApplicationName}-user-pool
      Schema:
        - Name: email
          AttributeDataType: String
          Mutable: false
          Required: true
        - Name: name
          AttributeDataType: String
          Mutable: true
          Required: true
        - Name: dataAccess
          AttributeDataType: String
          Mutable: true
          Required: false
      LambdaConfig:
        PreTokenGenerationConfig:
          LambdaArn: !GetAtt PreTokenGenerationFunction.Arn
          LambdaVersion: V2_0

Cedar Policy Schema with Context Attributes

Our Cedar schema defines the structure of principals, resources, actions, and their relationships. For ABAC, we need to update the schema to include context attributes:

{
  "${AVPNameSpace}": {
    "actions": {
      "get /rider": {
        "appliesTo": {
          "context": {
            "type": "Record",
            "attributes": {
              "dataAccess": {
                "type": "String",
                "required": true
              }
            }
          },
          "principalTypes": ["CognitoUser"],
          "resourceTypes": ["Application"]
        }
      }
      // Other actions...
    }
  }
}

This schema enables us to use the dataAccess attribute in our authorization policies.

Policies that Combine RBAC and ABAC

Our Cedar policies now combine RBAC (role-based) and ABAC (attribute-based) logic:

permit(
  principal in ${AVPNameSpace}::CognitoUserGroup::"${UserPoolId}|Admin",
  action in [
    ${AVPNameSpace}::Action::"get /rider",
    ${AVPNameSpace}::Action::"get /riders",
    ${AVPNameSpace}::Action::"get /trainer",
    // more actions...
  ],
  resource
)
when {
  // ABAC condition using attributes
  context.dataAccess == "" ||
  (context.token has "custom:dataAccess" &&
   context.token["custom:dataAccess"].contains(context.dataAccess))
};

This policy allows users in the Admin group to access endpoints (RBAC part), but only if one of these conditions is true (ABAC part):

• No specific data access check is needed (empty context.dataAccess), OR
• The user has the required data access attribute in their token

Let's break down the condition a bit for better understanding.

• If context.dataAccess is empty, we don't enforce data access restrictions
• Otherwise, we check if the user's token has the custom:dataAccess claim that contains the required value

Context-Aware Authorization Logic

Our PDP Lambda function now handles context-based authorization. It extracts relevant attributes from requests and includes them in the AVP authorization call. The below function has been simplified and some parts has been removed, check the source code for a full version.

def validate_permission(event):
    jwt_token = event["jwt_token"]
    resource = event["resource"]
    action = event["action"]
    resource_tags = {}

    if "resource_tags" in event:
        resource_tags = event["resource_tags"]

    parsed_token = decode_token(jwt_token)

    try:
        action_id = f"{action.lower()} {resource.lower()}"
        user_principal = parsed_token["sub"]

        # Call AVP with context for the authorization decision
        context = {
            "dataAccess": resource_tags.get("Data", "")
        }

        auth_response = avp_client.is_authorized_with_token(
            accessToken=jwt_token,
            policyStoreId=policy_store_id,
            action={"actionType": action_type, "actionId": action_id},
            resource={"entityType": resource_type, "entityId": resource_id},
            context=context
        )

        # Cache the result with the context hash
        store_auth_cache(user_principal, action_id, context_hash, auth_response)

        # Generate the response
        response_body = generate_access(
            user_principal,
            "Allow" if auth_response["decision"].upper() == "ALLOW" else "Deny",
            action,
            resource
        )

        return {
            "statusCode": 200,
            "body": json.dumps(response_body),
            "headers": {"Content-Type": "application/json"},
        }

    except Exception as e:
        print(f"Error validating permissions: {e}")

    response_body = generate_access(parsed_token["sub"], "Deny", action, resource)

    return {
        "statusCode": 200,
        "body": json.dumps(response_body),
        "headers": {"Content-Type": "application/json"},
    }

With this implementation, our PDP now evaluates both role-based rules (RBAC) and attribute-based conditions (ABAC).

Context-Aware Caching

One of the challenges with adding ABAC is that authorization decisions now depend not just on who you are and what you're trying to do, but also on context. This means our caching strategy needs to be updated.

We'll include context information in our cache key by creating a hash of the context attributes:

def generate_context_hash(context_dict):
    sorted_items = sorted(context_dict.items())
    context_str = json.dumps(sorted_items)
    return hashlib.md5(context_str.encode()).hexdigest()

def get_auth_cache(principal, action, context_hash):
    try:
        response = dynamodb_client.get_item(
            TableName=table_name, 
            Key={
                "PK": {"S": principal}, 
                "SK": {"S": f"{action}#{context_hash}"}
            }
        )

        item = response.get("Item")
        if not item:
            return None

        ttl = int(item.get("TTL")["N"])
        if ttl and ttl < int(time.time() * 1000):  # TTL is in milliseconds
            return None

        return item.get("Effect")["S"]

    except Exception as e:
        print(f"Error getting auth cache: {e}")
        return None

def store_auth_cache(principal, action, context_hash, auth_response):
    try:
        ttl = int((datetime.now(timezone.utc) + timedelta(hours=12)).timestamp() * 1000)

        effect = "Allow" if auth_response.get("decision", "").upper() == "ALLOW" else "Deny"

        dynamodb_client.put_item(
            TableName=table_name,
            Item={
                "PK": {"S": principal},
                "SK": {"S": f"{action}#{context_hash}"},
                "TTL": {"N": str(ttl)},
                "Effect": {"S": effect},
            },
        )

    except Exception as e:
        print(f"Error storing auth cache: {e}")

By including a hash of the context in our cache key, we ensure that authorization decisions are correctly cached and retrieved based on both the user's role and the relevant context attributes.

Modifying the PEP to Pass Context

Our PEP (Lambda Authorizer) also needs a small update to pass context information to the PDP, this is a very basic implementation, and in a production system you would for sure have something more sophisticated, like looking at resources Tags on AWS resources.

def get_resource_tags_for_path(path):
    if (
        path.startswith("/unicorn")
        or path.startswith("/rider")
        or path.startswith("/trainer")
    ):
        return {"Data": "Unicorn"}
    elif path.startswith("/race"):
        return {"Data": "Races"}
    else:
        return {}

# Inside the handler function
resource_tags = get_resource_tags_for_path(path)

data = {
    "jwt_token": token,
    "resource": path,
    "action": method,
    "resource_tags": resource_tags,
}

This function determines what data category a resource belongs to based on its path, and passes that information to the PDP for context-aware authorization.

Testing the ABAC Implementation

Now that we have our implementation in place, it's time to test it! Here's how we can test our combined RBAC + ABAC solution.

First we need to create or update a user in different Cognito groups (Admin, Trainer, Rider). We also need to assign different dataAccess attributes to these users (e.g., "Unicorn", "Races").

Next navigate to the webpage deployed with the CloudFront distribution and inspect the JWT tokens, cookies.

If we copy the access token and decode that, I use jwt.io, we can see that my user has the claim custom:dataAccess that our PEP and PDP will use later for permissions.

Now we can use a tool like Bruno or Postman to include the access token in the calls to the API. Of course, depending on how you have configured the access you should get either an Allow or Deny back from the PDP resulting in different results from the API. Below is some example views where I use Bruno to call the API with and without access.

Conclusion

By combining RBAC and ABAC with Amazon Verified Permissions, we've created a powerful, flexible authorization system that can handle complex access control requirements. This approach provides the simplicity of role-based access while enabling the fine-grained control of attribute-based decisions.

This implementation approach can be used in an multi-tenant SaaS applications where data segregation, based on tenant, and context-aware access control are critical. With AVP handling the policy evaluation and our serverless architecture for enforcement, we have a scalable, maintainable solution for even the most demanding authorization needs.

Source Code

The entire setup, with detailed deployment instructions, and all the code can be found on Serverless Handbook PEP and PDP

Final Words

Don't forget to follow me on LinkedIn and X for more content, and read rest of my Blogs

As Werner says! Now Go Build!