浏览代码

bufreader, RegisterFormat updated

Viktor Sokolov 1 月之前
父节点
当前提交
87f96eae61
共有 4 个文件被更改,包括 44 次插入90 次删除
  1. 13 38
      imagedetect/detect.go
  2. 17 42
      imagedetect/registry.go
  3. 11 7
      imagedetect/registry_test.go
  4. 3 3
      imagedetect/svg.go

+ 13 - 38
imagedetect/detect.go

@@ -3,48 +3,23 @@ package imagedetect
 import (
 	"io"
 
+	"github.com/imgproxy/imgproxy/v3/bufreader"
 	"github.com/imgproxy/imgproxy/v3/imagetype_new"
 )
 
+const (
+	// maxDetectionLimit is maximum bytes detectors allowed to read from the source
+	maxDetectionLimit = 32768
+)
+
 // Detect attempts to detect the image type from a reader.
 // It first tries magic byte detection, then custom detectors in registration order
 func Detect(r io.Reader) (imagetype_new.Type, error) {
-	// Start with 64 bytes to cover magic bytes
-	buf := make([]byte, 64)
-
-	n, err := io.ReadFull(r, buf)
-	if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
-		return imagetype_new.Unknown, err
-	}
-
-	data := buf[:n]
-
-	// First try magic byte detection
-	for _, magic := range registry.magicBytes {
-		if hasMagicBytes(data, magic) {
-			return magic.Type, nil
-		}
-	}
-
-	// Then try custom detectors
-	for _, detector := range registry.detectors {
-		// Check if we have enough bytes for this detector
-		if len(data) < detector.BytesNeeded {
-			// Need to read more data
-			additionalBytes := detector.BytesNeeded - len(data)
-			extraBuf := make([]byte, additionalBytes)
-			extraN, err := io.ReadFull(r, extraBuf)
-
-			// It's fine if we can't read required number of bytes
-			if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
-				return imagetype_new.Unknown, err
-			}
-
-			// Extend our data buffer
-			data = append(data, extraBuf[:extraN]...)
-		}
+	br := bufreader.New(io.LimitReader(r, maxDetectionLimit))
 
-		if typ, err := detector.Func(data); err == nil && typ != imagetype_new.Unknown {
+	for _, fn := range registry.detectors {
+		br.Rewind()
+		if typ, err := fn(br); err == nil && typ != imagetype_new.Unknown {
 			return typ, nil
 		}
 	}
@@ -54,12 +29,12 @@ func Detect(r io.Reader) (imagetype_new.Type, error) {
 
 // hasMagicBytes checks if the data matches a magic byte signature
 // Supports '?' characters in signature which match any byte
-func hasMagicBytes(data []byte, magic MagicBytes) bool {
-	if len(data) < len(magic.Signature) {
+func hasMagicBytes(data []byte, magic []byte) bool {
+	if len(data) < len(magic) {
 		return false
 	}
 
-	for i, c := range magic.Signature {
+	for i, c := range magic {
 		if c != data[i] && c != '?' {
 			return false
 		}

+ 17 - 42
imagedetect/registry.go

@@ -1,32 +1,16 @@
 package imagedetect
 
 import (
-	"sync"
-
+	"github.com/imgproxy/imgproxy/v3/bufreader"
 	"github.com/imgproxy/imgproxy/v3/imagetype_new"
 )
 
 // DetectFunc is a function that detects the image type from byte data
-type DetectFunc func(b []byte) (imagetype_new.Type, error)
-
-// MagicBytes represents a magic byte signature for image type detection
-// Signature can contain '?' characters which match any byte
-type MagicBytes struct {
-	Signature []byte
-	Type      imagetype_new.Type
-}
-
-// Detector represents a registered detector function with its byte requirements
-type Detector struct {
-	Func        DetectFunc
-	BytesNeeded int
-}
+type DetectFunc func(r bufreader.ReadPeeker) (imagetype_new.Type, error)
 
 // Registry manages the registration and execution of image type detectors
 type Registry struct {
-	mu         sync.RWMutex
-	detectors  []Detector
-	magicBytes []MagicBytes
+	detectors []DetectFunc
 }
 
 // Global registry instance
@@ -35,41 +19,32 @@ var registry = &Registry{}
 // RegisterDetector registers a custom detector function
 // Detectors are tried in the order they were registered
 func RegisterDetector(detector DetectFunc, bytesNeeded int) {
-	registry.mu.Lock()
-	defer registry.mu.Unlock()
-	registry.detectors = append(registry.detectors, Detector{
-		Func:        detector,
-		BytesNeeded: bytesNeeded,
-	})
+	registry.RegisterDetector(detector, bytesNeeded)
 }
 
 // RegisterMagicBytes registers magic bytes for a specific image type
 // Magic byte detectors are always tried before custom detectors
 func RegisterMagicBytes(signature []byte, typ imagetype_new.Type) {
-	registry.mu.Lock()
-	defer registry.mu.Unlock()
-	registry.magicBytes = append(registry.magicBytes, MagicBytes{
-		Signature: signature,
-		Type:      typ,
-	})
+	registry.RegisterMagicBytes(signature, typ)
 }
 
 // RegisterDetector registers a custom detector function on this registry instance
 func (r *Registry) RegisterDetector(detector DetectFunc, bytesNeeded int) {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-	r.detectors = append(r.detectors, Detector{
-		Func:        detector,
-		BytesNeeded: bytesNeeded,
-	})
+	r.detectors = append(r.detectors, detector)
 }
 
 // RegisterMagicBytes registers magic bytes for a specific image type on this registry instance
 func (r *Registry) RegisterMagicBytes(signature []byte, typ imagetype_new.Type) {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-	r.magicBytes = append(r.magicBytes, MagicBytes{
-		Signature: signature,
-		Type:      typ,
+	r.detectors = append(r.detectors, func(r bufreader.ReadPeeker) (imagetype_new.Type, error) {
+		b, err := r.Peek(len(signature))
+		if err != nil {
+			return imagetype_new.Unknown, err
+		}
+
+		if hasMagicBytes(b, signature) {
+			return typ, nil
+		}
+
+		return imagetype_new.Unknown, nil
 	})
 }

+ 11 - 7
imagedetect/registry_test.go

@@ -3,6 +3,7 @@ package imagedetect
 import (
 	"testing"
 
+	"github.com/imgproxy/imgproxy/v3/bufreader"
 	"github.com/imgproxy/imgproxy/v3/imagetype_new"
 	"github.com/stretchr/testify/require"
 )
@@ -12,8 +13,12 @@ func TestRegisterDetector(t *testing.T) {
 	testRegistry := &Registry{}
 
 	// Create a test detector function
-	testDetector := func(data []byte) (imagetype_new.Type, error) {
-		if len(data) >= 2 && data[0] == 0xFF && data[1] == 0xD8 {
+	testDetector := func(r bufreader.ReadPeeker) (imagetype_new.Type, error) {
+		b, err := r.Peek(2)
+		if err != nil {
+			return imagetype_new.Unknown, err
+		}
+		if len(b) >= 2 && b[0] == 0xFF && b[1] == 0xD8 {
 			return imagetype_new.JPEG, nil
 		}
 		return imagetype_new.Unknown, newUnknownFormatError()
@@ -24,20 +29,19 @@ func TestRegisterDetector(t *testing.T) {
 
 	// Verify the detector is registered
 	require.Len(t, testRegistry.detectors, 1)
-	require.Equal(t, 64, testRegistry.detectors[0].BytesNeeded)
-	require.NotNil(t, testRegistry.detectors[0].Func)
+	require.NotNil(t, testRegistry.detectors[0])
 }
 
 func TestRegisterMagicBytes(t *testing.T) {
 	// Create a test registry to avoid interfering with global state
 	testRegistry := &Registry{}
 
+	require.Empty(t, testRegistry.detectors)
+
 	// Register magic bytes for JPEG using the method
 	jpegMagic := []byte{0xFF, 0xD8}
 	testRegistry.RegisterMagicBytes(jpegMagic, imagetype_new.JPEG)
 
 	// Verify the magic bytes are registered
-	require.Len(t, testRegistry.magicBytes, 1)
-	require.Equal(t, jpegMagic, testRegistry.magicBytes[0].Signature)
-	require.Equal(t, imagetype_new.JPEG, testRegistry.magicBytes[0].Type)
+	require.Len(t, testRegistry.detectors, 1)
 }

+ 3 - 3
imagedetect/svg.go

@@ -1,9 +1,9 @@
 package imagedetect
 
 import (
-	"bytes"
 	"strings"
 
+	"github.com/imgproxy/imgproxy/v3/bufreader"
 	"github.com/imgproxy/imgproxy/v3/imagetype_new"
 
 	"github.com/tdewolff/parse/v2"
@@ -15,8 +15,8 @@ func init() {
 	RegisterDetector(IsSVG, 1000)
 }
 
-func IsSVG(b []byte) (imagetype_new.Type, error) {
-	l := xml.NewLexer(parse.NewInput(bytes.NewReader(b)))
+func IsSVG(r bufreader.ReadPeeker) (imagetype_new.Type, error) {
+	l := xml.NewLexer(parse.NewInput(r))
 
 	for {
 		tt, _ := l.Next()