Skip to content

Commit acbb0eb

Browse files
djshow832djshow832
authored
config, api: add a config to turn off traffic replay on TiDB Cloud (#738)
1 parent 3ced6ae commit acbb0eb

File tree

3 files changed

+76
-13
lines changed

3 files changed

+76
-13
lines changed

‎lib/config/proxy.go‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -22,15 +22,16 @@ var (
2222
)
2323

2424
type Config struct {
25-
Proxy ProxyServer `yaml:"proxy,omitempty" toml:"proxy,omitempty" json:"proxy,omitempty"`
26-
API API `yaml:"api,omitempty" toml:"api,omitempty" json:"api,omitempty"`
27-
Advance Advance `yaml:"advance,omitempty" toml:"advance,omitempty" json:"advance,omitempty"`
28-
Workdir string `yaml:"workdir,omitempty" toml:"workdir,omitempty" json:"workdir,omitempty"`
29-
Security Security `yaml:"security,omitempty" toml:"security,omitempty" json:"security,omitempty"`
30-
Log Log `yaml:"log,omitempty" toml:"log,omitempty" json:"log,omitempty"`
31-
Balance Balance `yaml:"balance,omitempty" toml:"balance,omitempty" json:"balance,omitempty"`
32-
Labels map[string]string `yaml:"labels,omitempty" toml:"labels,omitempty" json:"labels,omitempty"`
33-
HA HA `yaml:"ha,omitempty" toml:"ha,omitempty" json:"ha,omitempty"`
25+
Proxy ProxyServer `yaml:"proxy,omitempty" toml:"proxy,omitempty" json:"proxy,omitempty"`
26+
API API `yaml:"api,omitempty" toml:"api,omitempty" json:"api,omitempty"`
27+
Advance Advance `yaml:"advance,omitempty" toml:"advance,omitempty" json:"advance,omitempty"`
28+
Workdir string `yaml:"workdir,omitempty" toml:"workdir,omitempty" json:"workdir,omitempty"`
29+
Security Security `yaml:"security,omitempty" toml:"security,omitempty" json:"security,omitempty"`
30+
Log Log `yaml:"log,omitempty" toml:"log,omitempty" json:"log,omitempty"`
31+
Balance Balance `yaml:"balance,omitempty" toml:"balance,omitempty" json:"balance,omitempty"`
32+
Labels map[string]string `yaml:"labels,omitempty" toml:"labels,omitempty" json:"labels,omitempty"`
33+
HA HA `yaml:"ha,omitempty" toml:"ha,omitempty" json:"ha,omitempty"`
34+
EnableTrafficReplay bool `yaml:"enable-traffic-replay,omitempty" toml:"enable-traffic-replay,omitempty" json:"enable-traffic-replay,omitempty"`
3435
}
3536

3637
type KeepAlive struct {
@@ -137,6 +138,8 @@ func NewConfig() *Config {
137138

138139
cfg.Balance = DefaultBalance()
139140

141+
cfg.EnableTrafficReplay = true
142+
140143
return &cfg
141144
}
142145

‎pkg/server/api/traffic.go‎

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,17 @@ import (
1818
func (h *Server) registerTraffic(group *gin.RouterGroup) {
1919
group.POST("/capture", h.TrafficCapture)
2020
group.POST("/replay", h.TrafficReplay)
21-
group.POST("/cancel", h.TrafficStop)
21+
group.POST("/cancel", h.TrafficCancel)
2222
group.GET("/show", h.TrafficShow)
2323
}
2424

2525
func (h *Server) TrafficCapture(c *gin.Context) {
26+
globalCfg := h.mgr.CfgMgr.GetConfig()
27+
if !globalCfg.EnableTrafficReplay {
28+
c.String(http.StatusBadRequest, "traffic capture is disabled")
29+
return
30+
}
31+
2632
cfg := capture.CaptureConfig{}
2733
cfg.Output = c.PostForm("output")
2834
if durationStr := c.PostForm("duration"); durationStr != "" {
@@ -44,7 +50,7 @@ func (h *Server) TrafficCapture(c *gin.Context) {
4450
}
4551
}
4652
cfg.Compress = compress
47-
cfg.KeyFile = h.mgr.CfgMgr.GetConfig().Security.Encryption.KeyPath
53+
cfg.KeyFile = globalCfg.Security.Encryption.KeyPath
4854
if startTimeStr := c.PostForm("start-time"); startTimeStr != "" {
4955
startTime, err := time.Parse(time.RFC3339, startTimeStr)
5056
if err != nil {
@@ -64,6 +70,12 @@ func (h *Server) TrafficCapture(c *gin.Context) {
6470
}
6571

6672
func (h *Server) TrafficReplay(c *gin.Context) {
73+
globalCfg := h.mgr.CfgMgr.GetConfig()
74+
if !globalCfg.EnableTrafficReplay {
75+
c.String(http.StatusBadRequest, "traffic replay is disabled")
76+
return
77+
}
78+
6779
cfg := replay.ReplayConfig{}
6880
cfg.Input = c.PostForm("input")
6981
if speedStr := c.PostForm("speed"); speedStr != "" {
@@ -87,7 +99,7 @@ func (h *Server) TrafficReplay(c *gin.Context) {
8799
cfg.Username = c.PostForm("username")
88100
cfg.Password = c.PostForm("password")
89101
cfg.ReadOnly = strings.EqualFold(c.PostForm("readonly"), "true")
90-
cfg.KeyFile = h.mgr.CfgMgr.GetConfig().Security.Encryption.KeyPath
102+
cfg.KeyFile = globalCfg.Security.Encryption.KeyPath
91103

92104
if err := h.mgr.ReplayJobMgr.StartReplay(cfg); err != nil {
93105
c.String(http.StatusInternalServerError, err.Error())
@@ -96,12 +108,24 @@ func (h *Server) TrafficReplay(c *gin.Context) {
96108
c.String(http.StatusOK, "replay started")
97109
}
98110

99-
func (h *Server) TrafficStop(c *gin.Context) {
111+
func (h *Server) TrafficCancel(c *gin.Context) {
112+
globalCfg := h.mgr.CfgMgr.GetConfig()
113+
if !globalCfg.EnableTrafficReplay {
114+
c.String(http.StatusBadRequest, "traffic cancel is disabled")
115+
return
116+
}
117+
100118
result := h.mgr.ReplayJobMgr.Stop()
101119
c.String(http.StatusOK, result)
102120
}
103121

104122
func (h *Server) TrafficShow(c *gin.Context) {
123+
globalCfg := h.mgr.CfgMgr.GetConfig()
124+
if !globalCfg.EnableTrafficReplay {
125+
c.String(http.StatusBadRequest, "traffic show is disabled")
126+
return
127+
}
128+
105129
result := h.mgr.ReplayJobMgr.Jobs()
106130
c.String(http.StatusOK, result)
107131
}

‎pkg/server/api/traffic_test.go‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,42 @@ func cancelJob(t *testing.T, doHTTP doHTTPFunc) {
128128
})
129129
}
130130

131+
func TestDisableTrafficReplay(t *testing.T) {
132+
server, doHTTP := createServer(t)
133+
server.mgr.CfgMgr.GetConfig().EnableTrafficReplay = false
134+
135+
doHTTP(t, http.MethodPost, "/api/traffic/capture", httpOpts{
136+
reader: cli.GetFormReader(map[string]string{"output": "/tmp", "duration": "1h"}),
137+
header: map[string]string{"Content-Type": "application/x-www-form-urlencoded"},
138+
}, func(t *testing.T, r *http.Response) {
139+
require.Equal(t, http.StatusBadRequest, r.StatusCode)
140+
all, err := io.ReadAll(r.Body)
141+
require.NoError(t, err)
142+
require.Equal(t, "traffic capture is disabled", string(all))
143+
})
144+
doHTTP(t, http.MethodPost, "/api/traffic/replay", httpOpts{
145+
reader: cli.GetFormReader(map[string]string{"input": "/tmp"}),
146+
header: map[string]string{"Content-Type": "application/x-www-form-urlencoded"},
147+
}, func(t *testing.T, r *http.Response) {
148+
require.Equal(t, http.StatusBadRequest, r.StatusCode)
149+
all, err := io.ReadAll(r.Body)
150+
require.NoError(t, err)
151+
require.Equal(t, "traffic replay is disabled", string(all))
152+
})
153+
doHTTP(t, http.MethodPost, "/api/traffic/cancel", httpOpts{}, func(t *testing.T, r *http.Response) {
154+
require.Equal(t, http.StatusBadRequest, r.StatusCode)
155+
all, err := io.ReadAll(r.Body)
156+
require.NoError(t, err)
157+
require.Equal(t, "traffic cancel is disabled", string(all))
158+
})
159+
doHTTP(t, http.MethodGet, "/api/traffic/show", httpOpts{}, func(t *testing.T, r *http.Response) {
160+
require.Equal(t, http.StatusBadRequest, r.StatusCode)
161+
all, err := io.ReadAll(r.Body)
162+
require.NoError(t, err)
163+
require.Equal(t, "traffic show is disabled", string(all))
164+
})
165+
}
166+
131167
var _ manager.JobManager = (*mockReplayJobManager)(nil)
132168

133169
type mockReplayJobManager struct {

0 commit comments

Comments
 (0)