From 78c76e4d770a7d39ed856e51b4c9c3e5a8fd5e06 Mon Sep 17 00:00:00 2001
From: Saeed Rasooli <saeed.gnu@gmail.com>
Date: Fri, 14 Jul 2023 12:14:11 +0430
Subject: [PATCH] add REPASSGEN_MAX_LENGTH to prevent out-of-memory during fuzz
 testing

---
 lib/fuzz_test.go |  7 ++++++-
 lib/repeat.go    |  3 +++
 lib/state.go     | 34 +++++++++++++++++++++++++---------
 3 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/lib/fuzz_test.go b/lib/fuzz_test.go
index e8ba578..410f6b7 100644
--- a/lib/fuzz_test.go
+++ b/lib/fuzz_test.go
@@ -1,6 +1,7 @@
 package passgen_test
 
 import (
+	"os"
 	"testing"
 	"time"
 
@@ -18,8 +19,12 @@ func FuzzGenerate(f *testing.F) {
 	for _, tc := range testcases {
 		f.Add(tc) // Use f.Add to provide a seed corpus
 	}
+	// define a max output length to prevent out-of-memory and crash
+	os.Setenv("REPASSGEN_MAX_LENGTH", "500")
 	f.Fuzz(func(t *testing.T, pattern string) {
-		// TODO: define a max output length to prevent out of memory and crash
+		if len(pattern) > 100 {
+			return
+		}
 		out, _, err := passgen.Generate(passgen.GenerateInput{
 			Pattern: []rune(pattern),
 		})
diff --git a/lib/repeat.go b/lib/repeat.go
index 0908e11..f6ef194 100644
--- a/lib/repeat.go
+++ b/lib/repeat.go
@@ -10,6 +10,9 @@ func (g *repeatGenerator) Generate(s *State) error {
 	count := g.count
 	for i := int64(0); i < count; i++ {
 		err := child.Generate(s)
+		if s.tooLong() {
+			break
+		}
 		if err != nil {
 			return err
 		}
diff --git a/lib/state.go b/lib/state.go
index 4fb3e07..d6f1647 100644
--- a/lib/state.go
+++ b/lib/state.go
@@ -3,6 +3,8 @@ package passgen
 import (
 	"fmt"
 	"log"
+	"os"
+	"strconv"
 )
 
 // SharedState is the shared part of State
@@ -13,6 +15,8 @@ type SharedState struct {
 	errorMarkLen   int
 	patternEntropy float64
 	lastGroupId    uint64
+
+	maxOutputLength int
 }
 
 // State is lex inputs, output and temp state
@@ -46,15 +50,11 @@ func (s *State) moveBack(chars uint64) {
 	s.absPos -= chars
 }
 
-// func (s *State) moveBackAbs(chars uint64) {
-// 	if s.absPos < chars {
-// 		log.Printf("moveBack(%v) with absPos=%v", chars, s.absPos)
-// 		return
-// 	}
-// 	s.absPos -= chars
-// }
-
 func (s *State) addOutputOne(c rune) {
+	if s.tooLong() {
+		s.lastGen = nil
+		return
+	}
 	s.lastGen = &staticStringGenerator{str: []rune{c}}
 	s.lastGen.Generate(s)
 }
@@ -69,7 +69,14 @@ func (s *State) addOutputNonRepeatable(data []rune) {
 	s.output = append(s.output, data...)
 }
 
+func (s *State) tooLong() bool {
+	return s.maxOutputLength > 0 && len(s.output) > s.maxOutputLength
+}
+
 func (s *State) end() bool {
+	if s.tooLong() {
+		return true
+	}
 	return s.inputPos >= uint64(len(s.input))
 }
 
@@ -124,10 +131,19 @@ func (s *State) errorUnknown(msg string, args ...interface{}) error {
 
 // NewSharedState is factory function for SharedState
 func NewSharedState() *SharedState {
-	return &SharedState{
+	ss := &SharedState{
 		groupsOutput: map[uint64][]rune{},
 		errorMarkLen: 1,
 	}
+	maxLengthStr := os.Getenv("REPASSGEN_MAX_LENGTH")
+	if maxLengthStr != "" {
+		maxLength, err := strconv.ParseInt(maxLengthStr, 10, 64)
+		if err != nil {
+			panic("invalid REPASSGEN_MAX_LENGTH: must be an integer")
+		}
+		ss.maxOutputLength = int(maxLength)
+	}
+	return ss
 }
 
 // NewState is factory function for State