latency-stress.yml

---
#==================================================
# SSM Run Command Document:
# Defines the configuration and the steps to
# run by SSM Run Command on Ec2 Instances
#==================================================
description: |
  ## What does this document do?
  It adds latency to the network traffic using tc. If tc is not already installed on the instance, this SSM document will install it together with other dependencies listed below.
  This SSM document supports Amazon Linux and Ubuntu operating systems only.

  ## Dependencies installed by this SSM Document
  * tc
  * atd
  Note: This SSM document does not uninstall dependencies once installed. If you don't want this SSM document to install any dependencies, set InstallDependencies to False and bake the dependencies into your EC2 instance. For example, by using image-builder (https://aws.amazon.com/image-builder/).

  ## Input Parameters
  * Interface: The interface to add delay to (default: eth0).
  * DelayMilliseconds: The delay - in milliseconds - to add to the interface (default: 200).
  * DurationSeconds: (Required) The duration - in seconds - of the test.
  * InstallDependencies: If set to True, Systems Manager installs the required dependencies on the target instances. (default: True).

schemaVersion: '2.2'
parameters:
  Interface:
    type: String
    description: 'The interface to add delay to (default: eth0).'
    default: eth0
    allowedPattern: ^[0-9a-zA-Z\-]{1,15}$
  DelayMilliseconds:
    type: String
    description: 'The delay - in milliseconds - to add to the interface (default:
      200).'
    default: '200'
    allowedPattern: ^[0-9]+$
  DurationSeconds:
    type: String
    description: (Required) The duration - in seconds - of the fault injection.
    allowedPattern: ^[0-9]+$
  InstallDependencies:
    type: String
    description: 'If set to True, Systems Manager installs the required dependencies
      on the target instances. (default: True).'
    default: 'True'
    allowedValues:
    - 'True'
    - 'False'
mainSteps:
- action: aws:runShellScript
  name: InstallDependencies
  precondition:
    StringEquals:
    - platformType
    - Linux
  description: |
    ## Parameter: InstallDependencies
    If set to True, this step installs the required dependecy via operating system's repository. It supports both
    Debian (apt) and CentOS (yum) based package managers.
  inputs:
    onFailure: exit
    runCommand:
    - |
      #!/bin/bash

      if [[ "$( which tc 2>/dev/null )" && "$( which atd 2>/dev/null )" ]] ; then echo Dependency is already installed. ; exit ; fi
      if  [[ "{{ InstallDependencies }}" == True ]] ; then
        echo "Installing required dependencies"
        if [ -f  "/etc/system-release" ] ; then
          if cat /etc/system-release | grep -i 'Amazon Linux release 2023' ; then
            sudo dnf -y install iproute-tc jq at
            sudo dnf -y install bind-utils
          elif cat /etc/system-release | grep -i 'Amazon Linux' ; then
            sudo amazon-linux-extras install testing
            sudo yum -y install tc at
          else
            echo "There was a problem installing dependencies."
            exit 1
          fi
        elif cat /etc/issue | grep -i Ubuntu ; then
          sudo apt-get update -y
          sudo DEBIAN_FRONTEND=noninteractive sudo apt-get install -y iproute2 at
        else
          echo "Exiting - This SSM document supports Amazon Linux and Ubuntu operating systems only."
          exit 1
        fi
      else
        echo "Dependencies are not installed - Please set InstallDependencies to True."
        exit 1
      fi
- action: aws:runShellScript
  name: FaultInjection
  precondition:
    StringEquals:
    - platformType
    - Linux
  description: |
    ## Parameters: Interface, DelayMilliseconds and DurationSeconds
    This step adds a `DelayMilliseconds` delay to `Interface` for the given `DurationSeconds`, using the `tc` (Traffic Control) command.
    The script will inject latency on the network, and wait for the given duration to remove that. It has two rollback mechanisms in place:
    * It will listen for exit signals (SIGINT and SIGTERM), and will stop the latency injection if any of them is received.
    * It will periodically enqueue rollback checks into a queue (using `at` command). This way, if the command is stopped and the rollback
    was not executed, the enqueued commands will try to stop it. (for example, if the command is stopped using kill -9). This is
    a safety check to avoid latency to remain injected after the script is no longer running.
  inputs:
    maxAttempts: 1
    runCommand:
    - |
      #!/bin/bash

      INTERFACE={{ Interface }}
      DELAYMILLISECONDS={{ DelayMilliseconds }}
      DURATION={{ DurationSeconds }}

      if ! [[ "$( ip a ls $INTERFACE 2>/dev/null )" ]] ; then
        echo "Interface $INTERFACE does not exist."
        exit 1
      fi

      if [ $DURATION -lt 1 ] || [ $DURATION -gt 43200 ]; then
        echo "Parameter DurationSeconds must be between 1 and 43200, was: '$DURATION'"
        exit 1
      fi

      if [ $DELAYMILLISECONDS -lt 1 ] ; then
        echo "Parameter Delay must be a positive value"
        exit 1
      fi

      FAULT_NAME="Run-Network-Latency"

      MAX_FLAG_AGE_SECONDS=5
      ATTEMPT_ROLLBACK_AT_SECONDS=10
      STOP_TIME=$(( $(date +%s) + $DURATION ))

      IP_CIDR_REGEX="^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\/(3[0-2]|[1-2][0-9]|[0-9]))$"
      IP_REGEX="^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])$"
      IMDS_IP=169.254.169.254

      # Get the REGION of the instance
      EC2_REGION=$(
        TOKEN=`curl -s -X PUT "http://$IMDS_IP/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30"` &&
        curl -s -H "X-aws-ec2-metadata-token: $TOKEN" http://$IMDS_IP/latest/dynamic/instance-identity/document |
        grep region | awk -F\" '{print $4}'
      )

      IP_ADDRESSES_SSM_ENDPOINTS=()

      DELETE_MSG="#delete-after-fault-injection"

      ENDPOINTS=(
          "ssmmessages.$EC2_REGION.amazonaws.com"
          "ssm.$EC2_REGION.amazonaws.com"
          "ec2messages.$EC2_REGION.amazonaws.com"
          $IMDS_IP
      )

      force_dns()
      {
          # Force any new DNS resolution for that host to the latest known IPs
          params=("$@")
          hostname=${params[0]}   # Hostname is expected as first argument
          ips=("${params[@]:1}")  # The rest will be the IPs
          for ip in $(echo "${ips[@]}"); do
            echo "$ip $hostname $DELETE_MSG" >> /etc/hosts
          done
      }

      get_ips()
      {
          # Returns if the input is a valid IP/CIDR string, empty otherwise
          local var_ip_cidr=$(grep -E "$IP_CIDR_REGEX" <<< "$1")
          [  -z "$var_ip_cidr" ] || { echo $var_ip_cidr ; return ;}

          # Returns if the input is a valid IP string, empty otherwise
          local var_ip=$(grep -E "$IP_REGEX" <<< "$1")
          [  -z "$var_ip" ] || { echo $var_ip ; return ;}

          # Returns IPs of the domain and force DNS resolution to that IP, empty otherwise
          # Since `dig` can return different IPs, we call it 10 times to maximize the chances of covering as many IPs as possible for the provided domain
          local var_dns=()
          for i in $(seq 1 10)
          do
              dig_output=( $(dig +short $1 | grep -v '[[:alpha:]]') )
              var_dns=("${var_dns[@]}" "${dig_output[@]}")
          done
          var_unique_dns=( $(printf '%s\n' "${var_dns[@]}" | sort -u) )
          [  -z "$var_unique_dns" ] || { echo "${var_unique_dns[@]}" ; force_dns $1 "${var_unique_dns[@]}" ; return ;}
      }

      # Getting IPs from SSM endpoints
      for e in ${ENDPOINTS[*]}; do
          var_endpoint=$(get_ips $e)
          IP_ADDRESSES_SSM_ENDPOINTS=("${IP_ADDRESSES_SSM_ENDPOINTS[@]}" "${var_endpoint[@]}")
      done

      # Adds Latency to the network interface
      read -r -d '' ENABLE_FAULT_COMMAND <<'EOF'
        # Exit if FIS network fault is already running
        test_file_exit() {
            if [ "$(ls "$1" 2>/dev/null | wc -l)" -ge "1" ]; then { echo "Fault might be already running (Found flag file matching "$1"). Exiting..." 1>&2 ; exit 1; } ; fi;
        }
        test_file_exit /var/lib/amazon/ssm/Run-Network-*.flag

        echo "Injecting fault..."
        # Send all traffic by default to the band 1:3
        tc qdisc add dev $INTERFACE root handle 1: prio priomap 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
        for k in ${IP_ADDRESSES_SSM_ENDPOINTS[*]}; do
            # Redirect the matching ip addresses to the band 0 where is no filters
            tc filter add dev $INTERFACE protocol ip parent 1:0 prio 1 u32 match ip dst $k flowid 1:1
        done
        # Creates the qdisc in the band 3, affecting all the rest of the traffic
        tc qdisc add dev $INTERFACE parent 1:3 handle 10: netem delay "${DELAYMILLISECONDS}"ms
      EOF

      # Removes Packet Loss from the network interface
      read -r -d '' DISABLE_FAULT_COMMAND <<EOF
        echo "Rolling back..."
        tc filter del dev $INTERFACE prio 1
        tc qdisc del dev $INTERFACE parent 1:3 handle 10:
        tc qdisc del dev $INTERFACE root handle 1: prio
        sed -ie "/$DELETE_MSG/d" /etc/hosts
      EOF

      STOP_TIME=$(( $(date +%s) + $DURATION ))

      MAX_FLAG_AGE_SECONDS=5
      ATTEMPT_ROLLBACK_AT_SECONDS=10

      RANDOM_STRING=$(cat /dev/urandom | tr -dc 'a-zA-Z0-9' | head -c 32)
      FLAG_PATH="/var/lib/amazon/ssm/$FAULT_NAME-$RANDOM_STRING.flag"
      ROLLBACK_PATH="/var/lib/amazon/ssm/$FAULT_NAME-$RANDOM_STRING-Rollback.sh"

      # Creating a file with rollback check command to be executed by atd
      cat << EOF > "$ROLLBACK_PATH"
      #!/bin/bash
      # Dont sleep if run in terminal
      if ! [ -t 0 ] ; then
        sleep $ATTEMPT_ROLLBACK_AT_SECONDS
      fi
      if ! [ -f "$FLAG_PATH" ] || [ "\$(( \$(date +%s) - \$(stat -c "%Y" "$FLAG_PATH") ))" -gt $MAX_FLAG_AGE_SECONDS ] || [ -t 0 ]; then
        $DISABLE_FAULT_COMMAND
        rm -f "$FLAG_PATH"
        # Dont delete rollback script if run in terminal
        if ! [ -t 0 ] ; then
          rm -f "$ROLLBACK_PATH"
        fi
      fi
      EOF
      echo "Temporary rollback file created: $ROLLBACK_PATH"

      # Enqueue a rollback check after $ATTEMPT_ROLLBACK_AT_SECONDS seconds
      schedule_rollback_attempt() {
        echo "bash $ROLLBACK_PATH" | at now 2> >(sed '/warning: commands will be executed using \/bin\/sh/d' )
      }

      # Function to delete the flag file and rollback the fault injection
      rollback() {
        rm "$FLAG_PATH"
        rm "$ROLLBACK_PATH"
        eval "$DISABLE_FAULT_COMMAND" ; STATUS=$?
        echo Rollback done.
        exit $STATUS
      }

      # Binding the rollback function to these exit signals
      trap rollback INT
      trap rollback TERM

      echo "Making sure atd daemon is running"
      # atd must be running in order to use "at" later
      atd || { echo Failed to run atd daemon, exiting... 1>&2 ; exit 1; }

      echo "Scheduling rollback"
      schedule_rollback_attempt

      # Injecting fault
      echo "Enabling fault injection"
      eval "$ENABLE_FAULT_COMMAND"

      # For the duration of the injection, the flag file is updated, and a rollback check is enqueued
      while [[ $(date +%s) -lt $STOP_TIME ]] ; do
        touch "$FLAG_PATH"
        schedule_rollback_attempt
        sleep $MAX_FLAG_AGE_SECONDS
      done

      # After the desired duration, the fault injection is removed
      rollback