config, api: add a config to turn off traffic replay on TiDB Cloud (#738) (#800)

Signed-off-by: ti-chi-bot <[email protected]>
Co-authored-by: djshow832 <[email protected]>

Author: ti-chi-bot

lib/config/proxy.go

@@ -22,15 +22,16 @@ var (
 )
 
 type Config struct {
-	Proxy    ProxyServer       `yaml:"proxy,omitempty" toml:"proxy,omitempty" json:"proxy,omitempty"`
-	API      API               `yaml:"api,omitempty" toml:"api,omitempty" json:"api,omitempty"`
-	Advance  Advance           `yaml:"advance,omitempty" toml:"advance,omitempty" json:"advance,omitempty"`
-	Workdir  string            `yaml:"workdir,omitempty" toml:"workdir,omitempty" json:"workdir,omitempty"`
-	Security Security          `yaml:"security,omitempty" toml:"security,omitempty" json:"security,omitempty"`
-	Log      Log               `yaml:"log,omitempty" toml:"log,omitempty" json:"log,omitempty"`
-	Balance  Balance           `yaml:"balance,omitempty" toml:"balance,omitempty" json:"balance,omitempty"`
-	Labels   map[string]string `yaml:"labels,omitempty" toml:"labels,omitempty" json:"labels,omitempty"`
-	HA       HA                `yaml:"ha,omitempty" toml:"ha,omitempty" json:"ha,omitempty"`
+	Proxy               ProxyServer       `yaml:"proxy,omitempty" toml:"proxy,omitempty" json:"proxy,omitempty"`
+	API                 API               `yaml:"api,omitempty" toml:"api,omitempty" json:"api,omitempty"`
+	Advance             Advance           `yaml:"advance,omitempty" toml:"advance,omitempty" json:"advance,omitempty"`
+	Workdir             string            `yaml:"workdir,omitempty" toml:"workdir,omitempty" json:"workdir,omitempty"`
+	Security            Security          `yaml:"security,omitempty" toml:"security,omitempty" json:"security,omitempty"`
+	Log                 Log               `yaml:"log,omitempty" toml:"log,omitempty" json:"log,omitempty"`
+	Balance             Balance           `yaml:"balance,omitempty" toml:"balance,omitempty" json:"balance,omitempty"`
+	Labels              map[string]string `yaml:"labels,omitempty" toml:"labels,omitempty" json:"labels,omitempty"`
+	HA                  HA                `yaml:"ha,omitempty" toml:"ha,omitempty" json:"ha,omitempty"`
+	EnableTrafficReplay bool              `yaml:"enable-traffic-replay,omitempty" toml:"enable-traffic-replay,omitempty" json:"enable-traffic-replay,omitempty"`
 }
 
 type KeepAlive struct {
@@ -164,6 +165,8 @@ func NewConfig() *Config {
 
 	cfg.Balance = DefaultBalance()
 
+	cfg.EnableTrafficReplay = true
+
 	return &cfg
 }
 

pkg/server/api/traffic.go

@@ -16,11 +16,17 @@ import (
 func (h *Server) registerTraffic(group *gin.RouterGroup) {
 	group.POST("/capture", h.TrafficCapture)
 	group.POST("/replay", h.TrafficReplay)
-	group.POST("/cancel", h.TrafficStop)
+	group.POST("/cancel", h.TrafficCancel)
 	group.GET("/show", h.TrafficShow)
 }
 
 func (h *Server) TrafficCapture(c *gin.Context) {
+	globalCfg := h.mgr.CfgMgr.GetConfig()
+	if !globalCfg.EnableTrafficReplay {
+		c.String(http.StatusBadRequest, "traffic capture is disabled")
+		return
+	}
+
 	cfg := capture.CaptureConfig{}
 	cfg.Output = c.PostForm("output")
 	if durationStr := c.PostForm("duration"); durationStr != "" {
@@ -40,6 +46,12 @@ func (h *Server) TrafficCapture(c *gin.Context) {
 }
 
 func (h *Server) TrafficReplay(c *gin.Context) {
+	globalCfg := h.mgr.CfgMgr.GetConfig()
+	if !globalCfg.EnableTrafficReplay {
+		c.String(http.StatusBadRequest, "traffic replay is disabled")
+		return
+	}
+
 	cfg := replay.ReplayConfig{}
 	cfg.Input = c.PostForm("input")
 	if speedStr := c.PostForm("speed"); speedStr != "" {
@@ -60,12 +72,24 @@ func (h *Server) TrafficReplay(c *gin.Context) {
 	c.String(http.StatusOK, "replay started")
 }
 
-func (h *Server) TrafficStop(c *gin.Context) {
+func (h *Server) TrafficCancel(c *gin.Context) {
+	globalCfg := h.mgr.CfgMgr.GetConfig()
+	if !globalCfg.EnableTrafficReplay {
+		c.String(http.StatusBadRequest, "traffic cancel is disabled")
+		return
+	}
+
 	result := h.mgr.ReplayJobMgr.Stop()
 	c.String(http.StatusOK, result)
 }
 
 func (h *Server) TrafficShow(c *gin.Context) {
+	globalCfg := h.mgr.CfgMgr.GetConfig()
+	if !globalCfg.EnableTrafficReplay {
+		c.String(http.StatusBadRequest, "traffic show is disabled")
+		return
+	}
+
 	result := h.mgr.ReplayJobMgr.Jobs()
 	c.String(http.StatusOK, result)
 }

pkg/server/api/traffic_test.go

@@ -88,6 +88,42 @@ func TestTraffic(t *testing.T) {
 	})
 }
 
+func TestDisableTrafficReplay(t *testing.T) {
+	server, doHTTP := createServer(t)
+	server.mgr.CfgMgr.GetConfig().EnableTrafficReplay = false
+
+	doHTTP(t, http.MethodPost, "/api/traffic/capture", httpOpts{
+		reader: cli.GetFormReader(map[string]string{"output": "/tmp", "duration": "1h"}),
+		header: map[string]string{"Content-Type": "application/x-www-form-urlencoded"},
+	}, func(t *testing.T, r *http.Response) {
+		require.Equal(t, http.StatusBadRequest, r.StatusCode)
+		all, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+		require.Equal(t, "traffic capture is disabled", string(all))
+	})
+	doHTTP(t, http.MethodPost, "/api/traffic/replay", httpOpts{
+		reader: cli.GetFormReader(map[string]string{"input": "/tmp"}),
+		header: map[string]string{"Content-Type": "application/x-www-form-urlencoded"},
+	}, func(t *testing.T, r *http.Response) {
+		require.Equal(t, http.StatusBadRequest, r.StatusCode)
+		all, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+		require.Equal(t, "traffic replay is disabled", string(all))
+	})
+	doHTTP(t, http.MethodPost, "/api/traffic/cancel", httpOpts{}, func(t *testing.T, r *http.Response) {
+		require.Equal(t, http.StatusBadRequest, r.StatusCode)
+		all, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+		require.Equal(t, "traffic cancel is disabled", string(all))
+	})
+	doHTTP(t, http.MethodGet, "/api/traffic/show", httpOpts{}, func(t *testing.T, r *http.Response) {
+		require.Equal(t, http.StatusBadRequest, r.StatusCode)
+		all, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+		require.Equal(t, "traffic show is disabled", string(all))
+	})
+}
+
 var _ manager.JobManager = (*mockReplayJobManager)(nil)
 
 type mockReplayJobManager struct {