From 6d8f44a6e28fc9bcca29fbfe1ac6040c6c24978c Mon Sep 17 00:00:00 2001
From: Shizun Ge <shizunge@gmail.com>
Date: Wed, 1 Nov 2023 20:03:24 -0700
Subject: [PATCH] Use curl to get the docker hub rate when possible.

The busybox wget would consume a docker hub rate.
---
 src/docker_hub_rate.sh | 35 ++++++++++++++++++++++++++++-------
 1 file changed, 28 insertions(+), 7 deletions(-)

diff --git a/src/docker_hub_rate.sh b/src/docker_hub_rate.sh
index ffd2650..2806d14 100755
--- a/src/docker_hub_rate.sh
+++ b/src/docker_hub_rate.sh
@@ -15,20 +15,41 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 #
 
+docker_hub_rate_token() {
+  local IMAGE="${1:-ratelimitpreview/test}"
+  local TOKEN_URL="https://auth.docker.io/token?service=registry.docker.io&scope=repository:${IMAGE}:pull"
+  if curl --version 1>/dev/null 2>&1; then
+    curl -s ${TOKEN_URL}
+    return $?
+  fi
+  wget -qO-  ${TOKEN_URL}
+}
+
+docker_hub_rate_read_rate() {
+  local IMAGE="${1:-ratelimitpreview/test}"
+  local TOKEN="${2}"
+  [ -z "${TOKEN}" ] && echo "[GET TOKEN ERROR]" && return 1
+  local HEADER="Authorization: Bearer ${TOKEN}"
+  local URL="https://registry-1.docker.io/v2/${IMAGE}/manifests/latest"
+  if curl --version 1>/dev/null 2>&1; then
+    curl --head -H "${HEADER}" "${URL}" 2>&1
+    return $?
+  fi
+  # Add `--spider`` implies that you want to send a HEAD request (as opposed to GET or POST).
+  # The `busybox wget` does not send a HEAD request, thus it will consume a docker hub rate.
+  wget -qS --spider --header="${HEADER}" -O /dev/null "${URL}" 2>&1
+}
+
 docker_hub_rate() {
   local IMAGE="${1:-ratelimitpreview/test}"
   local RESPONSE=
-  if ! RESPONSE=$(wget -qO- "https://auth.docker.io/token?service=registry.docker.io&scope=repository:${IMAGE}:pull"); then
+  if ! RESPONSE=$(docker_hub_rate_token "${IMAGE}"); then
     echo "[GET TOKEN RESPONSE ERROR]"
     return 1
   fi
   local TOKEN=
   TOKEN=$(echo "${RESPONSE}" | sed 's/.*"token":"\([^"]*\).*/\1/')
-  [ -z "${TOKEN}" ] && echo "[GET TOKEN ERROR]" && return 1
-  local HEADER="Authorization: Bearer ${TOKEN}"
-  local URL="https://registry-1.docker.io/v2/${IMAGE}/manifests/latest"
-  # adding --spider implies that you want to send a HEAD request (as opposed to GET or POST).
-  if ! RESPONSE=$(wget -qS --spider --header="${HEADER}" -O /dev/null "${URL}" 2>&1); then
+  if ! RESPONSE=$(docker_hub_rate_read_rate "${IMAGE}" "${TOKEN}"); then
     if echo "${RESPONSE}" | grep -q "Too Many Requests" ; then
       echo "0"
       return 0
@@ -40,4 +61,4 @@ docker_hub_rate() {
   RATE=$(echo "${RESPONSE}" | sed -n 's/.*ratelimit-remaining: \([0-9]*\).*/\1/p' )
   [ -z "${RATE}" ] && echo "[GET RATE ERROR]" && return 1
   echo "${RATE}"
-}
\ No newline at end of file
+}