From 4edce7cbdb20452d421ef378581a6be1453a39b8 Mon Sep 17 00:00:00 2001 From: Breezewish Date: Tue, 23 Nov 2021 21:14:42 +0800 Subject: [PATCH 1/2] refactor: Switch to use ziputil, netutil, reflectutil and fileswap --- .golangci.yml | 1 + .../clusterinfo/hostinfo/cluster_config.go | 4 +- .../clusterinfo/hostinfo/cluster_hardware.go | 6 +- .../clusterinfo/hostinfo/cluster_load.go | 4 +- pkg/apiserver/debugapi/service.go | 35 +++--- pkg/apiserver/logsearch/pack.go | 4 +- pkg/apiserver/profiling/router.go | 5 +- pkg/apiserver/slowquery/model.go | 3 +- pkg/apiserver/statement/models.go | 3 +- pkg/apiserver/utils/fs_stream.go | 119 ------------------ pkg/apiserver/utils/reflections.go | 32 ----- pkg/apiserver/utils/reflections_test.go | 39 ------ pkg/apiserver/utils/zip.go | 56 --------- pkg/utils/host/host.go | 38 ------ pkg/utils/topology/pd.go | 4 +- pkg/utils/topology/store.go | 6 +- pkg/utils/topology/tidb.go | 4 +- util/rest/fileswap/server_test.go | 39 +++++- 18 files changed, 82 insertions(+), 320 deletions(-) delete mode 100644 pkg/apiserver/utils/fs_stream.go delete mode 100644 pkg/apiserver/utils/reflections.go delete mode 100644 pkg/apiserver/utils/reflections_test.go delete mode 100644 pkg/apiserver/utils/zip.go delete mode 100644 pkg/utils/host/host.go diff --git a/.golangci.yml b/.golangci.yml index 4696b676fe..cef0a2efb0 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -2,6 +2,7 @@ run: skip-dirs: - swaggerspec - pkg/uiserver + - ui timeout: 2m issues: diff --git a/pkg/apiserver/clusterinfo/hostinfo/cluster_config.go b/pkg/apiserver/clusterinfo/hostinfo/cluster_config.go index 1a3fd71c23..e9c298588f 100644 --- a/pkg/apiserver/clusterinfo/hostinfo/cluster_config.go +++ b/pkg/apiserver/clusterinfo/hostinfo/cluster_config.go @@ -7,7 +7,7 @@ import ( "gorm.io/gorm" - "github.com/pingcap/tidb-dashboard/pkg/utils/host" + "github.com/pingcap/tidb-dashboard/util/netutil" ) type clusterConfigModel struct { @@ -29,7 +29,7 @@ func FillInstances(db *gorm.DB, m InfoMap) error { } for _, row := range rows { - hostname, _, err := host.ParseHostAndPortFromAddress(row.Instance) + hostname, _, err := netutil.ParseHostAndPortFromAddress(row.Instance) if err != nil { continue } diff --git a/pkg/apiserver/clusterinfo/hostinfo/cluster_hardware.go b/pkg/apiserver/clusterinfo/hostinfo/cluster_hardware.go index 78a1f926d3..c9e83bd5ea 100644 --- a/pkg/apiserver/clusterinfo/hostinfo/cluster_hardware.go +++ b/pkg/apiserver/clusterinfo/hostinfo/cluster_hardware.go @@ -9,7 +9,7 @@ import ( "gorm.io/gorm" - "github.com/pingcap/tidb-dashboard/pkg/utils/host" + "github.com/pingcap/tidb-dashboard/util/netutil" ) // Used to deserialize from JSON_VALUE. @@ -46,7 +46,7 @@ func FillFromClusterHardwareTable(db *gorm.DB, m InfoMap) error { tiFlashDisks := make([]clusterTableModel, 0) for _, row := range rows { - hostname, _, err := host.ParseHostAndPortFromAddress(row.Instance) + hostname, _, err := netutil.ParseHostAndPortFromAddress(row.Instance) if err != nil { continue } @@ -126,7 +126,7 @@ func FillFromClusterHardwareTable(db *gorm.DB, m InfoMap) error { } // Back fill TiFlash instances for instance, de := range tiFlashDiskInfo { - hostname, _, err := host.ParseHostAndPortFromAddress(instance) + hostname, _, err := netutil.ParseHostAndPortFromAddress(instance) if err != nil { panic(err) } diff --git a/pkg/apiserver/clusterinfo/hostinfo/cluster_load.go b/pkg/apiserver/clusterinfo/hostinfo/cluster_load.go index 6854caae2a..0540b4d9f1 100644 --- a/pkg/apiserver/clusterinfo/hostinfo/cluster_load.go +++ b/pkg/apiserver/clusterinfo/hostinfo/cluster_load.go @@ -8,7 +8,7 @@ import ( "gorm.io/gorm" - "github.com/pingcap/tidb-dashboard/pkg/utils/host" + "github.com/pingcap/tidb-dashboard/util/netutil" ) // Used to deserialize from JSON_VALUE. @@ -40,7 +40,7 @@ func FillFromClusterLoadTable(db *gorm.DB, m InfoMap) error { } for _, row := range rows { - hostname, _, err := host.ParseHostAndPortFromAddress(row.Instance) + hostname, _, err := netutil.ParseHostAndPortFromAddress(row.Instance) if err != nil { continue } diff --git a/pkg/apiserver/debugapi/service.go b/pkg/apiserver/debugapi/service.go index 131cb89732..5a6d704d38 100644 --- a/pkg/apiserver/debugapi/service.go +++ b/pkg/apiserver/debugapi/service.go @@ -14,8 +14,8 @@ import ( "github.com/pingcap/tidb-dashboard/pkg/apiserver/debugapi/endpoint" "github.com/pingcap/tidb-dashboard/pkg/apiserver/user" - "github.com/pingcap/tidb-dashboard/pkg/apiserver/utils" "github.com/pingcap/tidb-dashboard/util/rest" + "github.com/pingcap/tidb-dashboard/util/rest/fileswap" ) const ( @@ -34,11 +34,15 @@ func registerRouter(r *gin.RouterGroup, auth *user.AuthService, s *Service) { type Service struct { Client *endpoint.Client + + fSwap *fileswap.Handler } func newService(hp httpClientParam) *Service { return &Service{ Client: endpoint.NewClient(newHTTPClient(hp), endpointDefs), + + fSwap: fileswap.New(), } } @@ -82,27 +86,31 @@ func (s *Service) RequestEndpoint(c *gin.Context) { } defer res.Response.Body.Close() //nolint:errcheck - ext := getExtFromContentTypeHeader(res.Header.Get("Content-Type")) - fileName := fmt.Sprintf("%s_%d%s", req.EndpointID, time.Now().Unix(), ext) - - writer, token, err := utils.FSPersist(utils.FSPersistConfig{ - TokenIssuer: tokenIssuer, - TokenExpire: time.Minute * 5, // Note: the expire time should include request time. - TempFilePattern: "debug_api", - DownloadFileName: fileName, - }) + writer, err := s.fSwap.NewFileWriter("debug_api") if err != nil { _ = c.Error(err) return } - defer writer.Close() //nolint:errcheck + defer func() { + _ = writer.Close() + }() + _, err = io.Copy(writer, res.Response.Body) if err != nil { _ = c.Error(err) return } - c.String(http.StatusOK, token) + ext := getExtFromContentTypeHeader(res.Header.Get("Content-Type")) + fileName := fmt.Sprintf("%s_%d%s", req.EndpointID, time.Now().Unix(), ext) + downloadToken, err := writer.GetDownloadToken(fileName, time.Minute*5) + if err != nil { + // This shall never happen + _ = c.Error(err) + return + } + + c.String(http.StatusOK, downloadToken) } // @Summary Download a finished request result @@ -112,8 +120,7 @@ func (s *Service) RequestEndpoint(c *gin.Context) { // @Failure 500 {object} rest.ErrorResponse // @Router /debug_api/download [get] func (s *Service) Download(c *gin.Context) { - token := c.Query("token") - utils.FSServe(c, token, tokenIssuer) + s.fSwap.HandleDownloadRequest(c) } // @Summary Get all endpoints diff --git a/pkg/apiserver/logsearch/pack.go b/pkg/apiserver/logsearch/pack.go index e2559d9e48..84be383963 100644 --- a/pkg/apiserver/logsearch/pack.go +++ b/pkg/apiserver/logsearch/pack.go @@ -9,8 +9,8 @@ import ( "github.com/pingcap/log" "go.uber.org/zap" - "github.com/pingcap/tidb-dashboard/pkg/apiserver/utils" "github.com/pingcap/tidb-dashboard/util/rest" + "github.com/pingcap/tidb-dashboard/util/ziputil" ) func serveTaskForDownload(task *TaskModel, c *gin.Context) { @@ -41,7 +41,7 @@ func serveMultipleTaskForDownload(tasks []*TaskModel, c *gin.Context) { c.Writer.Header().Set("Content-type", "application/octet-stream") c.Writer.Header().Set("Content-Disposition", "attachment; filename=\"logs.zip\"") - err := utils.StreamZipPack(c.Writer, filePaths, false) + err := ziputil.WriteZipFromFiles(c.Writer, filePaths, false) if err != nil { log.Error("Stream zip pack failed", zap.Error(err)) } diff --git a/pkg/apiserver/profiling/router.go b/pkg/apiserver/profiling/router.go index 6510a9cf8b..7038b5d926 100644 --- a/pkg/apiserver/profiling/router.go +++ b/pkg/apiserver/profiling/router.go @@ -17,6 +17,7 @@ import ( "github.com/pingcap/tidb-dashboard/pkg/apiserver/utils" "github.com/pingcap/tidb-dashboard/pkg/config" "github.com/pingcap/tidb-dashboard/util/rest" + "github.com/pingcap/tidb-dashboard/util/ziputil" ) // Register register the handlers to the service. @@ -223,7 +224,7 @@ func (s *Service) downloadGroup(c *gin.Context) { fileName := fmt.Sprintf("profiling_pack_%d.zip", taskGroupID) c.Writer.Header().Set("Content-type", "application/octet-stream") c.Writer.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s\"", fileName)) - err = utils.StreamZipPack(c.Writer, filePathes, true) + err = ziputil.WriteZipFromFiles(c.Writer, filePathes, true) if err != nil { log.Error("Stream zip pack failed", zap.Error(err)) } @@ -262,7 +263,7 @@ func (s *Service) downloadSingle(c *gin.Context) { fileName := fmt.Sprintf("profiling_%d.zip", taskID) c.Writer.Header().Set("Content-type", "application/octet-stream") c.Writer.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s\"", fileName)) - err = utils.StreamZipPack(c.Writer, []string{task.FilePath}, true) + err = ziputil.WriteZipFromFiles(c.Writer, []string{task.FilePath}, true) if err != nil { log.Error("Stream zip pack failed", zap.Error(err)) } diff --git a/pkg/apiserver/slowquery/model.go b/pkg/apiserver/slowquery/model.go index 491ec9d372..75b1ec109c 100644 --- a/pkg/apiserver/slowquery/model.go +++ b/pkg/apiserver/slowquery/model.go @@ -6,6 +6,7 @@ import ( "github.com/thoas/go-funk" "github.com/pingcap/tidb-dashboard/pkg/apiserver/utils" + "github.com/pingcap/tidb-dashboard/util/reflectutil" ) type Model struct { @@ -96,7 +97,7 @@ type Field struct { } func getFieldsAndTags() (slowQueryFields []Field) { - fields := utils.GetFieldsAndTags(Model{}, []string{"gorm", "proj", "json"}) + fields := reflectutil.GetFieldsAndTags(Model{}, []string{"gorm", "proj", "json"}) for _, f := range fields { sqf := Field{ diff --git a/pkg/apiserver/statement/models.go b/pkg/apiserver/statement/models.go index acd9e6a58d..96ba009f5a 100644 --- a/pkg/apiserver/statement/models.go +++ b/pkg/apiserver/statement/models.go @@ -10,6 +10,7 @@ import ( "gorm.io/gorm/schema" "github.com/pingcap/tidb-dashboard/pkg/apiserver/utils" + "github.com/pingcap/tidb-dashboard/util/reflectutil" ) // TimeRange represents a range of time. @@ -135,7 +136,7 @@ type Field struct { var gormDefaultNamingStrategy = schema.NamingStrategy{} func getFieldsAndTags() (stmtFields []Field) { - fields := utils.GetFieldsAndTags(Model{}, []string{"related", "agg", "json"}) + fields := reflectutil.GetFieldsAndTags(Model{}, []string{"related", "agg", "json"}) for _, f := range fields { sf := Field{ diff --git a/pkg/apiserver/utils/fs_stream.go b/pkg/apiserver/utils/fs_stream.go deleted file mode 100644 index d2a1be4f30..0000000000 --- a/pkg/apiserver/utils/fs_stream.go +++ /dev/null @@ -1,119 +0,0 @@ -// Copyright 2021 PingCAP, Inc. Licensed under Apache-2.0. - -package utils - -import ( - "encoding/hex" - "encoding/json" - "fmt" - "io" - "io/ioutil" - "os" - "time" - - "github.com/gin-gonic/gin" - "github.com/gtank/cryptopasta" - "github.com/minio/sio" - - "github.com/pingcap/tidb-dashboard/util/rest" -) - -type FSPersistConfig struct { - TokenIssuer string - TokenExpire time.Duration - TempFilePattern string - DownloadFileName string -} - -type FSPersistTokenBody struct { - TempFileName string - TempFileKeyInHex string - DownloadFileName string -} - -// FSPersist returns a writer and corresponding download token that will persist a stream in the FS in an encrypted way. -func FSPersist(config FSPersistConfig) (io.WriteCloser, string, error) { - file, err := ioutil.TempFile("", config.TempFilePattern) - if err != nil { - return nil, "", err - } - - key := cryptopasta.NewEncryptionKey()[:] - w, err := sio.EncryptWriter(file, sio.Config{ - Key: key, - }) - if err != nil { - _ = file.Close() - _ = os.Remove(file.Name()) - return nil, "", err - } - - keyInHex := hex.EncodeToString(key) - tokenBody := FSPersistTokenBody{ - TempFileName: file.Name(), - TempFileKeyInHex: keyInHex, - DownloadFileName: config.DownloadFileName, - } - tokenBodyStr, err := json.Marshal(tokenBody) - if err != nil { - _ = file.Close() - _ = os.Remove(file.Name()) - return nil, "", err - } - // TODO: Maybe better to generate the token after `w.Close()`. - token, err := NewJWTStringWithExpire(config.TokenIssuer, string(tokenBodyStr), config.TokenExpire) - if err != nil { - _ = file.Close() - _ = os.Remove(file.Name()) - return nil, "", err - } - - // Note: we intentionally keep the temp file not removed, so that it can be downloaded later. - return w, token, nil -} - -// FSServe serves the persisted file in the FS to the user by using the download token. -// The file will be removed after it is served to the user. -func FSServe(c *gin.Context, token string, requiredIssuer string) { - str, err := ParseJWTString(requiredIssuer, token) - if err != nil { - _ = c.Error(rest.ErrBadRequest.Wrap(err, "Invalid download request")) - return - } - var tokenBody FSPersistTokenBody - err = json.Unmarshal([]byte(str), &tokenBody) - if err != nil { - _ = c.Error(err) - return - } - - file, err := os.Open(tokenBody.TempFileName) - if err != nil { - if os.IsNotExist(err) { - // It is possible that token is reused. In this case, raise invalid request error. - _ = c.Error(rest.ErrBadRequest.Wrap(err, "Download file not found")) - } else { - _ = c.Error(err) - } - return - } - defer file.Close() // #nosec - defer os.Remove(tokenBody.TempFileName) // #nosec - - key, err := hex.DecodeString(tokenBody.TempFileKeyInHex) - if err != nil { - _ = c.Error(err) - return - } - - c.Writer.Header().Set("Content-type", "application/octet-stream") - c.Writer.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, tokenBody.DownloadFileName)) - - _, err = sio.Decrypt(c.Writer, file, sio.Config{ - Key: key, - }) - if err != nil { - _ = c.Error(err) - return - } -} diff --git a/pkg/apiserver/utils/reflections.go b/pkg/apiserver/utils/reflections.go deleted file mode 100644 index aef070b0b2..0000000000 --- a/pkg/apiserver/utils/reflections.go +++ /dev/null @@ -1,32 +0,0 @@ -// Copyright 2021 PingCAP, Inc. Licensed under Apache-2.0. - -package utils - -import ( - "reflect" -) - -// GetFieldsAndTags return fields' tags assign by `tags` parameter. -func GetFieldsAndTags(obj interface{}, tags []string) []Field { - fieldTags := []Field{} - t := reflect.TypeOf(obj) - fNum := t.NumField() - for i := 0; i < fNum; i++ { - f := Field{Tags: map[string]string{}} - structField := t.Field(i) - - f.Name = structField.Name - for _, tagName := range tags { - f.Tags[tagName] = structField.Tag.Get(tagName) - } - - fieldTags = append(fieldTags, f) - } - - return fieldTags -} - -type Field struct { - Name string - Tags map[string]string -} diff --git a/pkg/apiserver/utils/reflections_test.go b/pkg/apiserver/utils/reflections_test.go deleted file mode 100644 index 3e1d421e3c..0000000000 --- a/pkg/apiserver/utils/reflections_test.go +++ /dev/null @@ -1,39 +0,0 @@ -// Copyright 2021 PingCAP, Inc. Licensed under Apache-2.0. - -package utils - -import ( - "testing" - - "github.com/stretchr/testify/assert" -) - -type MyStruct struct { - FirstField string `matched:"first tag" value:"whatever"` - SecondField string `matched:"second tag" value:"another whatever"` -} - -func TestGetFieldTags(t *testing.T) { - rst := GetFieldsAndTags(MyStruct{}, []string{"matched", "value"}) - - assert.Equal(t, rst, []Field{ - { - Name: "FirstField", - Tags: map[string]string{ - "matched": "first tag", - "value": "whatever", - }, - }, - { - - Name: "SecondField", - Tags: map[string]string{ - "matched": "second tag", - "value": "another whatever", - }, - }, - }) -} - -// // TODO: support nested struct -// func TestGetFieldTags_with_nested_struct(t *testing.T) {} diff --git a/pkg/apiserver/utils/zip.go b/pkg/apiserver/utils/zip.go deleted file mode 100644 index 31e0c081b2..0000000000 --- a/pkg/apiserver/utils/zip.go +++ /dev/null @@ -1,56 +0,0 @@ -// Copyright 2021 PingCAP, Inc. Licensed under Apache-2.0. - -package utils - -import ( - "archive/zip" - "io" - "os" - "path/filepath" -) - -func StreamZipPack(w io.Writer, files []string, needCompress bool) error { - pack := zip.NewWriter(w) - defer pack.Close() - - for _, file := range files { - err := streamZipFile(pack, file, needCompress) - if err != nil { - return err - } - } - - return nil -} - -func streamZipFile(zipPack *zip.Writer, file string, needCompress bool) error { - f, err := os.Open(filepath.Clean(file)) - if err != nil { - return err - } - defer f.Close() // #nosec - - fileInfo, err := f.Stat() - if err != nil { - return err - } - - zipMethod := zip.Store // no compress - if needCompress { - zipMethod = zip.Deflate // compress - } - zipFile, err := zipPack.CreateHeader(&zip.FileHeader{ - Name: fileInfo.Name(), - Method: zipMethod, - }) - if err != nil { - return err - } - - _, err = io.Copy(zipFile, f) - if err != nil { - return err - } - - return nil -} diff --git a/pkg/utils/host/host.go b/pkg/utils/host/host.go deleted file mode 100644 index e7de6872f4..0000000000 --- a/pkg/utils/host/host.go +++ /dev/null @@ -1,38 +0,0 @@ -// Copyright 2021 PingCAP, Inc. Licensed under Apache-2.0. - -package host - -import ( - "fmt" - "net" - "net/url" - "strconv" - "strings" -) - -// address should be like "ip:port" as "127.0.0.1:2379". -// return error if string is not like "ip:port". -func ParseHostAndPortFromAddress(address string) (string, uint, error) { - host, port, err := net.SplitHostPort(address) - if err != nil { - return "", 0, fmt.Errorf("invalid address: %v", err) - } - portNumeric, err := strconv.Atoi(port) - if err != nil || portNumeric == 0 { - return "", 0, fmt.Errorf("invalid address: invalid port") - } - return strings.ToLower(host), uint(portNumeric), nil -} - -// address should be like "protocol://ip:port" as "http://127.0.0.1:2379". -func ParseHostAndPortFromAddressURL(urlString string) (string, uint, error) { - u, err := url.Parse(urlString) - if err != nil { - return "", 0, fmt.Errorf("invalid address: %v", err) - } - port, err := strconv.Atoi(u.Port()) - if err != nil || port == 0 { - return "", 0, fmt.Errorf("invalid address: invalid port") - } - return strings.ToLower(u.Hostname()), uint(port), nil -} diff --git a/pkg/utils/topology/pd.go b/pkg/utils/topology/pd.go index cf56b7369b..85bbd1580b 100644 --- a/pkg/utils/topology/pd.go +++ b/pkg/utils/topology/pd.go @@ -13,7 +13,7 @@ import ( "github.com/pingcap/tidb-dashboard/pkg/pd" "github.com/pingcap/tidb-dashboard/pkg/utils/distro" - "github.com/pingcap/tidb-dashboard/pkg/utils/host" + "github.com/pingcap/tidb-dashboard/util/netutil" ) func FetchPDTopology(pdClient *pd.Client) ([]PDInfo, error) { @@ -45,7 +45,7 @@ func FetchPDTopology(pdClient *pd.Client) ([]PDInfo, error) { for _, ds := range ds.Members { u := ds.ClientUrls[0] - hostname, port, err := host.ParseHostAndPortFromAddressURL(u) + hostname, port, err := netutil.ParseHostAndPortFromAddressURL(u) if err != nil { continue } diff --git a/pkg/utils/topology/store.go b/pkg/utils/topology/store.go index 6d2b5c057e..b6a894ca95 100644 --- a/pkg/utils/topology/store.go +++ b/pkg/utils/topology/store.go @@ -12,7 +12,7 @@ import ( "github.com/pingcap/tidb-dashboard/pkg/pd" "github.com/pingcap/tidb-dashboard/pkg/utils/distro" - "github.com/pingcap/tidb-dashboard/pkg/utils/host" + "github.com/pingcap/tidb-dashboard/util/netutil" ) // FetchStoreTopology returns TiKV info and TiFlash info. @@ -75,12 +75,12 @@ func FetchStoreLocation(pdClient *pd.Client) (*StoreLocation, error) { func buildStoreTopology(stores []store) []StoreInfo { nodes := make([]StoreInfo, 0, len(stores)) for _, v := range stores { - hostname, port, err := host.ParseHostAndPortFromAddress(v.Address) + hostname, port, err := netutil.ParseHostAndPortFromAddress(v.Address) if err != nil { log.Warn("Failed to parse store address", zap.Any("store", v)) continue } - _, statusPort, err := host.ParseHostAndPortFromAddress(v.StatusAddress) + _, statusPort, err := netutil.ParseHostAndPortFromAddress(v.StatusAddress) if err != nil { log.Warn("Failed to parse store status address", zap.Any("store", v)) continue diff --git a/pkg/utils/topology/tidb.go b/pkg/utils/topology/tidb.go index 979125e714..cd33dd7108 100644 --- a/pkg/utils/topology/tidb.go +++ b/pkg/utils/topology/tidb.go @@ -16,7 +16,7 @@ import ( "go.uber.org/zap" "github.com/pingcap/tidb-dashboard/pkg/utils/distro" - "github.com/pingcap/tidb-dashboard/pkg/utils/host" + "github.com/pingcap/tidb-dashboard/util/netutil" ) const tidbTopologyKeyPrefix = "/topology/tidb/" @@ -110,7 +110,7 @@ func parseTiDBInfo(address string, value []byte) (*TiDBInfo, error) { if err != nil { return nil, ErrInvalidTopologyData.Wrap(err, "%s info unmarshal failed", distro.Data("tidb")) } - hostname, port, err := host.ParseHostAndPortFromAddress(address) + hostname, port, err := netutil.ParseHostAndPortFromAddress(address) if err != nil { return nil, ErrInvalidTopologyData.Wrap(err, "%s info address parse failed", distro.Data("tidb")) } diff --git a/util/rest/fileswap/server_test.go b/util/rest/fileswap/server_test.go index 93d7fd6708..ac318d2b05 100644 --- a/util/rest/fileswap/server_test.go +++ b/util/rest/fileswap/server_test.go @@ -10,7 +10,11 @@ import ( "time" "github.com/gin-gonic/gin" + "github.com/joomcode/errorx" "github.com/stretchr/testify/require" + + "github.com/pingcap/tidb-dashboard/util/assertutil" + "github.com/pingcap/tidb-dashboard/util/rest" ) func (s *Handler) mustGetDownloadToken(t *testing.T, fileContent string, downloadFileName string, expireIn time.Duration) string { @@ -36,8 +40,8 @@ func TestDownload(t *testing.T) { require.Len(t, c.Errors, 0) require.Equal(t, http.StatusOK, r.Code) - require.Equal(t, r.Header().Get("Content-Disposition"), `attachment; filename="file.txt"`) - require.Equal(t, r.Header().Get("Content-Type"), `application/octet-stream`) + require.Equal(t, `attachment; filename="file.txt"`, r.Header().Get("Content-Disposition")) + require.Equal(t, `application/octet-stream`, r.Header().Get("Content-Type")) require.Equal(t, "foobar", r.Body.String()) // Download again @@ -47,6 +51,7 @@ func TestDownload(t *testing.T) { handler.HandleDownloadRequest(c) require.Len(t, c.Errors, 1) + require.True(t, errorx.IsOfType(c.Errors[0].Err, rest.ErrBadRequest)) require.Contains(t, c.Errors[0].Error(), "Download file not found") } @@ -61,6 +66,7 @@ func TestDownloadAnotherInstance(t *testing.T) { handler2.HandleDownloadRequest(c) require.Len(t, c.Errors, 1) + require.True(t, errorx.IsOfType(c.Errors[0].Err, rest.ErrBadRequest)) require.Contains(t, c.Errors[0].Error(), "Invalid download request") require.Contains(t, c.Errors[0].Error(), "download token is invalid") } @@ -78,6 +84,7 @@ func TestExpiredToken(t *testing.T) { handler.HandleDownloadRequest(c) require.Len(t, c.Errors, 1) + require.True(t, errorx.IsOfType(c.Errors[0].Err, rest.ErrBadRequest)) require.Contains(t, c.Errors[0].Error(), "Invalid download request") require.Contains(t, c.Errors[0].Error(), "download token is expired") } @@ -91,6 +98,34 @@ func TestNotAToken(t *testing.T) { handler.HandleDownloadRequest(c) require.Len(t, c.Errors, 1) + require.True(t, errorx.IsOfType(c.Errors[0].Err, rest.ErrBadRequest)) require.Contains(t, c.Errors[0].Error(), "Invalid download request") require.Contains(t, c.Errors[0].Error(), "download token is invalid") } + +func TestDownloadInMiddleware(t *testing.T) { + handler := New() + token := handler.mustGetDownloadToken(t, "abc", "myfile.bin", time.Second*5) + + engine := gin.New() + engine.Use(rest.ErrorHandlerFn()) + engine.GET("/download", handler.HandleDownloadRequest) + + // A normal request + r := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/download?token="+token, nil) + engine.ServeHTTP(r, req) + require.Equal(t, http.StatusOK, r.Code) + require.Equal(t, `attachment; filename="myfile.bin"`, r.Header().Get("Content-Disposition")) + require.Equal(t, `application/octet-stream`, r.Header().Get("Content-Type")) + require.Equal(t, "abc", r.Body.String()) + + // A request without token + r = httptest.NewRecorder() + req, _ = http.NewRequest("GET", "/download", nil) + engine.ServeHTTP(r, req) + require.Equal(t, http.StatusBadRequest, r.Code) + require.Equal(t, "", r.Header().Get("Content-Disposition")) + require.Equal(t, "application/json; charset=utf-8", r.Header().Get("Content-Type")) + assertutil.RequireJSONContains(t, r.Body.String(), `{"code":"common.bad_request", "error":true}`) +} From c6aa25389e3309f73eb39ef2de0eed2e2e78510b Mon Sep 17 00:00:00 2001 From: Breezewish Date: Tue, 23 Nov 2021 22:51:08 +0800 Subject: [PATCH 2/2] Fix lint --- pkg/apiserver/debugapi/service.go | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/pkg/apiserver/debugapi/service.go b/pkg/apiserver/debugapi/service.go index 5a6d704d38..89377a0d9f 100644 --- a/pkg/apiserver/debugapi/service.go +++ b/pkg/apiserver/debugapi/service.go @@ -18,10 +18,6 @@ import ( "github.com/pingcap/tidb-dashboard/util/rest/fileswap" ) -const ( - tokenIssuer = "debugAPI" -) - func registerRouter(r *gin.RouterGroup, auth *user.AuthService, s *Service) { ep := r.Group("/debug_api") ep.GET("/download", s.Download) @@ -34,15 +30,13 @@ func registerRouter(r *gin.RouterGroup, auth *user.AuthService, s *Service) { type Service struct { Client *endpoint.Client - - fSwap *fileswap.Handler + fSwap *fileswap.Handler } func newService(hp httpClientParam) *Service { return &Service{ Client: endpoint.NewClient(newHTTPClient(hp), endpointDefs), - - fSwap: fileswap.New(), + fSwap: fileswap.New(), } }