#!/bin/bash

VERSION=4.3.0
IPTABLES_OPTS=()

# Ingress Routing Daemon v4.3.0
# Copyright © 2020-2025 Struan Bartlett
# ----------------------------------------------------------------------
# Permission is hereby granted, free of charge, to any person 
# obtaining a copy of this software and associated documentation files 
# (the "Software"), to deal in the Software without restriction, 
# including without limitation the rights to use, copy, modify, merge, 
# publish, distribute, sublicense, and/or sell copies of the Software, 
# and to permit persons to whom the Software is furnished to do so, 
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be 
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS 
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN 
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 
# SOFTWARE.
# ----------------------------------------------------------------------
# Workaround for https://github.com/moby/moby/issues/25526

log() {
  if [ -n "${JOURNAL_STREAM}" ]; then
    # Launched with systemd and journald
    echo "$@"
    return
  fi

  [ -z "$_BASHPID" ] && _BASHPID="$BASHPID"
  local D=$(date +%Y-%m-%d.%H:%M:%S.%N)
  local S=$(printf "%s|%s|%05d|" "${D:0:26}" "$HOSTNAME" "$_BASHPID")
  echo "$@" | sed "s/^/$S /g"
}			

quit() {
  trap '' EXIT TERM INT
  log "Docker Ingress Routing Daemon received signal, propagating it to pgroup $$ ..."
  kill -TERM -$$
  log "Docker Ingress Routing Daemon exiting."
  exit 0
}

detect_ingress() {
  read INGRESS_SUBNET INGRESS_DEFAULT_GATEWAY \
    < <(docker network inspect ingress --format '{{(index .IPAM.Config 0).Subnet}} {{index (split (index .Containers "ingress-sbox").IPv4Address "/") 0}}' 2>/dev/null)

  [ -n "$INGRESS_SUBNET" ] && [ -n "$INGRESS_DEFAULT_GATEWAY" ] && nsenter --net=/var/run/docker/netns/ingress_sbox iptables -L >/dev/null && return 0
  
  return 1
}

route_ingress() {
  local ID="$1"
  local SERVICE="$2"

  if [ -z "$SERVICE" ]; then
    return
  fi

  if [ ${#SERVICES[@]} -gt 0 ] && ! [[ " ${SERVICES[@]} " =~ " $SERVICE " ]]; then
    log "Detected container launch for service '$SERVICE', with ID '$ID': unmatched service, so skipping."
    return
  fi

  local NID=$(docker inspect -f '{{.State.Pid}}' $ID)
  local CIF=$(nsenter -n -t $NID ip -brief addr show to $INGRESS_SUBNET | cut -d'@' -f1)

  if [ -z "$CIF" ]; then
    log "Detected container launch for service '$SERVICE', with ID '$ID' and NID '$NID': no ingress network interface found, so skipping."
    return
  fi

  local IPTABLE_COMMENT="docker-ingress-routing-daemon"
  if nsenter -n -t $NID iptables "${IPTABLES_OPTS[@]}" -t mangle -C OUTPUT -m comment --comment "$IPTABLE_COMMENT" 2>/dev/null; then
    log "Detected container for service '$SERVICE', with ID '$ID' and NID '$NID': mangle table already configured, so skipping."
    return
  fi
    
  log "Detected container launch for service '$SERVICE', with ID '$ID' and NID '$NID': ingress network interface $CIF found, so applying policy routing/firewall rules:"
   
  # 3. Map any connection mark on outgoing tcp or udp traffic to a firewall mark on the individual packets.
  #    These rules /could potentially/ be applied more selectively, according to --tcp-ports and --udp-ports, to make
  #    a marginal efficiency gain, but this is not necessary: as, if no connection mark has been set, because no
  #    TOS byte has been set by the load balancer, then none will be restored and legacy routing rules will apply.
  #    - See https://github.com/newsnowlabs/docker-ingress-routing-daemon/issues/11
  log "- Adding container mangle table iptables rules"
  nsenter -n -t $NID iptables "${IPTABLES_OPTS[@]}" -t mangle -A OUTPUT -m comment --comment "$IPTABLE_COMMENT"
  nsenter -n -t $NID iptables "${IPTABLES_OPTS[@]}" -t mangle -A OUTPUT -p udp -j CONNMARK --restore-mark
  nsenter -n -t $NID iptables "${IPTABLES_OPTS[@]}" -t mangle -A OUTPUT -p tcp -j CONNMARK --restore-mark

  # 3.1 Enable 'loose' rp_filter mode on interface $CIF (and 'all' as required by kernel
  #     see https://www.kernel.org/doc/Documentation/networking/ip-sysctl.txt)
  log "- Setting container sysctl net.ipv4.conf.all.rp_filter=2 net.ipv4.conf.$CIF.rp_filter=2"
  nsenter -n -t $NID sysctl -q net.ipv4.conf.all.rp_filter=2 net.ipv4.conf.$CIF.rp_filter=2

  local INDEX=1 NODE_IP
  for NODE_IP in ${INGRESS_NODE_GATEWAY_IPS[@]}
  do
    if [ "$NODE_IP" != "null" ]; then
      if [ -n "$INDEXED_IDS" ]; then
        NODE_ID=$INDEX
      else
        NODE_ID=$(echo $NODE_IP | cut -d'.' -f4)
      fi

      log "- Adding container policy routing/firewall rules for load-balancer #$NODE_ID with IP $NODE_IP"

      # 2. Map the TOS value on any incoming packets to a connection mark, using the same value.
      nsenter -n -t $NID iptables "${IPTABLES_OPTS[@]}" -t mangle -A PREROUTING -m tos --tos $NODE_ID/0xff -j CONNMARK --set-xmark $NODE_ID/0xffffffff

      # 4. Select the correct routing table to use, according to the firewall mark on the outgoing packet.
      nsenter -n -t $NID ip rule add from $INGRESS_SUBNET fwmark $NODE_ID lookup $NODE_ID prio 32700

      # 5. Route outgoing traffic to the correct node's ingress network IP, according to its firewall mark
      #    (which in turn came from its connection mark, its TOS value, and ultimately its IP).
      nsenter -n -t $NID ip route add table $NODE_ID default via $NODE_IP
    fi
    ((INDEX+=1))
  done
  log "- Finished configuring launched container"
}

usage() {
  echo "Docker Ingress Routing Daemon v$VERSION" >&2
  echo >&2
  echo "Usage: $0 [--install [OPTIONS] | --uninstall | --help]" >&2
  echo >&2
  echo "         --services <services>   - service names to whitelist (i.e. disable masquerading for)" >&2
  echo "           --tcp-ports <ports>   - TCP ports to whitelist (i.e. disable masquerading for)" >&2
  echo "           --udp-ports <ports>   - UDP ports to whitelist (i.e. disable masquerading for)" >&2
  echo "   --ingress-gateway-ips <ips>   - specify load-balancer ingress IPs" >&2
  echo "                 --preexisting   - optionally install rules where needed" >&2
  echo "                                   on preexisting containers (recommended)" >&2
  echo >&2
  echo "               --iptables-wait   - pass '--wait' option to iptables" >&2
  echo "   --iptables-wait-seconds <n>   - pass '--wait <n>' option to iptables" >&2
  echo >&2
  echo "              --no-performance   - disable performance optimisations" >&2
  echo "                 --indexed-ids   - use sequential ids for load balancers" >&2
  echo "                                   (forced where ingress subnet larger than /24)" >&2
  echo >&2
  echo "Notes:" >&2
  echo "- <services>, <ports> and <ips> may be comma-separated, space-separated," >&2
  echo "  or the relevant option specified multiple times;" >&2
  echo "- an IP of 'null' will be skipped, allowing old IPs to be retired without" >&2
  echo "  changing the indexes of serving IPs." >&2
  echo >&2
  
  if detect_ingress; then
    echo "Detected ingress subnet and node IP:" >&2
    echo "- Ingress subnet: $INGRESS_SUBNET" >&2
    echo "- This node's ingress network IP: $INGRESS_DEFAULT_GATEWAY" >&2
  fi
  
  echo >&2
  exit 1
}

SCRIPT_PATH=$(dirname $(realpath $0))
DOCKER=$(which docker)
INGRESS_NODE_GATEWAY_IPS=($INGRESS_NODE_GATEWAY_IPS) # Convert legacy env var to new array

while true
do
  case "$1" in
              --services) shift; SERVICES+=($(echo "$1" | tr ',' ' ')); shift; continue; ;;
             --tcp-ports) shift; TCP_PORTS+=($(echo "$1" | tr ',' ' ')); shift; continue; ;;
             --udp-ports) shift; UDP_PORTS+=($(echo "$1" | tr ',' ' ')); shift; continue; ;;
               --install) shift; INSTALL=1; continue; ;;
             --uninstall) shift; INSTALL=0; continue; ;;
   --ingress-gateway-ips) shift; INGRESS_NODE_GATEWAY_IPS+=($(echo "$1" | tr ',' ' ')); shift; continue; ;;
        --no-performance) shift; PERFORMANCE=0; continue; ;;
           --indexed-ids) shift; INDEXED_IDS=1; continue; ;;
           --preexisting) shift; PREEXISTING=1; continue; ;;
         --iptables-wait) shift; IPTABLES_WAIT=1; continue; ;;
 --iptables-wait-seconds) shift; IPTABLES_WAIT=1; IPTABLES_WAIT_SECONDS="$1"; shift; continue; ;;

               -h|--help) usage; ;;
                      '') break; ;;

                       *) usage; break; ;;
  esac
done

# Display usage, unless --install or --uninstall
[ -z "$INSTALL" ] && usage

# Add '--iptables-wait' or '--iptables-wait $IPTABLES_WAIT_SECONDS' to iptables options
[ -n "$IPTABLES_WAIT" ] && IPTABLES_OPTS+=(--wait $IPTABLES_WAIT_SECONDS)

# Convert arrays to comma-separated strings
TCPServicePortString=$(echo ${TCP_PORTS[@]} | tr ' ' ',')
UDPServicePortString=$(echo ${UDP_PORTS[@]} | tr ' ' ',')

# Prevent "WARNING: Error loading config file: .dockercfg: $HOME is not defined" messages
export HOME=$SCRIPT_PATH

if ! [ -x "$DOCKER" ]; then
  echo "Docker binary not found; exiting." >&2
  exit -1
fi

log "Docker Ingress Routing Daemon $VERSION starting ..."

if detect_ingress; then
  log "Detecting ingress network and node IP:"
  log "- Ingress subnet: $INGRESS_SUBNET"
  log "- This node's IP: $INGRESS_DEFAULT_GATEWAY"
else
  log "- Couldn't identify ingress network subnet or this node's ingress network IP; sleeping 1s, then exiting."
  sleep 1
  exit -1
fi

# Delete any relevant preexisting rules.
log "Cleaning up any stale load-balancer rules ..."

nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t nat -S | \
  grep -- '-m ipvs --ipvs -j ACCEPT' | \
  sed -r 's/^-A /-D /' | \
  while read RULE; \
  do
    log "- Deleting old rule: iptables "${IPTABLES_OPTS[@]}" -t nat $RULE"
    nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t nat $RULE
  done

nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t mangle -S | \
  grep -- '-j TOS --set-tos' | \
  sed -r 's/^-A /-D /' | \
  while read RULE; \
  do
    log "- Deleting old rule: iptables "${IPTABLES_OPTS[@]}" -t mangle $RULE"
    nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t mangle $RULE
  done

nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t raw -S | \
  grep -- '-j CT --notrack' | \
  sed -r 's/^-A /-D /' | \
  while read RULE; \
  do
    log "- Deleting old rule: iptables "${IPTABLES_OPTS[@]}" -t raw $RULE"
    nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t raw $RULE
  done

if [ "$INSTALL" = "0" ]; then
  log "Docker Ingress Routing Daemon iptables rules uninstalled, exiting."
  exit 0
fi

###############
# $INSTALL is 1
#

if [ -z "${INGRESS_NODE_GATEWAY_IPS[*]}" ]; then
  INGRESS_NET=$(echo $INGRESS_DEFAULT_GATEWAY | cut -d'.' -f1,2,3)
  INGRESS_NODE_GATEWAY_IPS=($INGRESS_NET.2 $INGRESS_NET.3 $INGRESS_NET.4 $INGRESS_NET.5 $INGRESS_NET.6 $INGRESS_NET.7 $INGRESS_NET.8 $INGRESS_NET.9)
  
  log "!!! -------------------------- WARNING ------------------------------------"
  log "!!! Assuming --ingress-gateway-ips ${INGRESS_NODE_GATEWAY_IPS[*]}"
  log "!!!"
  log "!!! Please compile a list of the ingress network IPs of each of your nodes"
  log "!!! that you will be using as a load-balancer."
  log "!!!"
  log "!!! You only have to do this once, or whenever you change your set of"
  log "!!! load-balancer nodes."
  log "!!!"
  log "!!! Then relaunch using:"
  log "!!! $0 --install --ingress-gateway-ips \"<Node Ingress IP List>\""
  log "!!! ----------------------------------------------------------------------"

fi

log "Enumerating load balancers from --ingress-gateway-ips ${INGRESS_NODE_GATEWAY_IPS[*]}"

# Force --indexed-ids where the ingress subnet is larger than /24
if [ $(echo $INGRESS_SUBNET | cut -d'/' -f2) -lt 24 ]; then
  log "- Ingress subnet is larger than /24, so enabling --indexed-ids"
  INDEXED_IDS=1
fi

# Look for $INGRESS_DEFAULT_GATEWAY in ${INGRESS_NODE_GATEWAY_IPS[@]}:
# - if we find it, set NODE_ID to the index+1
# - if we don't, leave NODE_ID unset
#
unset NODE_ID

INDEX=1
for NODE_IP in ${INGRESS_NODE_GATEWAY_IPS[@]}
do

  if [ -n "$INDEXED_IDS" ]; then
    ID=$INDEX
  else
    ID=$(echo $NODE_IP | cut -d'.' -f4)
  fi
    
  if [ $NODE_IP == $INGRESS_DEFAULT_GATEWAY ]; then
    log "- Load balancer $NODE_IP will have ID $ID <=== THIS NODE"
    NODE_ID=$ID
  elif [ $NODE_IP != "null" ]; then
    log "- Load balancer $NODE_IP will have ID $ID"
  fi
  
  ((INDEX+=1))
done

if [ -n "$NODE_ID" ]; then
  log "Installing ingress namespace iptables rules for this load balancer (ID $NODE_ID):"

  # Add a rule ahead of the ingress network SNAT rule, that will cause the SNAT rule to be skipped.
  if [ -z "$TCPServicePortString" ] && [ -z "$UDPServicePortString" ]; then
    log "- Adding ingress_sbox iptables nat rule: iptables "${IPTABLES_OPTS[@]}" -t nat -I POSTROUTING -d $INGRESS_SUBNET -m ipvs --ipvs -j ACCEPT"
    nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t nat -I POSTROUTING -d $INGRESS_SUBNET -m ipvs --ipvs -j ACCEPT

    # 1. Set TOS to NODE_ID in all outgoing packets to INGRESS_SUBNET
    log "- Adding ingress_sbox iptables mangle rule: iptables "${IPTABLES_OPTS[@]}" -t mangle -A POSTROUTING -d $INGRESS_SUBNET -j TOS --set-tos $NODE_ID/0xff"
    nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t mangle -A POSTROUTING -d $INGRESS_SUBNET -j TOS --set-tos $NODE_ID/0xff

    log "- Adding ingress_sbox connection tracking disable rule: iptables "${IPTABLES_OPTS[@]}" -t raw -I PREROUTING -j CT --notrack"
    nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t raw -I PREROUTING -j CT --notrack
  else

    if [ -n "$TCPServicePortString" ]; then
      log "- Adding ingress_sbox iptables nat rule: iptables "${IPTABLES_OPTS[@]}" -t nat -I POSTROUTING -d $INGRESS_SUBNET -p tcp -m multiport --dports $TCPServicePortString -m ipvs --ipvs -j ACCEPT"
      nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t nat -I POSTROUTING -d $INGRESS_SUBNET -p tcp -m multiport --dports $TCPServicePortString -m ipvs --ipvs -j ACCEPT

      # 1. Set TOS to NODE_ID in all outgoing packets to INGRESS_SUBNET
      log "- Adding ingress_sbox iptables mangle rule: iptables "${IPTABLES_OPTS[@]}" -t mangle -A POSTROUTING -d $INGRESS_SUBNET -p tcp -m multiport --dports $TCPServicePortString -j TOS --set-tos $NODE_ID/0xff"
      nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t mangle -A POSTROUTING -d $INGRESS_SUBNET -p tcp -m multiport --dports $TCPServicePortString -j TOS --set-tos $NODE_ID/0xff

      log "- Adding ingress_sbox connection tracking disable rule: iptables "${IPTABLES_OPTS[@]}" -t raw -I PREROUTING -p tcp -m multiport --dports $TCPServicePortString -j CT --notrack"
      nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t raw -I PREROUTING -p tcp -m multiport --dports $TCPServicePortString -j CT --notrack
    fi

    if [ -n "$UDPServicePortString" ]; then
      log "- Adding ingress_sbox iptables nat rule: iptables "${IPTABLES_OPTS[@]}" -t nat -I POSTROUTING -d $INGRESS_SUBNET -p udp -m multiport --dports $UDPServicePortString -m ipvs --ipvs -j ACCEPT"
      nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t nat -I POSTROUTING -d $INGRESS_SUBNET -p udp -m multiport --dports $UDPServicePortString -m ipvs --ipvs -j ACCEPT

      # 1. Set TOS to NODE_ID in all outgoing packets to INGRESS_SUBNET
      log "- Adding ingress_sbox iptables mangle rule: iptables "${IPTABLES_OPTS[@]}" -t mangle -A POSTROUTING -d $INGRESS_SUBNET -p udp -m multiport --dports $UDPServicePortString -j TOS --set-tos $NODE_ID/0xff"
      nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t mangle -A POSTROUTING -d $INGRESS_SUBNET -p udp -m multiport --dports $UDPServicePortString -j TOS --set-tos $NODE_ID/0xff

      log "- Adding ingress_sbox connection tracking disable rule: iptables -p udp -m multiport --dports $UDPServicePortString -j CT --notrack"
      nsenter --net=/var/run/docker/netns/ingress_sbox iptables "${IPTABLES_OPTS[@]}" -t raw -I PREROUTING -p udp -m multiport --dports $UDPServicePortString -j CT --notrack
    fi

  fi
  
else
  log "This node is not a specified load balancer; so skipping installing ingress namespace iptables rules"
  
fi

if [ "$PERFORMANCE" != "0" ]; then
  # Set sysctl variables
  log "Setting ingress_sbox namespace sysctl variables:"
  log "- Setting net.ipv4.vs.conn_reuse_mode=0 net.ipv4.vs.expire_nodest_conn=1 net.ipv4.vs.expire_quiescent_template=1"
  nsenter --net=/var/run/docker/netns/ingress_sbox sysctl -q net.ipv4.vs.conn_reuse_mode=0 net.ipv4.vs.expire_nodest_conn=1 net.ipv4.vs.expire_quiescent_template=1
  
  [ -f "/etc/sysctl.d/conntrack.conf" ] && \
    log "- Setting $(cat /etc/sysctl.d/conntrack.conf | tr '\012' ' ') <=== from /etc/sysctl.d/conntrack.conf" && \
    nsenter --net=/var/run/docker/netns/ingress_sbox sysctl -q --load=/etc/sysctl.d/conntrack.conf
  
  [ -f "/etc/sysctl.d/ipvs.conf" ] && \
    log "- Setting $(cat /etc/sysctl.d/ipvs.conf | tr '\012' ' ') <=== from /etc/sysctl.d/ipvs.conf" && \
    nsenter --net=/var/run/docker/netns/ingress_sbox sysctl -q --load=/etc/sysctl.d/ipvs.conf
fi

# Turn off job control
set +m

# Set lastpipe, so that the 'while read' runs in the main shell process,
# making the script more resilient to subprocess exit.
shopt -s lastpipe

trap quit EXIT TERM INT

if [ "$PREEXISTING" = "1" ]; then
  log "Checking for preexisting containers that need routing rules added ..."
  docker ps \
    --format '{{.ID}} {{.Label "com.docker.swarm.service.name"}}' \
    --filter 'status=running' \
    --filter "label=com.docker.swarm.service.name" | \
    while read ID SERVICE
    do
      route_ingress "$ID" "$SERVICE"
    done
fi
  
log "Launching docker event watcher to monitor for container launches (pgroup $$) ..."

# Watch for container start events, and configure policy routing rules on each container
# to ensure return path traffic for incoming connections is routed back via the correct interface
# and to the correct node from which the incoming connection was received.
docker events \
  --format '{{.Actor.ID}} {{index .Actor.Attributes "com.docker.swarm.service.name"}}' \
  --filter 'event=start' \
  --filter 'type=container' | \
  while read ID SERVICE
  do
    route_ingress "$ID" "$SERVICE"
  done
