Truncate response errors vs swallowing if too large

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")
+			}
 		}
 	}
 

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]..."))
+		})
+	})
+})