Skip to content

Commit ca5e3e1

Browse files
CabinfeverBnolouch
authored and
disksing
committed
api: Add Request Info Middleware (tikv#4526)
* close tikv#4494 Signed-off-by: Cabinfever_B <[email protected]> * close tikv#4494: add priority comment Signed-off-by: Cabinfever_B <[email protected]> * add request info middleware Signed-off-by: Cabinfever_B <[email protected]> * close tikv#4494 Signed-off-by: Cabinfever_B <[email protected]> * change service label getter and setter Signed-off-by: Cabinfever_B <[email protected]> * add lock Signed-off-by: Cabinfever_B <[email protected]> * add audit middleware Signed-off-by: Cabinfever_B <[email protected]> * add audit middleware Signed-off-by: Cabinfever_B <[email protected]> * tikv#4538 Signed-off-by: Cabinfever_B <[email protected]> * fix Signed-off-by: Cabinfever_B <[email protected]> * fix check Signed-off-by: Cabinfever_B <[email protected]> * fix check Signed-off-by: Cabinfever_B <[email protected]> * change service label placement Signed-off-by: Cabinfever_B <[email protected]> * merge master Signed-off-by: Cabinfever_B <[email protected]> * add const service label and could be dynamically turned on and off service middleware Signed-off-by: Cabinfever_B <[email protected]> * fix statics check Signed-off-by: Cabinfever_B <[email protected]> * fix statics check Signed-off-by: Cabinfever_B <[email protected]> * fix statics check Signed-off-by: Cabinfever_B <[email protected]> * fix statics check Signed-off-by: Cabinfever_B <[email protected]> * fix function name Signed-off-by: Cabinfever_B <[email protected]> * add benchmark test Signed-off-by: Cabinfever_B <[email protected]> * fix r.Body nil panic Signed-off-by: Cabinfever_B <[email protected]> * add benchmark test Signed-off-by: Cabinfever_B <[email protected]> * change export Signed-off-by: Cabinfever_B <[email protected]> * change export Signed-off-by: Cabinfever_B <[email protected]> * fix statics check Signed-off-by: Cabinfever_B <[email protected]> * fix typo problem Signed-off-by: Cabinfever_B <[email protected]> * change service label method Signed-off-by: Cabinfever_B <[email protected]> * fix statics Signed-off-by: Cabinfever_B <[email protected]> * fix router Signed-off-by: Cabinfever_B <[email protected]> * fix router Signed-off-by: Cabinfever_B <[email protected]> * add http proto Signed-off-by: Cabinfever_B <[email protected]> * add http proto Signed-off-by: Cabinfever_B <[email protected]> * change default value Signed-off-by: Cabinfever_B <[email protected]> * use middleware func shortname Signed-off-by: Cabinfever_B <[email protected]> * fix comment Signed-off-by: Cabinfever_B <[email protected]> * change const to iota Signed-off-by: Cabinfever_B <[email protected]> * change register method Signed-off-by: Cabinfever_B <[email protected]> * fix Signed-off-by: Cabinfever_B <[email protected]> * fix Signed-off-by: Cabinfever_B <[email protected]> * for test Signed-off-by: Cabinfever_B <[email protected]> * for test Signed-off-by: Cabinfever_B <[email protected]> * fix race Signed-off-by: Cabinfever_B <[email protected]> * fix race Signed-off-by: Cabinfever_B <[email protected]> Co-authored-by: ShuNing <[email protected]>
1 parent 9df9661 commit ca5e3e1

File tree

11 files changed

+614
-141
lines changed

11 files changed

+614
-141
lines changed

‎pkg/apiutil/apiutil.go

+10
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323
"strconv"
2424
"strings"
2525

26+
"github.com/gorilla/mux"
2627
"github.com/pingcap/errcode"
2728
"github.com/pingcap/errors"
2829
"github.com/pingcap/log"
@@ -190,3 +191,12 @@ func (rt *ComponentSignatureRoundTripper) RoundTrip(req *http.Request) (resp *ht
190191
resp, err = rt.proxied.RoundTrip(req)
191192
return
192193
}
194+
195+
// GetRouteName return mux route name registered
196+
func GetRouteName(req *http.Request) string {
197+
route := mux.CurrentRoute(req)
198+
if route != nil {
199+
return route.GetName()
200+
}
201+
return ""
202+
}

‎pkg/requestutil/context.go

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
// Copyright 2022 TiKV Project Authors.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package requestutil
16+
17+
import (
18+
"context"
19+
)
20+
21+
// The key type is unexported to prevent collisions
22+
type key int
23+
24+
const (
25+
// requestInfoKey is the context key for the request compoenent.
26+
requestInfoKey key = iota
27+
)
28+
29+
// WithRequestInfo returns a copy of parent in which the request info value is set
30+
func WithRequestInfo(parent context.Context, requestInfo RequestInfo) context.Context {
31+
return context.WithValue(parent, requestInfoKey, requestInfo)
32+
}
33+
34+
// RequestInfoFrom returns the value of the request info key on the ctx
35+
func RequestInfoFrom(ctx context.Context) (RequestInfo, bool) {
36+
requestInfo, ok := ctx.Value(requestInfoKey).(RequestInfo)
37+
return requestInfo, ok
38+
}

‎pkg/requestutil/context_test.go

+57
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
// Copyright 2022 TiKV Project Authors.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package requestutil
16+
17+
import (
18+
"context"
19+
"testing"
20+
21+
. "github.com/pingcap/check"
22+
)
23+
24+
func Test(t *testing.T) {
25+
TestingT(t)
26+
}
27+
28+
var _ = Suite(&testRequestContextSuite{})
29+
30+
type testRequestContextSuite struct {
31+
}
32+
33+
func (s *testRequestContextSuite) TestRequestInfo(c *C) {
34+
ctx := context.Background()
35+
_, ok := RequestInfoFrom(ctx)
36+
c.Assert(ok, Equals, false)
37+
ctx = WithRequestInfo(ctx,
38+
RequestInfo{
39+
ServiceLabel: "test label",
40+
Method: "POST",
41+
Component: "pdctl",
42+
IP: "localhost",
43+
URLParam: "{\"id\"=1}",
44+
BodyParam: "{\"state\"=\"Up\"}",
45+
TimeStamp: "2022",
46+
})
47+
result, ok := RequestInfoFrom(ctx)
48+
c.Assert(result, NotNil)
49+
c.Assert(ok, Equals, true)
50+
c.Assert(result.ServiceLabel, Equals, "test label")
51+
c.Assert(result.Method, Equals, "POST")
52+
c.Assert(result.Component, Equals, "pdctl")
53+
c.Assert(result.IP, Equals, "localhost")
54+
c.Assert(result.URLParam, Equals, "{\"id\"=1}")
55+
c.Assert(result.BodyParam, Equals, "{\"state\"=\"Up\"}")
56+
c.Assert(result.TimeStamp, Equals, "2022")
57+
}

‎pkg/requestutil/request_info.go

+70
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
// Copyright 2022 TiKV Project Authors.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package requestutil
16+
17+
import (
18+
"bytes"
19+
"encoding/json"
20+
"fmt"
21+
"io"
22+
"net/http"
23+
"time"
24+
25+
"github.com/tikv/pd/pkg/apiutil"
26+
)
27+
28+
// RequestInfo holds service information from http.Request
29+
type RequestInfo struct {
30+
ServiceLabel string
31+
Method string
32+
Component string
33+
IP string
34+
TimeStamp string
35+
URLParam string
36+
BodyParam string
37+
}
38+
39+
// GetRequestInfo returns request info needed from http.Request
40+
func GetRequestInfo(r *http.Request) RequestInfo {
41+
return RequestInfo{
42+
ServiceLabel: apiutil.GetRouteName(r),
43+
Method: fmt.Sprintf("%s/%s:%s", r.Proto, r.Method, r.URL.Path),
44+
Component: apiutil.GetComponentNameOnHTTP(r),
45+
IP: apiutil.GetIPAddrFromHTTPRequest(r),
46+
TimeStamp: time.Now().Local().String(),
47+
URLParam: getURLParam(r),
48+
BodyParam: getBodyParam(r),
49+
}
50+
}
51+
52+
func getURLParam(r *http.Request) string {
53+
buf, err := json.Marshal(r.URL.Query())
54+
if err != nil {
55+
return ""
56+
}
57+
return string(buf)
58+
}
59+
60+
func getBodyParam(r *http.Request) string {
61+
if r.Body == nil {
62+
return ""
63+
}
64+
// http request body is a io.Reader between bytes.Reader and strings.Reader, it only has EOF error
65+
buf, _ := io.ReadAll(r.Body)
66+
r.Body.Close()
67+
bodyParam := string(buf)
68+
r.Body = io.NopCloser(bytes.NewBuffer(buf))
69+
return bodyParam
70+
}

‎server/api/admin.go

+18
Original file line numberDiff line numberDiff line change
@@ -138,3 +138,21 @@ func (h *adminHandler) UpdateWaitAsyncTime(w http.ResponseWriter, r *http.Reques
138138
cluster.GetReplicationMode().UpdateMemberWaitAsyncTime(memberID)
139139
h.rd.JSON(w, http.StatusOK, nil)
140140
}
141+
142+
// @Tags admin
143+
// @Summary switch Service Middlewares including ServiceInfo, Audit and rate limit
144+
// @Param enable query string true "enable" Enums(true, false)
145+
// @Produce json
146+
// @Success 200 {string} string "Switching Service middleware is successful."
147+
// @Failure 400 {string} string "The input is invalid."
148+
// @Router /admin/service-middleware [POST]
149+
func (h *adminHandler) HanldeServiceMiddlewareSwitch(w http.ResponseWriter, r *http.Request) {
150+
enableStr := r.URL.Query().Get("enable")
151+
enable, err := strconv.ParseBool(enableStr)
152+
if err != nil {
153+
h.rd.JSON(w, http.StatusBadRequest, "The input is invalid.")
154+
return
155+
}
156+
h.svr.SetServiceMiddleware(enable)
157+
h.rd.JSON(w, http.StatusOK, "Switching Service middleware is successful.")
158+
}

‎server/api/admin_test.go

+42
Original file line numberDiff line numberDiff line change
@@ -173,3 +173,45 @@ func (s *testTSOSuite) TestResetTS(c *C) {
173173
c.Assert(err, NotNil)
174174
c.Assert(err.Error(), Equals, "\"invalid tso value\"\n")
175175
}
176+
177+
var _ = Suite(&testServiceSuite{})
178+
179+
type testServiceSuite struct {
180+
svr *server.Server
181+
cleanup cleanUpFunc
182+
}
183+
184+
func (s *testServiceSuite) SetUpSuite(c *C) {
185+
s.svr, s.cleanup = mustNewServer(c)
186+
mustWaitLeader(c, []*server.Server{s.svr})
187+
188+
mustBootstrapCluster(c, s.svr)
189+
mustPutStore(c, s.svr, 1, metapb.StoreState_Up, nil)
190+
}
191+
192+
func (s *testServiceSuite) TearDownSuite(c *C) {
193+
s.cleanup()
194+
}
195+
196+
func (s *testServiceSuite) TestSwitchServiceMiddleware(c *C) {
197+
urlPrefix := fmt.Sprintf("%s%s/api/v1/admin/service-middleware", s.svr.GetAddr(), apiPrefix)
198+
disableURL := fmt.Sprintf("%s?enable=false", urlPrefix)
199+
err := postJSON(testDialClient, disableURL, nil,
200+
func(res []byte, code int) {
201+
c.Assert(string(res), Equals, "\"Switching Service middleware is successful.\"\n")
202+
c.Assert(code, Equals, http.StatusOK)
203+
})
204+
205+
c.Assert(err, IsNil)
206+
c.Assert(s.svr.IsServiceMiddlewareEnabled(), Equals, false)
207+
208+
enableURL := fmt.Sprintf("%s?enable=true", urlPrefix)
209+
err = postJSON(testDialClient, enableURL, nil,
210+
func(res []byte, code int) {
211+
c.Assert(string(res), Equals, "\"Switching Service middleware is successful.\"\n")
212+
c.Assert(code, Equals, http.StatusOK)
213+
})
214+
215+
c.Assert(err, IsNil)
216+
c.Assert(s.svr.IsServiceMiddlewareEnabled(), Equals, true)
217+
}

‎server/api/middleware.go

+33
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,45 @@ import (
1818
"context"
1919
"net/http"
2020

21+
"github.com/pingcap/failpoint"
2122
"github.com/tikv/pd/pkg/errs"
23+
"github.com/tikv/pd/pkg/requestutil"
2224
"github.com/tikv/pd/server"
2325
"github.com/tikv/pd/server/cluster"
2426
"github.com/unrolled/render"
27+
"github.com/urfave/negroni"
2528
)
2629

30+
// requestInfoMiddleware is used to gather info from requsetInfo
31+
type requestInfoMiddleware struct {
32+
svr *server.Server
33+
}
34+
35+
func newRequestInfoMiddleware(s *server.Server) negroni.Handler {
36+
return &requestInfoMiddleware{svr: s}
37+
}
38+
39+
func (rm *requestInfoMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
40+
if !rm.svr.IsServiceMiddlewareEnabled() {
41+
next(w, r)
42+
return
43+
}
44+
45+
requestInfo := requestutil.GetRequestInfo(r)
46+
r = r.WithContext(requestutil.WithRequestInfo(r.Context(), requestInfo))
47+
48+
failpoint.Inject("addRequestInfoMiddleware", func() {
49+
w.Header().Add("service-label", requestInfo.ServiceLabel)
50+
w.Header().Add("body-param", requestInfo.BodyParam)
51+
w.Header().Add("url-param", requestInfo.URLParam)
52+
w.Header().Add("method", requestInfo.Method)
53+
w.Header().Add("component", requestInfo.Component)
54+
w.Header().Add("ip", requestInfo.IP)
55+
})
56+
57+
next(w, r)
58+
}
59+
2760
type clusterMiddleware struct {
2861
s *server.Server
2962
rd *render.Render

0 commit comments

Comments
 (0)