fix(zookeeper): fail closed on uncertain role probe

addons/zookeeper/scripts-ut-spec/roleprobe_spec.sh

@@ -122,5 +122,117 @@ Describe "ZooKeeper Startup Bash Script Tests"
       When call get_zk_role
       The output should eq "follower"
     End
+
+    It "returns observer when mode is observer"
+      get_zookeeper_mode() {
+        echo "observer"
+      }
+
+      When call get_zk_role
+      The status should be success
+      The output should eq "observer"
+    End
+  End
+
+  Describe "roleprobe.sh behavior"
+    setup_mock_nc() {
+      roleprobe_mock_bin="$(mktemp -d)"
+      roleprobe_original_path="$PATH"
+      cat > "$roleprobe_mock_bin/nc" <<'EOF'
+#!/bin/sh
+case "${ROLEPROBE_NC_CASE:-}" in
+  standalone|leader|follower|observer|looking)
+    printf 'Mode: %s\n' "$ROLEPROBE_NC_CASE"
+    ;;
+  ambiguous)
+    printf 'Mode: leader\nMode: follower\n'
+    ;;
+  empty)
+    ;;
+  refused)
+    exit 1
+    ;;
+  *)
+    exit 2
+    ;;
+esac
+EOF
+      chmod +x "$roleprobe_mock_bin/nc"
+      export PATH="$roleprobe_mock_bin:$PATH"
+    }
+
+    cleanup_mock_nc() {
+      PATH="$roleprobe_original_path"
+      export PATH
+      rm -rf "$roleprobe_mock_bin"
+      unset roleprobe_mock_bin roleprobe_original_path ROLEPROBE_NC_CASE
+    }
+
+    BeforeEach "setup_mock_nc"
+    AfterEach "cleanup_mock_nc"
+
+    It "maps standalone to leader"
+      export ROLEPROBE_NC_CASE="standalone"
+
+      When run command bash ../scripts/roleprobe.sh
+      The status should be success
+      The output should eq "leader"
+    End
+
+    It "publishes leader"
+      export ROLEPROBE_NC_CASE="leader"
+
+      When run command bash ../scripts/roleprobe.sh
+      The status should be success
+      The output should eq "leader"
+    End
+
+    It "publishes follower"
+      export ROLEPROBE_NC_CASE="follower"
+
+      When run command bash ../scripts/roleprobe.sh
+      The status should be success
+      The output should eq "follower"
+    End
+
+    It "publishes observer"
+      export ROLEPROBE_NC_CASE="observer"
+
+      When run command bash ../scripts/roleprobe.sh
+      The status should be success
+      The output should eq "observer"
+    End
+
+    It "fails closed on empty output"
+      export ROLEPROBE_NC_CASE="empty"
+
+      When run command bash ../scripts/roleprobe.sh
+      The status should be failure
+      The output should eq ""
+    End
+
+    It "fails closed on an uncertain mode"
+      export ROLEPROBE_NC_CASE="looking"
+
+      When run command bash ../scripts/roleprobe.sh
+      The status should be failure
+      The output should eq ""
+    End
+
+    It "fails closed when nc refuses the connection"
+      export ROLEPROBE_NC_CASE="refused"
+
+      When run command bash ../scripts/roleprobe.sh
+      The status should be failure
+      The output should eq ""
+    End
+
+    It "fails closed on ambiguous multi-line mode output"
+      export ROLEPROBE_NC_CASE="ambiguous"
+
+      When run command bash ../scripts/roleprobe.sh
+      The status should be failure
+      The output should eq ""
+    End
   End
-End
+End

addons/zookeeper/scripts/roleprobe.sh

@@ -1,29 +1,48 @@
 #!/bin/bash
 
+set -o pipefail
+
 get_zk_mode_from_script() {
   $ZOOBINDIR/zkServer.sh status
 }
 
 get_zookeeper_mode() {
+  local stat mode
   if command -v nc >/dev/null 2>&1; then
-    local stat
-    stat=$(echo srvr | nc 127.0.0.1 2181 | grep Mode)
-    echo "$stat" | awk '{print $2}'
+    if ! stat=$(echo srvr | nc 127.0.0.1 2181 | grep '^Mode:'); then
+      return 1
+    fi
   else
-    local stat
-    stat=$(get_zk_mode_from_script)
-    echo "$stat" | grep "Mode:" | awk '{print $2}'
+    if ! stat=$(get_zk_mode_from_script | grep '^Mode:'); then
+      return 1
+    fi
   fi
+
+  mode=$(echo "$stat" | awk '{print $2}')
+  case "$mode" in
+    standalone|leader|follower|observer)
+      printf "%s" "$mode"
+      ;;
+    *)
+      return 1
+      ;;
+  esac
 }
 
 get_zk_role() {
   local mode
-  mode=$(get_zookeeper_mode)
-  if [[ "$mode" == "standalone" ]]; then
-    printf "leader"
-  else
-    printf "%s" "$mode"
-  fi
+  mode=$(get_zookeeper_mode) || return 1
+  case "$mode" in
+    standalone)
+      printf "leader"
+      ;;
+    leader|follower|observer)
+      printf "%s" "$mode"
+      ;;
+    *)
+      return 1
+      ;;
+  esac
 }
 
 # This is magic for shellspec ut framework.
@@ -34,4 +53,4 @@ get_zk_role() {
 ${__SOURCED__:+false} : || return 0
 
 # main
-get_zk_role
+get_zk_role