@@ -15,6 +15,7 @@ import (
1515 "github.com/hashicorp/go-multierror"
1616 "github.com/hashicorp/vault/helper/namespace"
1717 "github.com/hashicorp/vault/limits"
18+ "github.com/hashicorp/vault/sdk/helper/jsonutil"
1819 "github.com/hashicorp/vault/sdk/logical"
1920 "github.com/hashicorp/vault/vault"
2021 "github.com/hashicorp/vault/vault/quotas"
@@ -24,25 +25,80 @@ var nonVotersAllowed = false
2425
2526func wrapMaxRequestSizeHandler (handler http.Handler , props * vault.HandlerProperties ) http.Handler {
2627 return http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
27- var maxRequestSize int64
28+ var maxRequestSize , maxJSONDepth , maxStringValueLength , maxObjectEntryCount , maxArrayElementCount int64
29+
2830 if props .ListenerConfig != nil {
2931 maxRequestSize = props .ListenerConfig .MaxRequestSize
32+ maxJSONDepth = props .ListenerConfig .CustomMaxJSONDepth
33+ maxStringValueLength = props .ListenerConfig .CustomMaxJSONStringValueLength
34+ maxObjectEntryCount = props .ListenerConfig .CustomMaxJSONObjectEntryCount
35+ maxArrayElementCount = props .ListenerConfig .CustomMaxJSONArrayElementCount
3036 }
37+
3138 if maxRequestSize == 0 {
3239 maxRequestSize = DefaultMaxRequestSize
3340 }
34- ctx := r .Context ()
35- originalBody := r .Body
41+ if maxJSONDepth == 0 {
42+ maxJSONDepth = CustomMaxJSONDepth
43+ }
44+ if maxStringValueLength == 0 {
45+ maxStringValueLength = CustomMaxJSONStringValueLength
46+ }
47+ if maxObjectEntryCount == 0 {
48+ maxObjectEntryCount = CustomMaxJSONObjectEntryCount
49+ }
50+ if maxArrayElementCount == 0 {
51+ maxArrayElementCount = CustomMaxJSONArrayElementCount
52+ }
53+
54+ jsonLimits := jsonutil.JSONLimits {
55+ MaxDepth : int (maxJSONDepth ),
56+ MaxStringValueLength : int (maxStringValueLength ),
57+ MaxObjectEntryCount : int (maxObjectEntryCount ),
58+ MaxArrayElementCount : int (maxArrayElementCount ),
59+ }
60+
61+ // If the payload is JSON, the VerifyMaxDepthStreaming function will perform validations.
62+ buf , err := jsonLimitsValidation (w , r , maxRequestSize , jsonLimits )
63+ if err != nil {
64+ respondError (w , http .StatusInternalServerError , err )
65+ return
66+ }
67+
68+ // Replace the body and update the context.
69+ // This ensures the request object is in a consistent state for all downstream handlers.
70+ // Because the original request body stream has been fully consumed by io.ReadAll,
71+ // we must replace it so that subsequent handlers can read the content.
72+ r .Body = newMultiReaderCloser (buf , r .Body )
73+ contextBody := r .Body
74+ ctx := logical .CreateContextOriginalBody (r .Context (), contextBody )
75+
3676 if maxRequestSize > 0 {
3777 r .Body = http .MaxBytesReader (w , r .Body , maxRequestSize )
3878 }
39- ctx = logical .CreateContextOriginalBody (ctx , originalBody )
4079 r = r .WithContext (ctx )
4180
4281 handler .ServeHTTP (w , r )
4382 })
4483}
4584
85+ func jsonLimitsValidation (w http.ResponseWriter , r * http.Request , maxRequestSize int64 , jsonLimits jsonutil.JSONLimits ) (* bytes.Buffer , error ) {
86+ // The TeeReader reads from the original body and writes a copy to our buffer.
87+ // We wrap the original body with a MaxBytesReader first to enforce the hard size limit.
88+ var limitedTeeReader io.Reader
89+ buf := & bytes.Buffer {}
90+ bodyReader := r .Body
91+ if maxRequestSize > 0 {
92+ bodyReader = http .MaxBytesReader (w , r .Body , maxRequestSize )
93+ }
94+ limitedTeeReader = io .TeeReader (bodyReader , buf )
95+ _ , err := jsonutil .VerifyMaxDepthStreaming (limitedTeeReader , jsonLimits )
96+ if err != nil {
97+ return nil , err
98+ }
99+ return buf , nil
100+ }
101+
46102func wrapRequestLimiterHandler (handler http.Handler , props * vault.HandlerProperties ) http.Handler {
47103 return http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
48104 request := r .WithContext (
0 commit comments