diff --git a/pkg/mcs/resourcemanager/metadataapi/config_service.go b/pkg/mcs/resourcemanager/metadataapi/config_service.go index 11fa31ee4e6..07bb8bb78d9 100644 --- a/pkg/mcs/resourcemanager/metadataapi/config_service.go +++ b/pkg/mcs/resourcemanager/metadataapi/config_service.go @@ -181,13 +181,19 @@ func decodeResourceGroup(body io.Reader, group *rmpb.ResourceGroup) error { if err != nil { return err } - legacyJSON, rawKeyspaceID, err := splitResourceGroupJSON(data) - if err != nil { - return err + var keyspaceID resourceGroupKeyspaceIDJSON + legacyGroup := resourceGroupJSON{ + resourceGroupWithoutKeyspaceID: (*resourceGroupWithoutKeyspaceID)(group), + KeyspaceID: resourceGroupKeyspaceIDField{value: &keyspaceID}, + KeyspaceIDCamel: resourceGroupKeyspaceIDField{value: &keyspaceID}, } // Keep the legacy encoding/json behavior for all existing ResourceGroup // fields. In particular, it matches JSON field names case-insensitively. - if err := json.Unmarshal(legacyJSON, group); err != nil { + if err := json.Unmarshal(data, &legacyGroup); err != nil { + var keyspaceIDErr *resourceGroupKeyspaceIDError + if errors.As(err, &keyspaceIDErr) { + return keyspaceIDErr + } // The updated ResourceGroup contains a protobuf oneof, so clients may // serialize the whole message as protobuf JSON. Retry strictly to // accept enum names and quoted 64-bit integers without silently @@ -196,44 +202,53 @@ func decodeResourceGroup(body io.Reader, group *rmpb.ResourceGroup) error { if protoErr := (&jsonpb.Unmarshaler{}).Unmarshal(bytes.NewReader(data), group); protoErr != nil { return fmt.Errorf("invalid resource group JSON: legacy JSON: %v; protobuf JSON: %w", err, protoErr) } - return validateResourceGroupKeyspaceID(group, rawKeyspaceID) + return validateResourceGroupKeyspaceID(group, keyspaceID.raw) } - if rawKeyspaceID != nil { - keyspaceID, err := decodeKeyspaceIDJSON(rawKeyspaceID) + if keyspaceID.raw != nil { + decodedKeyspaceID, err := decodeKeyspaceIDJSON(keyspaceID.raw) if err != nil { return err } - group.KeyspaceId = keyspaceID + group.KeyspaceId = decodedKeyspaceID } - return validateResourceGroupKeyspaceID(group, rawKeyspaceID) + return validateResourceGroupKeyspaceID(group, keyspaceID.raw) } -func splitResourceGroupJSON(data []byte) ([]byte, json.RawMessage, error) { - if isJSONNull(data) { - return data, nil, nil +// resourceGroupWithoutKeyspaceID lets encoding/json decode every legacy field +// directly. The explicit fields in resourceGroupJSON shadow KeyspaceId, whose +// protobuf oneof still requires compatibility handling. +type resourceGroupWithoutKeyspaceID rmpb.ResourceGroup + +type resourceGroupJSON struct { + *resourceGroupWithoutKeyspaceID + KeyspaceID resourceGroupKeyspaceIDField `json:"keyspace_id"` + KeyspaceIDCamel resourceGroupKeyspaceIDField `json:"keyspaceId"` +} + +type resourceGroupKeyspaceIDJSON struct { + raw json.RawMessage +} + +type resourceGroupKeyspaceIDField struct { + value *resourceGroupKeyspaceIDJSON +} + +// UnmarshalJSON records and normalizes a ResourceGroup keyspace ID without +// making a second copy of the rest of the request body. +func (f *resourceGroupKeyspaceIDField) UnmarshalJSON(data []byte) error { + if f.value.raw != nil { + return &resourceGroupKeyspaceIDError{errors.New("keyspace_id must be set only once")} } - // KeyspaceIDValue became a protobuf oneof, which encoding/json cannot decode. - // Remove it from the legacy payload and decode it separately with jsonpb. - fields, err := decodeJSONObjectFields(data) + raw, err := normalizeKeyspaceIDJSON(data) if err != nil { - return nil, nil, err - } - legacyFields := make([]jsonObjectField, 0, len(fields)) - var rawKeyspaceID json.RawMessage - for _, field := range fields { - if !isKeyspaceIDJSONField(field.name) { - legacyFields = append(legacyFields, field) - continue - } - if rawKeyspaceID != nil { - return nil, nil, errors.New("keyspace_id must be set only once") - } - rawKeyspaceID, err = normalizeKeyspaceIDJSON(field.value) - if err != nil { - return nil, nil, err - } + return &resourceGroupKeyspaceIDError{err} } - return marshalJSONObjectFields(legacyFields), rawKeyspaceID, nil + f.value.raw = raw + return nil +} + +type resourceGroupKeyspaceIDError struct { + error } type jsonObjectField struct { @@ -241,22 +256,6 @@ type jsonObjectField struct { value json.RawMessage } -func marshalJSONObjectFields(fields []jsonObjectField) []byte { - var buffer bytes.Buffer - buffer.WriteByte('{') - for i, field := range fields { - if i > 0 { - buffer.WriteByte(',') - } - name, _ := json.Marshal(field.name) - buffer.Write(name) - buffer.WriteByte(':') - buffer.Write(field.value) - } - buffer.WriteByte('}') - return buffer.Bytes() -} - func decodeJSONObjectFields(data []byte) ([]jsonObjectField, error) { var object map[string]json.RawMessage if err := json.Unmarshal(data, &object); err != nil { @@ -293,10 +292,6 @@ func decodeJSONObjectFields(data []byte) ([]jsonObjectField, error) { return fields, nil } -func isKeyspaceIDJSONField(name string) bool { - return strings.EqualFold(name, "keyspace_id") || strings.EqualFold(name, "keyspaceId") -} - func isJSONNull(data []byte) bool { return bytes.Equal(bytes.TrimSpace(data), []byte("null")) } diff --git a/pkg/mcs/resourcemanager/metadataapi/config_service_test.go b/pkg/mcs/resourcemanager/metadataapi/config_service_test.go index 1db1ca54b4f..0929d1d4826 100644 --- a/pkg/mcs/resourcemanager/metadataapi/config_service_test.go +++ b/pkg/mcs/resourcemanager/metadataapi/config_service_test.go @@ -22,6 +22,7 @@ import ( "fmt" "net/http" "net/http/httptest" + "strings" "testing" "github.com/gin-gonic/gin" @@ -162,6 +163,22 @@ func TestConfigServiceGroupCRUDAndErrorCodes(t *testing.T) { re.Equal(int64(987), storedGroup.RUSettings.RU.Settings.BurstLimit) } + largeJobType := strings.Repeat("x", 1<<20) + largeResourceGroupBody, err := json.Marshal(&rmpb.ResourceGroup{ + Name: "large_resource_group", + BackgroundSettings: &rmpb.BackgroundSettings{ + JobTypes: []string{largeJobType}, + }, + }) + re.NoError(err) + re.Greater(len(largeResourceGroupBody), 1<<20) + for _, method := range []string{http.MethodPost, http.MethodPut} { + resp = doRawResourceGroupRequest(handler, method, largeResourceGroupBody) + re.Equal(http.StatusOK, resp.Code, resp.Body.String()) + largeResourceGroup := store.groups[groupKey(constant.NullKeyspaceID, "large_resource_group")] + re.Equal(largeJobType, largeResourceGroup.Background.JobTypes[0]) + } + store.addErr = errors.New("add failed") resp = doJSONRequest(re, handler, http.MethodPost, "/resource-manager/api/v1/config/group", group) re.Equal(http.StatusInternalServerError, resp.Code) diff --git a/pkg/utils/requestutil/request_info.go b/pkg/utils/requestutil/request_info.go index cc5403f7232..fd2b0feeac6 100644 --- a/pkg/utils/requestutil/request_info.go +++ b/pkg/utils/requestutil/request_info.go @@ -15,11 +15,11 @@ package requestutil import ( - "bytes" "encoding/json" "fmt" "io" "net/http" + "strings" "time" "github.com/tikv/pd/pkg/utils/apiutil" @@ -47,6 +47,14 @@ func (info *RequestInfo) String() string { // GetRequestInfo returns request info needed from http.Request func GetRequestInfo(r *http.Request) RequestInfo { + info := GetRequestInfoWithoutBody(r) + info.CaptureBody(r) + return info +} + +// GetRequestInfoWithoutBody returns request info without consuming or buffering +// the request body. +func GetRequestInfoWithoutBody(r *http.Request) RequestInfo { ip, port := apiutil.GetIPPortFromHTTPRequest(r) return RequestInfo{ ServiceLabel: apiutil.GetRouteName(r), @@ -55,11 +63,16 @@ func GetRequestInfo(r *http.Request) RequestInfo { IP: ip, Port: port, URLParam: getURLParam(r), - BodyParam: getBodyParam(r), StartTimeStamp: time.Now().Unix(), } } +// CaptureBody consumes the request body into BodyParam and restores an +// equivalent body for the handler. +func (info *RequestInfo) CaptureBody(r *http.Request) { + info.BodyParam = getBodyParam(r) +} + func getURLParam(r *http.Request) string { buf, err := json.Marshal(r.URL.Query()) if err != nil { @@ -72,10 +85,12 @@ func getBodyParam(r *http.Request) string { if r.Body == nil { return "" } - // http request body is a io.Reader between bytes.Reader and strings.Reader, it only has EOF error buf, _ := io.ReadAll(r.Body) r.Body.Close() + // Restore the body from BodyParam so both views share one long-lived backing + // store. Restoring from buf would retain a second complete copy until the + // handler finishes. bodyParam := string(buf) - r.Body = io.NopCloser(bytes.NewBuffer(buf)) + r.Body = io.NopCloser(strings.NewReader(bodyParam)) return bodyParam } diff --git a/pkg/utils/requestutil/request_info_test.go b/pkg/utils/requestutil/request_info_test.go new file mode 100644 index 00000000000..fab99f195e0 --- /dev/null +++ b/pkg/utils/requestutil/request_info_test.go @@ -0,0 +1,73 @@ +// Copyright 2026 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. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package requestutil + +import ( + "io" + "net/http" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +type trackingReadCloser struct { + io.Reader + read bool + close bool +} + +func (r *trackingReadCloser) Read(p []byte) (int, error) { + r.read = true + return r.Reader.Read(p) +} + +func (r *trackingReadCloser) Close() error { + r.close = true + return nil +} + +func TestGetRequestInfoWithoutBodyDoesNotConsumeBody(t *testing.T) { + re := require.New(t) + body := &trackingReadCloser{Reader: strings.NewReader("request-body")} + req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1/test", body) + re.NoError(err) + + info := GetRequestInfoWithoutBody(req) + re.Empty(info.BodyParam) + re.False(body.read) + re.False(body.close) + + data, err := io.ReadAll(req.Body) + re.NoError(err) + re.Equal("request-body", string(data)) +} + +func TestGetRequestInfoRestoresExactBody(t *testing.T) { + re := require.New(t) + const requestBody = "request-body-\x00-\xff" + body := &trackingReadCloser{Reader: strings.NewReader(requestBody)} + req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1/test", body) + re.NoError(err) + + info := GetRequestInfo(req) + re.Equal(requestBody, info.BodyParam) + re.True(body.read) + re.True(body.close) + + data, err := io.ReadAll(req.Body) + re.NoError(err) + re.Equal([]byte(requestBody), data) +} diff --git a/server/api/middleware.go b/server/api/middleware.go index 203751e8a07..1c0a084cdec 100644 --- a/server/api/middleware.go +++ b/server/api/middleware.go @@ -17,6 +17,7 @@ package api import ( "context" "net/http" + "slices" "time" "github.com/unrolled/render" @@ -59,12 +60,13 @@ func newRequestInfoMiddleware(s *server.Server) negroni.Handler { } func (rm *requestInfoMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) { - if !rm.svr.GetServiceMiddlewarePersistOptions().IsAuditEnabled() && !rm.svr.GetServiceMiddlewarePersistOptions().IsRateLimitEnabled() { + if !rm.svr.GetServiceMiddlewarePersistOptions().IsAuditEnabled() && + !rm.svr.GetServiceMiddlewarePersistOptions().IsRateLimitEnabled() { next(w, r) return } - requestInfo := requestutil.GetRequestInfo(r) + requestInfo := requestutil.GetRequestInfoWithoutBody(r) r = r.WithContext(requestutil.WithRequestInfo(r.Context(), requestInfo)) failpoint.Inject("addRequestInfoMiddleware", func() { @@ -79,6 +81,38 @@ func (rm *requestInfoMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Reques next(w, r) } +func captureRequestBodyForAudit( + r *http.Request, + requestInfo *requestutil.RequestInfo, + labels *audit.BackendLabels, +) bool { + if requestInfo.BodyParam != "" || !auditNeedsRequestBody(labels) { + return false + } + requestInfo.CaptureBody(r) + return true +} + +func auditNeedsRequestBody(labels *audit.BackendLabels) bool { + return labels != nil && slices.Contains(labels.Labels, audit.LocalLogLabel) +} + +func prepareRequestForAudit( + r *http.Request, + requestInfo requestutil.RequestInfo, + hasRequestInfo bool, + labels *audit.BackendLabels, +) *http.Request { + contextNeedsUpdate := !hasRequestInfo + if captureRequestBodyForAudit(r, &requestInfo, labels) { + contextNeedsUpdate = true + } + if contextNeedsUpdate { + return r.WithContext(requestutil.WithRequestInfo(r.Context(), requestInfo)) + } + return r +} + type clusterMiddleware struct { s *server.Server rd *render.Render @@ -185,14 +219,14 @@ func (s *auditMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next requestInfo, ok := requestutil.RequestInfoFrom(r.Context()) if !ok { - requestInfo = requestutil.GetRequestInfo(r) + requestInfo = requestutil.GetRequestInfoWithoutBody(r) } - labels := s.svr.GetServiceAuditBackendLabels(requestInfo.ServiceLabel) if labels == nil { next(w, r) return } + r = prepareRequestForAudit(r, requestInfo, ok, labels) beforeNextBackends := make([]audit.Backend, 0) afterNextBackends := make([]audit.Backend, 0) @@ -234,7 +268,7 @@ func (s *rateLimitMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, } requestInfo, ok := requestutil.RequestInfoFrom(r.Context()) if !ok { - requestInfo = requestutil.GetRequestInfo(r) + requestInfo = requestutil.GetRequestInfoWithoutBody(r) } // There is no need to check whether rateLimiter is nil. CreateServer ensures that it is created diff --git a/server/api/middleware_test.go b/server/api/middleware_test.go new file mode 100644 index 00000000000..5df85044dfb --- /dev/null +++ b/server/api/middleware_test.go @@ -0,0 +1,101 @@ +// Copyright 2026 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. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package api + +import ( + "io" + "net/http" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/tikv/pd/pkg/audit" + "github.com/tikv/pd/pkg/utils/requestutil" +) + +func TestCaptureRequestBodyForAuditReadsBodyOnlyWhenNeeded(t *testing.T) { + testCases := []struct { + name string + labels *audit.BackendLabels + bodyParam string + expectBody string + expectRead bool + }{ + { + name: "prometheus-only", + labels: &audit.BackendLabels{Labels: []string{audit.PrometheusHistogram}}, + }, + { + name: "no-audit-backend", + }, + { + name: "local-log", + labels: &audit.BackendLabels{Labels: []string{audit.LocalLogLabel, audit.PrometheusHistogram}}, + expectBody: "request-body", + expectRead: true, + }, + { + name: "already-captured", + labels: &audit.BackendLabels{Labels: []string{audit.LocalLogLabel}}, + bodyParam: "captured-body", + expectBody: "captured-body", + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + re := require.New(t) + body := &trackingReadCloser{Reader: strings.NewReader("request-body")} + req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1/test", body) + re.NoError(err) + + info := requestutil.GetRequestInfoWithoutBody(req) + info.BodyParam = testCase.bodyParam + re.Equal(testCase.expectRead, captureRequestBodyForAudit(req, &info, testCase.labels)) + re.Equal(testCase.expectBody, info.BodyParam) + re.Equal(testCase.expectRead, body.reads > 0) + re.Equal(testCase.expectRead, body.closed) + + data, err := io.ReadAll(req.Body) + re.NoError(err) + re.Equal("request-body", string(data)) + }) + } +} + +func TestPrepareRequestForAuditCapturesBodyForExistingRequestInfo(t *testing.T) { + re := require.New(t) + body := &trackingReadCloser{Reader: strings.NewReader("request-body")} + req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1/test", body) + re.NoError(err) + + requestInfo := requestutil.GetRequestInfoWithoutBody(req) + req = req.WithContext(requestutil.WithRequestInfo(req.Context(), requestInfo)) + requestInfo, ok := requestutil.RequestInfoFrom(req.Context()) + re.True(ok) + + labels := &audit.BackendLabels{Labels: []string{audit.LocalLogLabel}} + req = prepareRequestForAudit(req, requestInfo, ok, labels) + requestInfo, ok = requestutil.RequestInfoFrom(req.Context()) + re.True(ok) + re.Equal("request-body", requestInfo.BodyParam) + re.Positive(body.reads) + re.True(body.closed) + + data, err := io.ReadAll(req.Body) + re.NoError(err) + re.Equal("request-body", string(data)) +} diff --git a/tests/server/api/api_test.go b/tests/server/api/api_test.go index bc595e09465..13eee6ca390 100644 --- a/tests/server/api/api_test.go +++ b/tests/server/api/api_test.go @@ -179,7 +179,7 @@ func (suite *middlewareTestSuite) TestRequestInfoMiddleware() { re.Equal("Profile", resp.Header.Get("service-label")) re.JSONEq("{\"seconds\":[\"1\"]}", resp.Header.Get("url-param")) - re.JSONEq("{\"testkey\":\"testvalue\"}", resp.Header.Get("body-param")) + re.Empty(resp.Header.Get("body-param")) re.Equal("HTTP/1.1/POST:/pd/api/v1/debug/pprof/profile", resp.Header.Get("method")) re.Equal("anonymous", resp.Header.Get("caller-id")) re.Equal("127.0.0.1", resp.Header.Get("ip"))