Merge pull request #5 from monrax/noworkers

magic.go

@@ -12,15 +12,50 @@ type job struct {
 	resultChan chan *FileType
 }
 
+// LookupConfig is a struct that contains configuration details to modify the default Lookup behavior
+type LookupConfig struct {
+	ConcurrencyEnabled bool // the search will be performed concurrently by multiple worker goroutines when this field is set to true. The search will be carried out by the calling goroutine if set to false.
+	WorkerCount        int  // number of worker goroutines to be spawned if concurrency is set to true. If set to -1, workerCount will be set to use all the available cores.
+}
+
 // ErrUnknown infers the file type cannot be determined by the provided magic bytes
 var ErrUnknown = fmt.Errorf("unknown file type")
 
 // Lookup looks up the file type based on the provided magic bytes. You should provide at least the first 1024 bytes of the file in this slice.
 // A magic.ErrUnknown will be returned if the file type is not known.
 func Lookup(bytes []byte) (*FileType, error) {
+	return lookup(bytes, true, -1)
+}
+
+// LookupWithConfig looks up the file type based on the provided magic bytes, and a given configuration. You should provide at least the first 1024 bytes of the file in this slice.
+// A magic.ErrUnknown will be returned if the file type is not known.
+func LookupWithConfig(bytes []byte, config LookupConfig) (*FileType, error) {
+	return lookup(bytes, config.ConcurrencyEnabled, config.WorkerCount)
+}
+
+// LookupSync lookups up the file type based on the provided magic bytes without spawning any additional goroutines. You should provide at least the first 1024 bytes of the file in this slice.
+// A magic.ErrUnknown will be returned if the file type is not known.
+func LookupSync(bytes []byte) (*FileType, error) {
+	return lookup(bytes, false, 0)
+}
+
+func lookup(bytes []byte, concurrent bool, workers int) (*FileType, error) {
+	// additional worker count check: avoid deadlock when worker count is set to zero
+	if !concurrent || workers == 0 {
+		for _, t := range types {
+			ft := t.check(bytes, 0)
+			if ft != nil {
+				return ft, nil
+			}
+		}
+		return nil, ErrUnknown
+	}
 
 	// use all available cores
 	workerCount := runtime.GOMAXPROCS(0)
+	if workers > -1 && workers < workerCount {
+		workerCount = workers
+	}
 	workChan := make(chan job)
 	resultChan := make(chan *FileType)
 

magic_test.go

@@ -24,3 +24,47 @@ func TestNestedLookup(t *testing.T) {
 	assert.Equal(t, fileType.Extension, "wav")
 
 }
+
+func TestLookupWithConfig(t *testing.T) {
+	payload := []byte{31, 139, 8, 0, 130, 139, 110, 100, 2, 255, 61, 143, 65, 14, 131, 32, 20, 68, 247}
+
+	workerCounts := []int{-1, 0, 1, 2, 10000}
+
+	lookupConfigs := func(counts []int) []LookupConfig {
+		configs := make([]LookupConfig, len(workerCounts)*2)
+		for i := 0; i < len(counts)*2; i++ {
+			configs[i] = LookupConfig{
+				WorkerCount: counts[i%len(counts)],
+			}
+			if i < len(counts) {
+				configs[i].ConcurrencyEnabled = true
+			}
+		}
+		return configs
+	}(workerCounts)
+
+	for _, config := range lookupConfigs {
+		fileType, err := LookupWithConfig(payload, config)
+		require.Nil(t, err)
+		assert.Equal(t, fileType.Extension, "gz")
+	}
+
+}
+
+func TestLookupSync(t *testing.T) {
+
+	fileType, err := LookupSync([]byte{0xa1, 0xb2, 0xc3, 0xd4, 0x00, 0x00, 0x00, 0x00})
+	require.Nil(t, err)
+
+	assert.Equal(t, fileType.Extension, "pcap")
+
+}
+
+func TestNestedLookupSync(t *testing.T) {
+
+	fileType, err := LookupSync([]byte{0x52, 0x49, 0x46, 0x46, 0, 0, 0, 0, 0x57, 0x41, 0x56, 0x45})
+	require.Nil(t, err)
+
+	assert.Equal(t, fileType.Extension, "wav")
+
+}