Repository: gabriel-vasile/mimetype
Pull request #158 (2021-06-06 04:42:28)

allow hierarchical extensions
This change addresses a usability problem with the current API, where defining a new mime type unconditionally makes it a child of the top node.
For example, registering a new sub-type of application/zip is a bit clunky, and essentially requires redefining the detection of the zip format at the extension point.

With this change, a slightly more generic Extend() entry point is defined at the individual MIME level, so that the previous scenario can now look like:

    mimetype.Lookup("application/zip").Extend(pred, "foo/zip", ".foo")

With this, the following 2 lines are equivalent:

    mimetype.Extend(pred, "foo", ".foo")
    mimetype.Lookup("application/octet-stream").Extend(pred, "foo", ".foo")



Patch:
diff --git a/mime.go b/mime.go
index 0def642..2e7f656 100644
--- a/mime.go
+++ b/mime.go
@@ -2,6 +2,7 @@ package mimetype
 
 import (
 	"mime"
+	"sync"
 
 	"github.com/gabriel-vasile/mimetype/internal/charset"
 	"github.com/gabriel-vasile/mimetype/internal/magic"
@@ -18,6 +19,8 @@ type MIME struct {
 	detector magic.Detector
 	children []*MIME
 	parent   *MIME
+
+	mu sync.RWMutex
 }
 
 // String returns the string representation of the MIME type, e.g., "application/zip".
@@ -56,6 +59,9 @@ func (m *MIME) Is(expectedMIME string) bool {
 	if expectedMIME == found {
 		return true
 	}
+
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	for _, alias := range m.aliases {
 		if alias == expectedMIME {
 			return true
@@ -85,13 +91,17 @@ func newMIME(
 }
 
 func (m *MIME) alias(aliases ...string) *MIME {
+	m.mu.Lock()
 	m.aliases = aliases
+	m.mu.Unlock()
 	return m
 }
 
 // match does a depth-first search on the signature tree. It returns the deepest
 // successful node for which all the children detection functions fail.
 func (m *MIME) match(in []byte, readLimit uint32) *MIME {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	for _, c := range m.children {
 		if c.detector(in, readLimit) {
 			return c.match(in, readLimit)
@@ -117,6 +127,8 @@ func (m *MIME) match(in []byte, readLimit uint32) *MIME {
 // flatten transforms an hierarchy of MIMEs into a slice of MIMEs.
 func (m *MIME) flatten() []*MIME {
 	out := []*MIME{m}
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	for _, c := range m.children {
 		out = append(out, c.flatten()...)
 	}
@@ -131,6 +143,8 @@ func (m *MIME) clone(ps map[string]string) *MIME {
 		clonedMIME = mime.FormatMediaType(m.mime, ps)
 	}
 
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	return &MIME{
 		mime:      clonedMIME,
 		aliases:   m.aliases,
@@ -151,3 +165,38 @@ func (m *MIME) cloneHierarchy(ps map[string]string) *MIME {
 
 	return ret
 }
+
+func (m *MIME) lookup(mime string) *MIME {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	for _, n := range append(m.aliases, m.mime) {
+		if n == mime {
+			return m
+		}
+	}
+
+	for _, c := range m.children {
+		if m := c.lookup(mime); m != nil {
+			return m
+		}
+	}
+	return nil
+}
+
+// Extend adds detection for a sub-format. The detector is a function
+// returning true when the raw input file satisfies a signature.
+// The sub-format will be detected if all the detectors in the parent chain return true.
+// The extension should include the leading dot, as in ".html".
+func (m *MIME) Extend(detector func(raw []byte, limit uint32) bool, mime, extension string, aliases ...string) {
+	c := &MIME{
+		mime:      mime,
+		extension: extension,
+		detector:  detector,
+		parent:    m,
+		aliases:   aliases,
+	}
+
+	m.mu.Lock()
+	m.children = append([]*MIME{c}, m.children...)
+	m.mu.Unlock()
+}
diff --git a/mimetype.go b/mimetype.go
index 154f69b..5c0b27a 100644
--- a/mimetype.go
+++ b/mimetype.go
@@ -6,11 +6,15 @@ import (
 	"io/ioutil"
 	"mime"
 	"os"
+	"sync"
 )
 
 // readLimit is the maximum number of bytes from the input used when detecting.
 var readLimit uint32 = 3072
 
+// rootMu guards the readLimit used when creating the detection buffer.
+var rootMu sync.RWMutex
+
 // Detect returns the MIME type found from the provided byte slice.
 //
 // The result is always a valid MIME type, with application/octet-stream
@@ -100,19 +104,14 @@ func SetLimit(limit uint32) {
 	rootMu.Unlock()
 }
 
-// Extend adds detection for other file formats. The detector is a function
-// returning true when the raw input file satisfies a  signature.
-// The extension should include the leading dot, as in ".html".
+// Extend adds detection for other file formats.
+// It is equivalent to calling Extend() on the root mime type "application/octet-stream".
 func Extend(detector func(raw []byte, limit uint32) bool, mime, extension string, aliases ...string) {
-	m := &MIME{
-		mime:      mime,
-		extension: extension,
-		detector:  detector,
-		parent:    root,
-		aliases:   aliases,
-	}
+	root.Extend(detector, mime, extension, aliases...)
+}
 
-	rootMu.Lock()
-	root.children = append([]*MIME{m}, root.children...)
-	rootMu.Unlock()
+// Lookup finds a MIME object by its string representation.
+// The representation can be the main mime type, or any of its aliases.
+func Lookup(mime string) *MIME {
+	return root.lookup(mime)
 }
diff --git a/tree.go b/tree.go
index a291696..757ffed 100644
--- a/tree.go
+++ b/tree.go
@@ -1,8 +1,6 @@
 package mimetype
 
 import (
-	"sync"
-
 	"github.com/gabriel-vasile/mimetype/internal/magic"
 )
 
@@ -29,9 +27,6 @@ var root = newMIME("application/octet-stream", "",
 	text,
 )
 
-// rootMu guards the root tree and the readLimit used when creating the detection buffer.
-var rootMu sync.RWMutex
-
 // The list of nodes appended to the root node.
 var (
 	xz   = newMIME("application/x-xz", ".xz", magic.Xz)


Test patch:
diff --git a/mimetype_test.go b/mimetype_test.go
index ef5fc79..d123eab 100644
--- a/mimetype_test.go
+++ b/mimetype_test.go
@@ -353,14 +353,15 @@ func TestExtend(t *testing.T) {
 		return bytes.HasPrefix(raw, []byte("foobar"))
 	}
 
-	mimetype.Extend(foobarDet, "text/foobar", ".fb")
+	txt := "text/plain"
+	mimetype.Lookup(txt).Extend(foobarDet, "text/foobar", ".fb")
 
 	mtype := mimetype.Detect([]byte("foobar file content"))
 	if !mtype.Is("text/foobar") {
 		t.Fatalf("extend error; expected text/foobar, got: %s", mtype)
 	}
-	if !mtype.Parent().Is("application/octet-stream") {
-		t.Fatalf("extend parent error; expected application/octet-stream, got: %s", mtype.Parent())
+	if !mtype.Parent().Is(txt) {
+		t.Fatalf("extend parent error; expected %s, got: %s", txt, mtype.Parent())
 	}
 }
 
diff --git a/tree_test.go b/tree_test.go
index 8d2415a..14350db 100644
--- a/tree_test.go
+++ b/tree_test.go
@@ -78,3 +78,55 @@ func TestMIMEFormat(t *testing.T) {
 		}
 	}
 }
+
+func TestLookup(t *testing.T) {
+	data := []struct {
+		mime string
+		m    *MIME
+	}{
+		{root.mime, root},
+		{zip.mime, zip},
+		{zip.aliases[0], zip},
+		{xlsx.mime, xlsx},
+	}
+
+	for _, tt := range data {
+		t.Run(fmt.Sprintf("lookup %s", tt.mime), func(t *testing.T) {
+			if m := Lookup(tt.mime); m != tt.m {
+				t.Fatalf("failed to lookup: %s", tt.mime)
+			}
+		})
+	}
+}
+
+func TestExtend(t *testing.T) {
+	data := []struct {
+		mime   string
+		ext    string
+		parent *MIME
+	}{
+		{"foo", ".foo", nil},
+		{"bar", ".bar", root},
+		{"baz", ".baz", zip},
+	}
+
+	for _, tt := range data {
+		t.Run(fmt.Sprintf("extending to %s", tt.mime), func(t *testing.T) {
+			extend := Extend
+			if tt.parent != nil {
+				extend = tt.parent.Extend
+			} else {
+				tt.parent = root
+			}
+
+			extend(func(raw []byte, limit uint32) bool { return false }, tt.mime, tt.ext)
+			m := Lookup(tt.mime)
+			if m == nil {
+				t.Fatalf("mime %s not found", tt.mime)
+			}
+			if m.parent != tt.parent {
+				t.Fatalf("mime %s has wrong parent: want %s, got %s", tt.mime, tt.parent.mime, m.parent.mime)
+			}
+		})
+	}
+}