DRIVERS-2945 AWS EKS Pod Identity (#655)

Co-authored-by: Noah Stapp <[email protected]>

Author: blink1073

.evergreen/auth_aws/README.md

@@ -36,9 +36,7 @@ for setting up local secrets handling are in the wiki page.
 The ECS test variant requires a slightly different approach, since we need to run the code in a container.
 
 Set up a `run-mongodb-aws-ecs-test.sh` script that will run on the container.  This script should be
-copied to `${DRIVERS_TOOLS}/.evergreen/auth_aws/src/.evergreen`.  The driver code and test code should
-be compiled if necessary, and then compressed into a `src.tgz` file that will be expanded and used in
-the container.
+copied to `${DRIVERS_TOOLS}/.evergreen/auth_aws/src/.evergreen`.
 
 ```bash
 # Set up the target directory.
@@ -52,6 +50,26 @@ cp ${PROJECT_DIRECTORY}/.evergreen/run-mongodb-aws-ecs-test.sh $ECS_SRC_DIR/.eve
 PROJECT_DIRECTORY="$ECS_SRC_DIR" MONGODB_BINARIES="/path/to/mongodb/bin" $AUTH_AWS_DIR/aws_setup.sh ecs
 ```
 
+## EKS Test Process
+
+The EKS Pod Identity test variant also requires a slightly different approach, since we need to run the code in a
+kubernetes pod. The MongoDB server version running in the k8s cluster can be set with `MONGODB_VERSION`.
+
+Set up a `run-mongodb-aws-eks-test.sh` script that will run on the pod.  This script should be
+copied to `${DRIVERS_TOOLS}/.evergreen/auth_aws/src/.evergreen`.
+
+```bash
+# Set up the target directory.
+EKS_SRC_DIR=${DRIVERS_TOOLS}/.evergreen/auth_aws/src
+mkdir -p $EKS_SRC_DIR/.evergreen
+# Move the test script to the correct location.
+cp ${PROJECT_DIRECTORY}/.evergreen/run-mongodb-aws-eks-test.sh $EKS_SRC_DIR/.evergreen
+# Driver-specific - compile/build code if needed.
+# Driver-specific - move artifacts needed for test to $EKS_SRC_DIR
+# Run the test
+PROJECT_DIRECTORY="$EKS_SRC_DIR" $AUTH_AWS_DIR/aws_setup.sh eks
+```
+
 ## Deprecated Scripts
 
 The top-level JavaScript files in this directory are deprecated and no longer needed when

.evergreen/auth_aws/aws_setup.sh

@@ -22,17 +22,29 @@ if [ ! -f "./secrets-export.sh" ]; then
 fi
 
 # Remove any AWS creds that might be set in the parent env.
-unset AWS_ACCESS_KEY_ID
-unset AWS_SECRET_ACCESS_KEY
-unset AWS_SESSION_TOKEN
+# We will need those creds for eks to set up the cluster.
+if [ $1 != "eks" ]; then
+    unset AWS_ACCESS_KEY_ID
+    unset AWS_SECRET_ACCESS_KEY
+    unset AWS_SESSION_TOKEN
+fi
 
 source ./secrets-export.sh
 
 if [ -f $SCRIPT_DIR/test-env.sh ]; then
     rm $SCRIPT_DIR/test-env.sh
 fi
 
+export PROJECT_DIRECTORY
 python aws_tester.py "$@"
+
+# Remove any AWS creds that might be set in the parent env.
+if [ $1 == "eks" ]; then
+    unset AWS_ACCESS_KEY_ID
+    unset AWS_SECRET_ACCESS_KEY
+    unset AWS_SESSION_TOKEN
+fi
+
 source $SCRIPT_DIR/test-env.sh
 
 popd

.evergreen/auth_aws/aws_tester.py

@@ -7,6 +7,7 @@
 import json
 import logging
 import os
+import random
 import subprocess
 import sys
 from functools import partial
@@ -26,9 +27,6 @@ def join(*parts):
 
 
 sys.path.insert(0, str(HERE / "lib"))
-from aws_assign_instance_profile import _assign_instance_policy
-from aws_assume_role import _assume_role
-from aws_assume_web_role import _assume_role_with_web_identity
 from util import get_key as _get_key
 
 ASSUMED_ROLE = "arn:aws:sts::557821124784:assumed-role/authtest_user_assume_role/*"
@@ -43,7 +41,7 @@ def join(*parts):
         CONFIG = json.load(fid)
         get_key = partial(_get_key, uppercase=False)
 except FileNotFoundError:
-    CONFIG = os.environ
+    CONFIG = os.environ.copy()
     get_key = partial(_get_key, uppercase=True)
 
 
@@ -72,6 +70,8 @@ def create_user(user, kwargs):
 
 
 def setup_assume_role():
+    from aws_assume_role import _assume_role
+
     # Assume the role to get temp creds.
     os.environ["AWS_ACCESS_KEY_ID"] = CONFIG[get_key("iam_auth_assume_aws_account")]
     os.environ["AWS_SECRET_ACCESS_KEY"] = CONFIG[
@@ -100,6 +100,8 @@ def setup_assume_role():
 
 def setup_ec2():
     # Create the user.
+    from aws_assign_instance_profile import _assign_instance_policy
+
     _assign_instance_policy()
     os.environ.pop("AWS_ACCESS_KEY_ID", None)
     os.environ.pop("AWS_SECRET_ACCESS_KEY", None)
@@ -179,6 +181,8 @@ def setup_env_creds():
 
 
 def setup_web_identity():
+    from aws_assume_web_role import _assume_role_with_web_identity
+
     # Unassign the instance profile.
     env = dict(
         AWS_ACCESS_KEY_ID=CONFIG[get_key("iam_auth_ec2_instance_account")],
@@ -235,6 +239,33 @@ def setup_web_identity():
     return dict(AWS_WEB_IDENTITY_TOKEN_FILE=token_file, AWS_ROLE_ARN=role_arn)
 
 
+def setup_eks_pod_identity():
+    if "PROJECT_DIRECTORY" not in os.environ:
+        raise ValueError("Please define a PROJECT_DIRECTORY")
+
+    test_path = (
+        Path(os.environ["PROJECT_DIRECTORY"])
+        / ".evergreen"
+        / "run-mongodb-aws-eks-test.sh"
+    )
+    if not test_path.exists():
+        raise ValueError(f"Please add the file {test_path}!")
+
+    # Set the name for the deployment.
+    name = f"mongodb-{random.randrange(0, 32767)}"
+    try:
+        subprocess.check_call(
+            ["bash", HERE / "lib" / "eks-pod-setup.sh", name], cwd=HERE / "lib"
+        )
+    finally:
+        # Tear down the EKS assets.
+        subprocess.check_call(
+            ["bash", HERE / "lib" / "eks-pod-teardown.sh", name], cwd=HERE / "lib"
+        )
+
+    return dict()
+
+
 def handle_creds(creds: dict):
     if "USER" in creds:
         USER = quote_plus(creds["USER"])
@@ -243,6 +274,8 @@ def handle_creds(creds: dict):
             MONGODB_URI = f"mongodb://{USER}:{PASS}@localhost"
         else:
             MONGODB_URI = f"mongodb://{USER}@localhost"
+    elif "MONGODB_URI" in creds:
+        MONGODB_URI = creds.pop("MONGODB_URI")
     else:
         MONGODB_URI = "mongodb://localhost"
     MONGODB_URI = f"{MONGODB_URI}/aws?authMechanism=MONGODB-AWS"
@@ -286,6 +319,9 @@ def main():
     run_web_identity_cmd = sub.add_parser("web-identity", help="Web identity test")
     run_web_identity_cmd.set_defaults(func=setup_web_identity)
 
+    run_eks_pod_identity_cmd = sub.add_parser("eks", help="EKS pod identity test")
+    run_eks_pod_identity_cmd.set_defaults(func=setup_eks_pod_identity)
+
     args = parser.parse_args()
     func_name = args.func.__name__.replace("setup_", "").replace("_", "-")
     LOGGER.info("Running aws_tester.py with %s...", func_name)

.evergreen/auth_aws/lib/eks-pod-run-self-test.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+set -eu
+
+echo "Installing dependencies ... begin"
+rm -rf .venv
+python3 -m venv .venv
+source .venv/bin/activate
+pip install -U -q pip "pymongo[aws]"
+echo "Installing dependencies ... end"
+
+# Run the Python Driver Self Test
+SCRIPT_DIR=$(realpath "$(dirname ${BASH_SOURCE[0]})")
+pushd $SCRIPT_DIR
+export MONGODB_URI=$1
+python test.py

.evergreen/auth_aws/lib/eks-pod-setup.sh

@@ -0,0 +1,120 @@
+#!/usr/bin/env bash
+set -eu
+
+# Write the secrets-export.sh file to the k8s/eks directory.
+EKS_DIR="../../k8s/eks"
+
+cat <<EOF >> $EKS_DIR/secrets-export.sh
+export EKS_CLUSTER_NAME=$EKS_CLUSTER_NAME
+export EKS_SERVICE_ACCOUNT_NAME=$EKS_SERVICE_ACCOUNT_NAME
+export EKS_REGION=$EKS_REGION
+EOF
+
+bash $EKS_DIR/setup.sh
+source $EKS_DIR/secrets-export.sh
+
+NAME="$1"
+MONGODB_URI="mongodb://${NAME}:27017"
+APP_LABEL=mongodb-deployment
+MONGODB_VERSION=${MONGODB_VERSION:-latest}
+
+. ../../ensure-binary.sh kubectl
+
+# Delete mongodb servers over one hour old in case they were not torn down.
+echo "Deleting old mongodb servers..."
+if [ "$(uname -s)" = "Darwin" ]; then
+    DATE="gdate"
+else
+    DATE="date"
+fi
+# shellcheck disable=SC2046
+kubectl get deployments -l app=$APP_LABEL -o go-template --template '{{range .items}}{{.metadata.name}} {{.metadata.creationTimestamp}}{{"\n"}}{{end}}' | awk '$2 <= "'$($DATE -d'now-1 hours' -Ins --utc | sed 's/+0000/Z/')'" { print $1 }' | xargs --no-run-if-empty kubectl delete deployment
+# shellcheck disable=SC2046
+kubectl get services -l app=$APP_LABEL -o go-template --template '{{range .items}}{{.metadata.name}} {{.metadata.creationTimestamp}}{{"\n"}}{{end}}' | awk '$2 <= "'$($DATE -d'now-1 hours' -Ins --utc | sed 's/+0000/Z/')'" { print $1 }' | xargs --no-run-if-empty kubectl delete service
+echo "Deleting old mongodb servers... done."
+
+cat <<EOF | kubectl apply -f -
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: ${NAME}
+  labels:
+    app: ${APP_LABEL}
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: ${NAME}
+  template:
+    metadata:
+      labels:
+        app: ${NAME}
+    spec:
+      containers:
+      - name: mongodb
+        image: mongodb/mongodb-enterprise-server:${MONGODB_VERSION}
+        ports:
+        - containerPort: 27017
+        env:
+        - name: MONGO_INITDB_ROOT_USERNAME
+          value: "bob"
+        - name: MONGO_INITDB_ROOT_PASSWORD
+          value: "pwd123"
+        - name: MONGODB_AWS_ACCOUNT_ARN
+          value: "${EKS_ROLE_ARN}"
+        args:
+        - "--setParameter"
+        - "authenticationMechanisms=MONGODB-AWS,SCRAM-SHA-256"
+EOF
+
+cat <<EOF | kubectl apply -f -
+apiVersion: v1
+kind: Service
+metadata:
+  name: ${NAME}
+  labels:
+    app: ${APP_LABEL}
+spec:
+  selector:
+    app: ${NAME}
+  ports:
+  - protocol: TCP
+    port: 27017
+    targetPort: 27017
+  type: ClusterIP
+EOF
+
+# Set up the server.
+echo "Setting up the server..."
+MONGODB_POD=$(kubectl get pods -l app=${NAME} -o jsonpath='{.items[0].metadata.name}')
+# Wait for the pod to be ready.
+kubectl wait --for=condition=Ready pod/${MONGODB_POD} --timeout=2000s
+kubectl exec ${MONGODB_POD} -- bash -c "rm -rf /tmp/test && mkdir /tmp/test"
+kubectl cp ./eks_pod_setup_server.js ${MONGODB_POD}:/tmp/test/setup_server.js
+kubectl exec ${MONGODB_POD} -- mongosh /tmp/test/setup_server.js
+echo "Setting up the server... done."
+
+# Run the self test.
+echo "Running self test on eks pod..."
+kubectl exec ${K8S_POD_NAME} -- bash -c "rm -rf /tmp/self-test && mkdir /tmp/self-test"
+kubectl cp ./eks-pod-run-self-test.sh ${K8S_POD_NAME}:/tmp/self-test/run-self-test.sh
+kubectl cp ./eks_pod_self_test.py ${K8S_POD_NAME}:/tmp/self-test/test.py
+kubectl exec ${K8S_POD_NAME} -- /tmp/self-test/run-self-test.sh $MONGODB_URI
+echo "Running self test on eks pod... done."
+
+# Set up driver test.
+echo "Setting up driver test files..."
+kubectl exec ${K8S_POD_NAME} -- bash -c "rm -rf /tmp/src"
+kubectl cp $PROJECT_DIRECTORY ${K8S_POD_NAME}:/tmp/src/
+echo "Setting up driver test files... done."
+
+echo "Running the driver test command... done."
+echo "export MONGODB_URI=${MONGODB_URI}" >> secrets-export.sh
+kubectl cp ./secrets-export.sh ${K8S_POD_NAME}:/tmp/src/secrets-export.sh
+echo "Setting up driver test files... done."
+
+# Run the driver test.
+echo "Running the driver test command..."
+MONGODB_URI="${MONGODB_URI}/aws?authMechanism=MONGODB-AWS"
+kubectl exec ${K8S_POD_NAME} -- bash -c "cd /tmp && source src/secrets-export.sh && bash src/.evergreen/run-mongodb-aws-eks-test.sh $MONGODB_URI"
+echo "Running the driver test command... done."

.evergreen/auth_aws/lib/eks-pod-teardown.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+set -eu
+
+EKS_APP_NAME=$1
+
+echo "Tearing down EKS assets..."
+. ../../ensure-binary.sh kubectl
+kubectl delete deployment $EKS_APP_NAME
+kubectl delete services $EKS_APP_NAME
+bash ../../k8s/eks/teardown.sh
+echo "Tearing down EKS assets... done."

.evergreen/auth_aws/lib/eks_pod_self_test.py

@@ -0,0 +1,12 @@
+import os
+
+from pymongo import MongoClient
+
+mongodb_uri = os.environ["MONGODB_URI"]
+
+print("Testing MONGODB-AWS on eks...")
+c = MongoClient(f"{mongodb_uri}/?authMechanism=MONGODB-AWS")
+c.aws.test.find_one({})
+c.close()
+print("Testing MONGODB-AWS on eks... done.")
+print("Self test complete!")

.evergreen/auth_aws/lib/eks_pod_setup_server.js

@@ -0,0 +1,10 @@
+(function() {
+"use strict";
+
+const AWS_ACCOUNT_ARN=process.env["MONGODB_AWS_ACCOUNT_ARN"];
+const admin = Mongo().getDB("admin");
+const external = admin.getMongo().getDB("$external");
+assert(admin.auth("bob", "pwd123"));
+
+external.runCommand({createUser: AWS_ACCOUNT_ARN, roles:[{role: 'read', db: "aws"}]});
+}());