From 7e659b23c90f792ddc3aeddddce330aeefea1f40 Mon Sep 17 00:00:00 2001
From: Chai Feng <chaifeng@chaifeng.com>
Date: Fri, 5 Oct 2018 16:26:07 +0800
Subject: [PATCH] update entrypoint: update or deny service rule

---
 docker-entrypoint.sh | 30 ++++++++++--------------------
 1 file changed, 10 insertions(+), 20 deletions(-)

diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh
index 2427c00..7578f85 100755
--- a/docker-entrypoint.sh
+++ b/docker-entrypoint.sh
@@ -6,26 +6,15 @@ set -euo pipefail
 ufw_docker_agent=ufw-docker-agent
 ufw_docker_agent_image="${ufw_docker_agent_image:-chaifeng/${ufw_docker_agent}:181003}"
 
-function ufw-update-rule-for-instance() {
-    name="$1"
-    port="$2"
+function ufw-allow-or-deny-service() {
+    declare id="$1"
+    declare port="$2"
 
-    declare -a opts
-    [[ "$port" = deny ]] && opts+=(delete)
-    opts+=(allow)
-
-    [[ "$port" = @(all|deny) ]] && port=""
-
-    run-ufw-docker "${opts[@]}" "${name}" "$port"
-}
-function ufw-update-service-instances() {
-    id="$1"
-    port="$2"
-
-    docker ps -qf "label=com.docker.swarm.service.id=${id}" |
-        while read name; do
-            ufw-update-rule-for-instance "${name}" "$port"
-        done
+    if [[ "$port" = deny ]]; then
+        run-ufw-docker delete allow "$id"
+    else
+        run-ufw-docker add-service-rule "$id" "$port"
+    fi
 }
 
 function update-ufw-rules() {
@@ -34,13 +23,14 @@ function update-ufw-rules() {
                      -e 's/="/ /' \
                      -e 's/"$//' |
         while read id port; do
-            ufw-update-service-instances "${id}" "${port}"
+            ufw-allow-or-deny-service "${id}" "${port}"
         done
 }
 
 function run-ufw-docker() {
     declare -a docker_opts=(run --rm -t --name ufw-docker-agent-"${RANDOM}"-$(date '+%Y%m%d%H%M%S')
          --cap-add NET_ADMIN --network host
+         --env DEBUG="$DEBUG"
          --env UFW_DOCKER_FORCE_ADD=yes
          -v /var/run/docker.sock:/var/run/docker.sock
          -v /etc/ufw:/etc/ufw "${ufw_docker_agent_image}" "$@")