From a83638f21b8bf26da14e061da9d46a3a1dbeb5b7 Mon Sep 17 00:00:00 2001
From: chenxiaobin19 <1025221611@qq.com>
Date: Tue, 20 Feb 2024 11:06:49 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AE=8C=E5=96=84port=E5=85=A5=E5=8F=82?=
 =?UTF-8?q?=E6=A0=A1=E9=AA=8C=E5=92=8C=E6=89=93=E5=8D=B0=E4=BF=A1=E6=81=AF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/bin/gs_retrieve/gs_retrieve.cpp | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/src/bin/gs_retrieve/gs_retrieve.cpp b/src/bin/gs_retrieve/gs_retrieve.cpp
index 887330cf3..273dc5cdb 100644
--- a/src/bin/gs_retrieve/gs_retrieve.cpp
+++ b/src/bin/gs_retrieve/gs_retrieve.cpp
@@ -709,6 +709,18 @@ XLogRecPtr DoradoFindMaxLsn(char *workingPath, char *returnMsg, int msgLen, pg_c
     return maxLsn;
 }
 
+static bool checkIsDigit(const char* arg)
+{
+    int i = 0;
+    while (arg[i] != '\0') {
+        if (!isdigit(arg[i])) {
+            return false;
+        }
+        i++;
+    }
+    return true;
+}
+
 static int getOptions(const int argc, char* const* argv)
 {
 #define NEWHOST 1
@@ -749,6 +761,11 @@ static int getOptions(const int argc, char* const* argv)
                 newHost = pg_strdup(optarg);
                 break;
             case NEWPORT:
+                check_env_value_c(optarg);
+                if (checkIsDigit(optarg) == 0) {
+                    write_stderr(_("%s: invalid port number \"%s\"\n"), progname, optarg);
+                    return -1;
+                }
                 newPort = atoi(optarg);
                 break;
             case OLDHOST:
@@ -759,6 +776,11 @@ static int getOptions(const int argc, char* const* argv)
                 oldHost = pg_strdup(optarg);
                 break;
             case OLDPORT:
+                check_env_value_c(optarg);
+                if (checkIsDigit(optarg) == 0) {
+                    write_stderr(_("%s: invalid port number \"%s\"\n"), progname, optarg);
+                    return -1;
+                }
                 oldPort = atoi(optarg);
                 break;
             case 'U':
@@ -837,7 +859,7 @@ int main(int argc, char **argv)
     }
 
     if (newPort <= 0) {
-        write_stderr(_("%s: no newport specified\n"), progname);
+        write_stderr(_("%s: no newport specified or it must be a true value\n"), progname);
         do_advice();
         exit(1);
     }
@@ -849,7 +871,7 @@ int main(int argc, char **argv)
     }
 
     if (oldPort <= 0) {
-        write_stderr(_("%s: no oldport specified\n"), progname);
+        write_stderr(_("%s: no oldport specified or it must be a true value\n"), progname);
         do_advice();
         exit(1);
     }