added panic recovery

scott-mescudi · Dec 11, 2024 · 6164ceb · 6164ceb
1 parent e9cfd39
commit 6164ceb
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 10 deletions.
diff --git a/image_core.go b/image_core.go
@@ -49,11 +49,22 @@ func (m *ExtractHandler) ExtractDataFromImage(coverImage image.Image, bitDepth u
 		return nil, err
 	}
 
-	var moddedData = make([]byte, 0)
-	for i := 4; i < lenData+4; i++ {
-		moddedData = append(moddedData, data[i])
-	}
-
+
+    var moddedData = make([]byte, 0, lenData) 
+    defer func() {
+        if r := recover(); r != nil {
+            moddedData = nil
+            err = fmt.Errorf("fatal error: %v", r)
+        }
+    }()
+
+    for i := 4; i < lenData+4; i++ {
+        if i >= len(data) {
+            return nil, fmt.Errorf("index out of range while accessing data: %d", i)
+        }
+        moddedData = append(moddedData, data[i])
+    }
+
 	return moddedData, nil
 }
 

diff --git a/image_embedder.go b/image_embedder.go
@@ -134,11 +134,21 @@ func (m *ExtractHandler) Decode(coverImage image.Image, bitDepth uint8, isDefaul
 		return nil, errors.New("extracted data length is zero")
 	}
 
-	// Retrieve the actual embedded data
-	var moddedData = make([]byte, 0, lenData)
-	for i := 4; i < lenData+4; i++ {
-		moddedData = append(moddedData, data[i])
-	}
+
+    var moddedData = make([]byte, 0, lenData) 
+    defer func() {
+        if r := recover(); r != nil {
+            moddedData = nil
+            err = fmt.Errorf("fatal error: %v", r)
+        }
+    }()
+
+    for i := 4; i < lenData+4; i++ {
+        if i >= len(data) {
+            return nil, fmt.Errorf("index out of range while accessing data: %d", i)
+        }
+        moddedData = append(moddedData, data[i])
+    }
 
 	// Decompress data if required
 	if isDefaultCompressed {