From d5207874e3c02516364993c866b09fad8f05bee3 Mon Sep 17 00:00:00 2001
From: bufferflies <1045931706@qq.com>
Date: Thu, 1 Feb 2024 15:21:43 +0800
Subject: [PATCH] reject modify store limit configuration if store limit
 version isn't v1

Signed-off-by: bufferflies <1045931706@qq.com>
---
 server/api/store.go                    | 16 ++++++++++++-
 tools/pd-ctl/tests/store/store_test.go | 32 ++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/server/api/store.go b/server/api/store.go
index 44e178c23fd..51ac4fb891c 100644
--- a/server/api/store.go
+++ b/server/api/store.go
@@ -429,6 +429,10 @@ func (h *storeHandler) SetStoreWeight(w http.ResponseWriter, r *http.Request) {
 // @Router   /store/{id}/limit [post]
 func (h *storeHandler) SetStoreLimit(w http.ResponseWriter, r *http.Request) {
 	rc := getCluster(r)
+	if version := rc.GetScheduleConfig().StoreLimitVersion; version != storelimit.VersionV1 {
+		h.rd.JSON(w, http.StatusBadRequest, fmt.Sprintf("current store limit version:%s not support set limit", version))
+		return
+	}
 	vars := mux.Vars(r)
 	storeID, errParse := apiutil.ParseUint64VarsField(vars, "id")
 	if errParse != nil {
@@ -529,6 +533,11 @@ func (h *storesHandler) RemoveTombStone(w http.ResponseWriter, r *http.Request)
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router   /stores/limit [post]
 func (h *storesHandler) SetAllStoresLimit(w http.ResponseWriter, r *http.Request) {
+	cfg := h.GetScheduleConfig()
+	if version := cfg.StoreLimitVersion; version != storelimit.VersionV1 {
+		h.rd.JSON(w, http.StatusBadRequest, fmt.Sprintf("current store limit version:%s not support get limit", version))
+		return
+	}
 	var input map[string]interface{}
 	if err := apiutil.ReadJSONRespondError(h.rd, w, r.Body, &input); err != nil {
 		return
@@ -609,7 +618,12 @@ func (h *storesHandler) SetAllStoresLimit(w http.ResponseWriter, r *http.Request
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router   /stores/limit [get]
 func (h *storesHandler) GetAllStoresLimit(w http.ResponseWriter, r *http.Request) {
-	limits := h.GetScheduleConfig().StoreLimit
+	cfg := h.GetScheduleConfig()
+	if version := cfg.StoreLimitVersion; version != storelimit.VersionV1 {
+		h.rd.JSON(w, http.StatusBadRequest, fmt.Sprintf("current store limit version:%s not support get limit", version))
+		return
+	}
+	limits := cfg.StoreLimit
 	includeTombstone := false
 	var err error
 	if includeStr := r.URL.Query().Get("include_tombstone"); includeStr != "" {
diff --git a/tools/pd-ctl/tests/store/store_test.go b/tools/pd-ctl/tests/store/store_test.go
index c8103414e9c..c7ddc33e9c6 100644
--- a/tools/pd-ctl/tests/store/store_test.go
+++ b/tools/pd-ctl/tests/store/store_test.go
@@ -38,6 +38,38 @@ import (
 	"go.etcd.io/etcd/pkg/transport"
 )
 
+func TestStoreLimitV2(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	cluster, err := pdTests.NewTestCluster(ctx, 1)
+	re.NoError(err)
+	err = cluster.RunInitialServers()
+	re.NoError(err)
+	cluster.WaitLeader()
+	pdAddr := cluster.GetConfig().GetClientURL()
+	cmd := ctl.GetRootCmd()
+
+	leaderServer := cluster.GetLeaderServer()
+	re.NoError(leaderServer.BootstrapCluster())
+	defer cluster.Destroy()
+
+	// store command
+	args := []string{"-u", pdAddr, "config", "set", "store-limit-version", "v2"}
+	_, err = tests.ExecuteCommand(cmd, args...)
+	re.NoError(err)
+
+	args = []string{"-u", pdAddr, "store", "limit"}
+	output, err := tests.ExecuteCommand(cmd, args...)
+	re.NoError(err)
+	re.Contains(string(output), "not support get limit")
+
+	args = []string{"-u", pdAddr, "store", "limit", "1", "10"}
+	output, err = tests.ExecuteCommand(cmd, args...)
+	re.NoError(err)
+	re.Contains(string(output), "not support set limit")
+}
+
 func TestStore(t *testing.T) {
 	re := require.New(t)
 	ctx, cancel := context.WithCancel(context.Background())