assume_role_with_web_identity doesn't work

[yaizkazani]

Good day and thanks for reading.

I have an issue with my Airflow connections that i'm trying to configure to use assume_role_with_web_identity method.

My Airflow is running in AKS and i'd like some Airflow connections to provide access to AWS resources (S3, SQS, etc.)

Right now i have this:

1. Airflow PODs are running with service account with azure.workload.identity/use: "true" label which make k8s hook to inject the token and some ENV vars into the PODs.
   (AZURE_FEDERATED_TOKEN_FILE=/var/run/secrets/azure/tokens/azure-identity-token is one of them)
2. If i decode the token i see this:

{
  "aud": [
    "api://AzureADTokenExchange"
  ],
  "exp": 1687526808,
  "iat": 1687523208,
  "iss": "https://southcentralus.oic.prod-aks.azure.com/AKS_OIDC_ISSUER",
  "kubernetes.io": {
    "namespace": "MY_AIRFLOW",
    "pod": {
      "name": "airflow-scheduler",
      "uid": "UID"
    },
    "serviceaccount": {
      "name": "airflow",
      "uid": "UID"
    }
  },
  "nbf": 1687523208,
  "sub": "system:serviceaccount:MY_AIRFLOW:airflow"
}

3. I have created OIDC provider in AWS + role that trust this OIDC provider and will allow anyone from correct namespace from AKS cluster to assume it:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Federated": "arn:aws:iam::MY_ACCOUNT:oidc-provider/southcentralus.oic.prod-aks.azure.com/AKS_OIDC_ISSUER"
            },
            "Action": "sts:AssumeRoleWithWebIdentity",
            "Condition": {
                "StringEquals": {
                    "southcentralus.oic.prod-aks.azure.com/AKS_OIDC_ISSUER/:sub": "system:serviceaccount:MY_AIRFLOW:airflow"
                }
            }
        }
    ]
}

4. If i run this command

aws sts assume-role-with-web-identity --role-arn arn:aws:iam::MY_ACCOUNT:role/MY_ROLE --region us-east-1 --role-session-name test --web-identity-token TOKEN_HERE while putting the token from any pod that has "airflow" service account from MY_AIRFLOW namespace - it works, no problems.

However when i do create Airflow connection with this in Extra:

{"role_arn": "arn:aws:iam::MY_ACCOUNT:role/MY_ROLE", "assume_role_method": "assume_role_with_web_identity", "assume_role_with_web_identity_federation": "file", "assume_role_with_web_identity_token_file": "/var/run/secrets/azure/tokens/azure-identity-token", "region_name": "us-east-1"}

It doesn't work.

We get errors related to missing creds:

[2023-07-20, 13:50:40 UTC] {base_aws.py:401} INFO - Airflow Connection: aws_conn_id=s3_remote
[2023-07-20, 13:50:40 UTC] {base.py:73} INFO - Using connection ID 's3_remote' for task execution.
[2023-07-20, 13:50:40 UTC] {base_aws.py:190} INFO - No credentials retrieved from Connection
[2023-07-20, 13:50:40 UTC] {base_aws.py:88} INFO - Retrieving region_name from Connection.extra_config['region_name']
[2023-07-20, 13:50:40 UTC] {base_aws.py:90} INFO - Creating session with aws_access_key_id=None region_name=us-east-1
[2023-07-20, 13:50:40 UTC] {base_aws.py:168} INFO - role_arn is arn:aws:iam::MY_ACCOUNT:role/MY_ROLE
[2023-07-20, 13:50:40 UTC] {base_aws.py:106} INFO - assume_role_method=assume_role_with_web_identity
[2023-07-20, 13:50:40 UTC] {base_aws.py:424} WARNING - Unable to use Airflow Connection for credentials.
[2023-07-20, 13:50:40 UTC] {base_aws.py:425} INFO - Fallback on boto3 credential strategy
[2023-07-20, 13:50:40 UTC] {base_aws.py:428} INFO - Creating session using boto3 credential strategy region_name=None
[2023-07-20, 13:50:40 UTC] {logging_mixin.py:149} WARNING - /home/airflow/.local/lib/python3.8/site-packages/airflow/providers/amazon/aws/hooks/base_aws.py:494 DeprecationWarning: client_type is deprecated. Set client_type from class attribute.
[2023-07-20, 13:50:41 UTC] {taskinstance.py:1824} ERROR - Task failed with exception
...
botocore.exceptions.NoCredentialsError: Unable to locate credentials

I'm not very familiar with Airflow and i rely on our Data Engineers here, however I can't get rid of the feeling that this might be something simple i've missed here.

I have configured Airflow connection according to docs, PODs have correct token that works, what am I missing here? Any help is appreciated, thanks!

[yaizkazani]

Update:

We've created a DAG to test direct access with boto3:

import boto3
from airflow import DAG
from airflow.operators.python import task
from airflow.providers.amazon.aws.hooks.s3 import S3Hook

import dags.common.defaults as dag_defaults
from dags.common.constants import DagTags

s3_path = "MY_PATH"
s3_bucket = "MY_BUCKET"


@task(task_id="read_s3_with_hook")
def read_s3_hoook():
    print("Try s3 Read & Write")
    s3_hook = S3Hook("s3_remote")
    s3_hook.load_string(string_data="test data", bucket_name=s3_bucket, key=s3_path, replace=True)
    key_value = s3_hook.get_key(bucket_name=s3_bucket, key=s3_path)

    print("s3 Read & Write works")
    print(key_value)


@task(task_id="read_s3_with_boto")
def read_s3_boto():
    token = ""
    with open("/var/run/secrets/azure/tokens/azure-identity-token", mode="r") as token_file:
        token = token_file.read()

    client = boto3.client("sts", region_name="us-east-1")
    response = client.assume_role_with_web_identity(
        DurationSeconds=3600,
        RoleArn="arn:aws:iam::MY_ACCOUNT:role/MY_ROLE",
        RoleSessionName="test",
        WebIdentityToken=token,
    )

    access_key = response["Credentials"]["AccessKeyId"]
    secret_key = response["Credentials"]["SecretAccessKey"]
    session_token = response["Credentials"]["SessionToken"]

    s3_client = boto3.client(
        "s3", aws_access_key_id=access_key, aws_secret_access_key=secret_key, aws_session_token=session_token
    )

    s3_file = s3_client.get_object(Bucket="MY_BUCKET", Key="testing/MY_FILE.xlsx")

    print(s3_file)


with DAG(
    dag_id="test_s3_read_write",
    default_args=dag_defaults.default_args(dag_defaults.Team.MY_TEAM),
    schedule=None,
    catchup=False,
    tags=[
        DagTags.TeamTags.MY_TEAM.value,
    ],
) as dag:
    s3_hook_operator = read_s3_hoook()
    s3_boto_operator = read_s3_boto()

So s3_boto_operator step - works just fine
s3_hook_operator - fails with "Unable to locate credentials"

[yaizkazani]

It seems that our AWS provider version is too old to support this. I will double check and close this if this is true.

[yaizkazani]

Sorry for this to be open - we have some slow processes updating providers. I will update this ASAP

[Taragolis]

Another workaround could be use boto3 default credential strategy with environment variables or profiles: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html#assume-role-with-web-identity-provider

Option 1: Environment variables

You need to inject this variables to your pod
AWS_ROLE_ARN=arn:aws:iam::MY_ACCOUNT:role/MY_ROLE
AWS_WEB_IDENTITY_TOKEN_FILE=/var/run/secrets/azure/tokens/azure-identity-token

So in this case your AWS Connection should be something like

{"region_name": "us-east-1"}

Option 2: AWS config file

Note

This is less flexible because required create files, and depend on your deployment process this could turns into the hell.
In the other hand it could allow "infinity" assume role chain

First of all you need to create ~/.aws/config (this is default patch which could be change by set AWS_CONFIG_FILE env var) with at least this config

[profile awesome-profile]
role_arn=arn:aws:iam::MY_ACCOUNT:role/MY_ROLE
web_identity_token_file=/var/run/secrets/azure/tokens/azure-identity-token

So in this case your AWS Connection should be something like

{
  "profile_name": "awesome-profile",
  "region_name": "us-east-1"
}

---

Important

If you choose one of this methods you should not provide role_arn into your AWS Connection, unless you need to switch from arn:aws:iam::MY_ACCOUNT:role/MY_ROLE to another role

[yaizkazani]

Thank you, this is definitely something to be aware of.

[yaizkazani]

Also i will close this as we have tested and it works - just need to update Airflow AWS provider to version that support that feature

[Thamban777]

assume_role_with_web_identity doesn't work on eks 1.30

EKS Verison : 1.30

my-values.yaml:

env:

• name: AWS_ROLE_ARN
  value: arn:aws:iam::accountID:role/eksctl-roll
• name: AWS_WEB_IDENTITY_TOKEN_FILE
  value: /var/run/secrets/eks.amazonaws.com/serviceaccount/token
• name: AWS_DEFAULT_REGION
  value: us-east-2
• name: AWS_REGION
  value: us-east-2

Volumes for all airflow containers

volumes:

• name: aws-iam-token
  projected:
  sources:
  • serviceAccountToken:
    audience: sts.amazonaws.com
    expirationSeconds: 86400
    path: token

VolumeMounts for all airflow containers

volumeMounts:
- name: aws-iam-token
mountPath: /var/run/secrets/eks.amazonaws.com/serviceaccount
readOnly: true

extraInitContainers:

• name: copy-dags
  image: public.ecr.aws/aws-cli/aws-cli:2.11.18
  command:
  • /bin/sh
  • -c
  • |
    AWS_EC2_METADATA_DISABLED=true aws s3 sync s3://mys3/apps/airflow/dags/ /opt/airflow/dags
    env:
  • name: AWS_ROLE_ARN
    value: arn:aws:iam::accountID:role/eksctl-role
  • name: AWS_WEB_IDENTITY_TOKEN_FILE
    value: /var/run/secrets/eks.amazonaws.com/serviceaccount/token
  • name: AWS_DEFAULT_REGION
    value: us-east-2
  • name: AWS_REGION
    value: us-east-2
  • name: AWS_STS_REGIONAL_ENDPOINTS
    value: regional
  • name: AWS_EC2_METADATA_DISABLED
    value: "true"
    volumeMounts:
  • name: aws-iam-token
    mountPath: /var/run/secrets/eks.amazonaws.com/serviceaccount
    readOnly: true
  • name: dags
    mountPath: /opt/airflow/dags