swagger: block swagger url if disbale swagger server (#6785)

close #6786 /swagger/* return 404 Signed-off-by: bufferflies <[email protected]>
tikv · Jul 12, 2023 · c07c333 · c07c333
1 parent 2b9c9d9
commit c07c333
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 38 deletions.
diff --git a/cmd/pd-server/main.go b/cmd/pd-server/main.go
@@ -212,7 +212,10 @@ func start(cmd *cobra.Command, args []string, services ...string) {
 
 	// Creates server.
 	ctx, cancel := context.WithCancel(context.Background())
-	serviceBuilders := []server.HandlerBuilder{api.NewHandler, apiv2.NewV2Handler, swaggerserver.NewHandler, autoscaling.NewHandler}
+	serviceBuilders := []server.HandlerBuilder{api.NewHandler, apiv2.NewV2Handler, autoscaling.NewHandler}
+	if swaggerserver.Enabled() {
+		serviceBuilders = append(serviceBuilders, swaggerserver.NewHandler)
+	}
 	serviceBuilders = append(serviceBuilders, dashboard.GetServiceBuilders()...)
 	svr, err := server.CreateServer(ctx, cfg, services, serviceBuilders...)
 	if err != nil {

diff --git a/pkg/swaggerserver/swagger_handler.go b/pkg/swaggerserver/swagger_handler.go
diff --git a/pkg/swaggerserver/swaggerserver.go b/pkg/swaggerserver/swaggerserver.go
@@ -12,12 +12,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+//go:build swagger_server
+// +build swagger_server
+
 package swaggerserver
 
 import (
 	"context"
 	"net/http"
 
+	httpSwagger "github.com/swaggo/http-swagger"
+	_ "github.com/tikv/pd/docs/swagger"
 	"github.com/tikv/pd/pkg/utils/apiutil"
 	"github.com/tikv/pd/server"
 )
@@ -33,9 +38,14 @@ var (
 	}
 )
 
+// Enabled return true if swagger server is disabled.
+func Enabled() bool {
+	return true
+}
+
 // NewHandler creates a HTTP handler for Swagger.
 func NewHandler(context.Context, *server.Server) (http.Handler, apiutil.APIServiceGroup, error) {
 	swaggerHandler := http.NewServeMux()
-	swaggerHandler.Handle(swaggerPrefix, handler())
+	swaggerHandler.Handle(swaggerPrefix, httpSwagger.Handler())
 	return swaggerHandler, swaggerServiceGroup, nil
 }
diff --git a/pkg/swaggerserver/empty_handler.go → pkg/swaggerserver/swaggerserver_disable.go b/pkg/swaggerserver/empty_handler.go → pkg/swaggerserver/swaggerserver_disable.go
@@ -1,4 +1,4 @@
-// Copyright 2020 TiKV Project Authors.
+// Copyright 2023 TiKV Project Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -18,12 +18,19 @@
 package swaggerserver
 
 import (
-	"io"
+	"context"
 	"net/http"
+
+	"github.com/tikv/pd/pkg/utils/apiutil"
+	"github.com/tikv/pd/server"
 )
 
-func handler() http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		_, _ = io.WriteString(w, "Swagger UI is not built. Try `make` with `SWAGGER=1`.\n")
-	})
+// Enabled return false if swagger server is disabled.
+func Enabled() bool {
+	return false
+}
+
+// NewHandler creates a HTTP handler for Swagger.
+func NewHandler(context.Context, *server.Server) (http.Handler, apiutil.APIServiceGroup, error) {
+	return nil, apiutil.APIServiceGroup{}, nil
 }
diff --git a/tests/cluster.go b/tests/cluster.go
@@ -97,7 +97,10 @@ func createTestServer(ctx context.Context, cfg *config.Config, services []string
 	if err != nil {
 		return nil, err
 	}
-	serviceBuilders := []server.HandlerBuilder{api.NewHandler, apiv2.NewV2Handler, swaggerserver.NewHandler, autoscaling.NewHandler}
+	serviceBuilders := []server.HandlerBuilder{api.NewHandler, apiv2.NewV2Handler, autoscaling.NewHandler}
+	if swaggerserver.Enabled() {
+		serviceBuilders = append(serviceBuilders, swaggerserver.NewHandler)
+	}
 	serviceBuilders = append(serviceBuilders, dashboard.GetServiceBuilders()...)
 	svr, err := server.CreateServer(ctx, cfg, services, serviceBuilders...)
 	if err != nil {

diff --git a/tests/server/api/api_test.go b/tests/server/api/api_test.go
@@ -377,6 +377,15 @@ func (suite *middlewareTestSuite) TestRateLimitMiddleware() {
 	}
 }
 
+func (suite *middlewareTestSuite) TestSwaggerUrl() {
+	leader := suite.cluster.GetServer(suite.cluster.GetLeader())
+	req, _ := http.NewRequest(http.MethodGet, leader.GetAddr()+"/swagger/ui/index", nil)
+	resp, err := dialClient.Do(req)
+	suite.NoError(err)
+	suite.True(resp.StatusCode == http.StatusNotFound)
+	resp.Body.Close()
+}
+
 func (suite *middlewareTestSuite) TestAuditPrometheusBackend() {
 	leader := suite.cluster.GetServer(suite.cluster.GetLeader())
 	input := map[string]interface{}{