1 files changed, 58 insertions, 0 deletions
diff --git a/forged/internal/bare/limit.go b/forged/internal/bare/limit.go
new file mode 100644
index 0000000..212bc05
--- /dev/null
+++ b/forged/internal/bare/limit.go
@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright (c) 2025 Drew Devault <https://drewdevault.com>
+
+package bare
+
+import (
+	"errors"
+	"io"
+)
+
+var (
+	maxUnmarshalBytes uint64 = 1024 * 1024 * 32 /* 32 MiB */
+	maxArrayLength    uint64 = 1024 * 4         /* 4096 elements */
+	maxMapSize        uint64 = 1024
+)
+
+// MaxUnmarshalBytes sets the maximum size of a message decoded by unmarshal.
+// By default, this is set to 32 MiB.
+func MaxUnmarshalBytes(bytes uint64) {
+	maxUnmarshalBytes = bytes
+}
+
+// MaxArrayLength sets maximum number of elements in array. Defaults to 4096 elements
+func MaxArrayLength(length uint64) {
+	maxArrayLength = length
+}
+
+// MaxMapSize sets maximum size of map. Defaults to 1024 key/value pairs
+func MaxMapSize(size uint64) {
+	maxMapSize = size
+}
+
+// Use MaxUnmarshalBytes to prevent this error from occuring on messages which
+// are large by design.
+var ErrLimitExceeded = errors.New("Maximum message size exceeded")
+
+// Identical to io.LimitedReader, except it returns our custom error instead of
+// EOF if the limit is reached.
+type limitedReader struct {
+	R io.Reader
+	N uint64
+}
+
+func (l *limitedReader) Read(p []byte) (n int, err error) {
+	if l.N <= 0 {
+		return 0, ErrLimitExceeded
+	}
+	if uint64(len(p)) > l.N {
+		p = p[0:l.N]
+	}
+	n, err = l.R.Read(p)
+	l.N -= uint64(n)
+	return
+}
+
+func newLimitedReader(r io.Reader) *limitedReader {
+	return &limitedReader{r, maxUnmarshalBytes}
+}