diff --git a/handler/perform_handler_with_json.go b/handler/perform_handler_with_json.go index cca657f27..19927939a 100644 --- a/handler/perform_handler_with_json.go +++ b/handler/perform_handler_with_json.go @@ -2,6 +2,7 @@ package handler import ( "encoding/json" + "unicode/utf8" bosherr "github.com/cloudfoundry/bosh-utils/errors" boshlog "github.com/cloudfoundry/bosh-utils/logger" @@ -10,6 +11,7 @@ import ( const ( mbusHandlerLogTag = "MBus Handler" responseMaxLengthErrMsg = "Response exceeded maximum allowed length" + truncationMarker = "\n...[truncated]...\n" UnlimitedResponseLength = -1 ) @@ -53,6 +55,102 @@ func BuildErrorWithJSON(msg string, logger boshlog.Logger) ([]byte, error) { return respJSON, nil } +// truncateExceptionResponse truncates the message of excResp so that its JSON +// serialization fits within maxLength bytes. It preserves both the beginning and +// the end of the message, inserting truncationMarker in between, so that the +// high-level error context and the root cause are both visible. +// Returns the truncated response and true if successful, or false if it cannot +// be made to fit. +func truncateExceptionResponse(excResp exceptionResponse, maxLength int) (exceptionResponse, bool) { + // Probe the fixed JSON overhead by marshaling just the truncationMarker as the + // message. Because the `message` field uses omitempty, marshaling an empty struct + // would omit the key entirely and undercount the overhead. + // e.g. {"exception":{"message":"...[truncated]..."}} → fixedOverhead = len(probe) - len(marker) + probe := exceptionResponse{} + probe.Exception.Message = truncationMarker + probeJSON, err := json.Marshal(probe) + if err != nil { + return exceptionResponse{}, false + } + + fixedOverhead := len(probeJSON) - len(truncationMarker) + maxMsgBytes := maxLength - fixedOverhead + if maxMsgBytes <= len(truncationMarker) { + return exceptionResponse{}, false + } + + msg := excResp.Exception.Message + + // Short-circuit only if the actual marshaled response already fits. We can't + // rely on len(msg) alone: JSON escaping (e.g. `"` → `\"`) can expand the + // message past the limit even when its raw byte length is within maxMsgBytes. + if fullJSON, err := json.Marshal(excResp); err == nil && len(fullJSON) <= maxLength { + return excResp, true + } + + // Split the available content bytes evenly between head and tail. + // Tail gets the extra byte when contentBytes is odd, since the root cause + // is typically found at the end of the error chain. + // + // We loop and re-marshal because json.Marshal expands escapable characters + // (e.g. `"` → `\"`, `\` → `\\`, non-ASCII → `\uXXXX`), meaning the raw + // byte estimate can undercount the actual JSON size. Each iteration scales + // contentBytes down proportionally until the marshaled result fits. + contentBytes := maxMsgBytes - len(truncationMarker) + + // Clamp to the message length so head/tail slicing never reads out of + // bounds. This matters when a short-but-heavily-escaped message reaches the + // loop: contentBytes can exceed len(msg), which would otherwise produce a + // negative tail index. + if contentBytes > len(msg) { + contentBytes = len(msg) + } + + for contentBytes > 0 { + headBytes := contentBytes / 2 + tailBytes := contentBytes - headBytes + + // Trim head backwards to a valid UTF-8 rune boundary. + head := msg[:headBytes] + for len(head) > 0 && !utf8.ValidString(head) { + head = head[:len(head)-1] + } + + // Trim tail forwards past any UTF-8 continuation bytes (0x80–0xBF) + // that would be invalid at the start of a string. + tail := msg[len(msg)-tailBytes:] + for len(tail) > 0 && tail[0]&0xC0 == 0x80 { + tail = tail[1:] + } + + candidate := excResp + candidate.Exception.Message = head + truncationMarker + tail + + candidateJSON, err := json.Marshal(candidate) + if err != nil { + return exceptionResponse{}, false + } + if len(candidateJSON) <= maxLength { + return candidate, true + } + + // The JSON exceeded the limit due to character escaping. Reduce + // contentBytes proportionally based on the observed expansion so the + // next iteration converges toward the target size. + newContentBytes := contentBytes * maxLength / len(candidateJSON) + if newContentBytes >= contentBytes { + // Integer division can round the proportional estimate back to the + // current value when candidateJSON is only marginally over the + // limit; step down by one to guarantee forward progress. + contentBytes-- + } else { + contentBytes = newContentBytes + } + } + + return exceptionResponse{}, false +} + func marshalResponse(response Response, maxResponseLength int, logger boshlog.Logger) ([]byte, error) { respJSON, err := json.Marshal(response) if err != nil { @@ -65,18 +163,30 @@ func marshalResponse(response Response, maxResponseLength int, logger boshlog.Lo } if len(respJSON) > maxResponseLength { - respJSON, err = json.Marshal(response.Shorten()) + shortened := response.Shorten() + + respJSON, err = json.Marshal(shortened) if err != nil { logger.Error(mbusHandlerLogTag, "Failed to marshal response: %s", err.Error()) return respJSON, bosherr.WrapError(err, "Marshalling JSON response") } - } - if len(respJSON) > maxResponseLength { - respJSON, err = BuildErrorWithJSON(responseMaxLengthErrMsg, logger) - if err != nil { - logger.Error(mbusHandlerLogTag, "Failed to build 'max length exceeded' response: %s", err.Error()) - return respJSON, bosherr.WrapError(err, "Building error") + if len(respJSON) > maxResponseLength { + if excResp, ok := shortened.(exceptionResponse); ok { + if truncated, ok := truncateExceptionResponse(excResp, maxResponseLength); ok { + respJSON, err = json.Marshal(truncated) + if err == nil && len(respJSON) <= maxResponseLength { + logger.Warn(mbusHandlerLogTag, "Response too large, exception message truncated to fit within %d bytes", maxResponseLength) + return respJSON, nil + } + } + } + + respJSON, err = BuildErrorWithJSON(responseMaxLengthErrMsg, logger) + if err != nil { + logger.Error(mbusHandlerLogTag, "Failed to build 'max length exceeded' response: %s", err.Error()) + return respJSON, bosherr.WrapError(err, "Building error") + } } } diff --git a/handler/perform_handler_with_json_test.go b/handler/perform_handler_with_json_test.go new file mode 100644 index 000000000..e9f208fbd --- /dev/null +++ b/handler/perform_handler_with_json_test.go @@ -0,0 +1,183 @@ +package handler_test + +import ( + "encoding/json" + "errors" + "strings" + "unicode/utf8" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + boshlog "github.com/cloudfoundry/bosh-utils/logger" + + . "github.com/cloudfoundry/bosh-agent/v2/handler" +) + +var _ = Describe("PerformHandlerWithJSON", func() { + var logger boshlog.Logger + + BeforeEach(func() { + logger = boshlog.NewLogger(boshlog.LevelNone) + }) + + Context("when the response is within the size limit", func() { + It("returns the full response unchanged", func() { + respBytes, _, err := PerformHandlerWithJSON( + []byte(`{"method":"ping","arguments":[],"reply_to":"fake-reply"}`), + func(req Request) Response { + return NewValueResponse("pong") + }, + 1024, + logger, + ) + Expect(err).NotTo(HaveOccurred()) + Expect(string(respBytes)).To(Equal(`{"value":"pong"}`)) + }) + }) + + Context("when the response exceeds the size limit", func() { + buildLargeExceptionHandler := func(msg string) Func { + return func(req Request) Response { + return NewExceptionResponse(errors.New(msg)) + } + } + + It("truncates the exception message preserving head and tail around the marker", func() { + const maxLen = 150 + // Construct a message with a distinct head and tail so we can verify both survive. + head := strings.Repeat("H", 500) + tail := strings.Repeat("T", 500) + msg := head + strings.Repeat("M", 500) + tail + + respBytes, _, err := PerformHandlerWithJSON( + []byte(`{"method":"big","arguments":[],"reply_to":"fake-reply"}`), + buildLargeExceptionHandler(msg), + maxLen, + logger, + ) + Expect(err).NotTo(HaveOccurred()) + Expect(len(respBytes)).To(BeNumerically("<=", maxLen)) + + var parsed map[string]map[string]string + Expect(json.Unmarshal(respBytes, &parsed)).To(Succeed()) + result := parsed["exception"]["message"] + + Expect(result).To(ContainSubstring("\n...[truncated]...\n"), "truncation marker should be on its own line") + Expect(result).To(HavePrefix("H"), "should preserve the beginning of the message") + Expect(result).To(HaveSuffix("T"), "should preserve the end of the message") + Expect(utf8.ValidString(result)).To(BeTrue(), "truncated message should be valid UTF-8") + }) + + It("produces JSON that is exactly at or under the limit", func() { + const maxLen = 200 + + respBytes, _, err := PerformHandlerWithJSON( + []byte(`{"method":"big","arguments":[],"reply_to":"fake-reply"}`), + buildLargeExceptionHandler(strings.Repeat("x", maxLen*5)), + maxLen, + logger, + ) + Expect(err).NotTo(HaveOccurred()) + Expect(len(respBytes)).To(BeNumerically("<=", maxLen)) + }) + + It("falls back to the generic error when response cannot be truncated to fit", func() { + // maxLen is so small that even the truncation marker alone won't fit. + const maxLen = 10 + + respBytes, _, err := PerformHandlerWithJSON( + []byte(`{"method":"big","arguments":[],"reply_to":"fake-reply"}`), + buildLargeExceptionHandler(strings.Repeat("x", 1000)), + maxLen, + logger, + ) + Expect(err).NotTo(HaveOccurred()) + Expect(string(respBytes)).To(ContainSubstring("Response exceeded maximum allowed length")) + }) + + It("preserves valid UTF-8 when truncating a multi-byte message", func() { + const maxLen = 100 + // Build a message entirely of 3-byte UTF-8 runes ('€' = 0xE2 0x82 0xAC). + multiByteMsg := strings.Repeat("€", maxLen) + + respBytes, _, err := PerformHandlerWithJSON( + []byte(`{"method":"big","arguments":[],"reply_to":"fake-reply"}`), + buildLargeExceptionHandler(multiByteMsg), + maxLen, + logger, + ) + Expect(err).NotTo(HaveOccurred()) + Expect(len(respBytes)).To(BeNumerically("<=", maxLen)) + + var parsed map[string]map[string]string + Expect(json.Unmarshal(respBytes, &parsed)).To(Succeed()) + result := parsed["exception"]["message"] + Expect(utf8.ValidString(result)).To(BeTrue(), "result should be valid UTF-8") + }) + + It("respects the limit when the message contains JSON-escapable characters", func() { + // `"` marshals to `\"` (2 bytes), so a naive byte-count would produce + // JSON roughly twice as large as expected. The handler must either + // successfully truncate within the limit or fall back to the generic + // error — it must never return oversized or invalid JSON. + const maxLen = 200 + escapableMsg := strings.Repeat(`"`, maxLen*10) + + respBytes, _, err := PerformHandlerWithJSON( + []byte(`{"method":"big","arguments":[],"reply_to":"fake-reply"}`), + buildLargeExceptionHandler(escapableMsg), + maxLen, + logger, + ) + Expect(err).NotTo(HaveOccurred()) + Expect(len(respBytes)).To(BeNumerically("<=", maxLen), "response must not exceed the size limit") + + var parsed map[string]interface{} + Expect(json.Unmarshal(respBytes, &parsed)).To(Succeed(), "response must be valid JSON") + + // Accept either a truncated exception message or the generic fallback. + if exc, ok := parsed["exception"].(map[string]interface{}); ok { + msg, _ := exc["message"].(string) + Expect(utf8.ValidString(msg)).To(BeTrue(), "exception message must be valid UTF-8") + } + }) + + It("truncates rather than falling back when the raw message fits but its escaped form does not", func() { + // The message's raw byte length fits within the limit, but every byte + // is an escapable quote, so its marshaled (escaped) size is ~2x larger + // and exceeds the limit. The handler must still produce a truncated + // exception message (with the marker), not the generic fallback. + const maxLen = 120 + escapableMsg := strings.Repeat(`"`, maxLen-40) + + respBytes, _, err := PerformHandlerWithJSON( + []byte(`{"method":"big","arguments":[],"reply_to":"fake-reply"}`), + buildLargeExceptionHandler(escapableMsg), + maxLen, + logger, + ) + Expect(err).NotTo(HaveOccurred()) + Expect(len(respBytes)).To(BeNumerically("<=", maxLen), "response must not exceed the size limit") + Expect(string(respBytes)).NotTo(ContainSubstring("Response exceeded maximum allowed length"), "should truncate, not fall back to the generic error") + + var parsed map[string]map[string]string + Expect(json.Unmarshal(respBytes, &parsed)).To(Succeed(), "response must be valid JSON") + Expect(parsed["exception"]["message"]).To(ContainSubstring("...[truncated]..."), "should contain truncation marker") + }) + + It("does not truncate when UnlimitedResponseLength is set", func() { + const bigMsgSize = 2 * 1024 * 1024 // 2 MB + + respBytes, _, err := PerformHandlerWithJSON( + []byte(`{"method":"big","arguments":[],"reply_to":"fake-reply"}`), + buildLargeExceptionHandler(strings.Repeat("x", bigMsgSize)), + UnlimitedResponseLength, + logger, + ) + Expect(err).NotTo(HaveOccurred()) + Expect(len(respBytes)).To(BeNumerically(">", bigMsgSize)) + Expect(string(respBytes)).NotTo(ContainSubstring("...[truncated]...")) + }) + }) +})