feat(05-03): add VerifyAll ants worker pool for parallel verification
- VerifyAll(ctx, findings, reg, workers) returns a result channel closed after all findings are processed or ctx is cancelled. - Default worker count of 10 when workers <= 0. - Missing providers yield StatusUnknown with 'provider not found' error. - Graceful context cancellation stops dispatch while still draining inflight.
This commit is contained in:
@@ -8,8 +8,10 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/panjf2000/ants/v2"
|
||||||
"github.com/salvacybersec/keyhunter/pkg/engine"
|
"github.com/salvacybersec/keyhunter/pkg/engine"
|
||||||
"github.com/salvacybersec/keyhunter/pkg/providers"
|
"github.com/salvacybersec/keyhunter/pkg/providers"
|
||||||
"github.com/tidwall/gjson"
|
"github.com/tidwall/gjson"
|
||||||
@@ -18,6 +20,9 @@ import (
|
|||||||
// DefaultTimeout is the per-call verification timeout when none is configured.
|
// DefaultTimeout is the per-call verification timeout when none is configured.
|
||||||
const DefaultTimeout = 10 * time.Second
|
const DefaultTimeout = 10 * time.Second
|
||||||
|
|
||||||
|
// DefaultWorkers is the fallback worker-pool size for VerifyAll.
|
||||||
|
const DefaultWorkers = 10
|
||||||
|
|
||||||
// maxMetadataBody caps how much of a JSON response we read for metadata extraction.
|
// maxMetadataBody caps how much of a JSON response we read for metadata extraction.
|
||||||
const maxMetadataBody = 1 << 20 // 1 MiB
|
const maxMetadataBody = 1 << 20 // 1 MiB
|
||||||
|
|
||||||
@@ -137,6 +142,75 @@ func (v *HTTPVerifier) Verify(ctx context.Context, f engine.Finding, p providers
|
|||||||
return res
|
return res
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// VerifyAll runs verification for every finding through an ants worker pool
|
||||||
|
// of the given size (or DefaultWorkers when workers <= 0). The returned
|
||||||
|
// channel is closed once every finding has been processed or ctx is cancelled.
|
||||||
|
//
|
||||||
|
// Findings whose provider is not present in reg are emitted as
|
||||||
|
// Result{Status: StatusUnknown, Error: "provider not found in registry"}
|
||||||
|
// rather than silently dropped.
|
||||||
|
func (v *HTTPVerifier) VerifyAll(ctx context.Context, findings []engine.Finding, reg *providers.Registry, workers int) <-chan Result {
|
||||||
|
if workers <= 0 {
|
||||||
|
workers = DefaultWorkers
|
||||||
|
}
|
||||||
|
out := make(chan Result, len(findings))
|
||||||
|
|
||||||
|
pool, err := ants.NewPool(workers)
|
||||||
|
if err != nil {
|
||||||
|
go func() {
|
||||||
|
defer close(out)
|
||||||
|
for _, f := range findings {
|
||||||
|
out <- Result{
|
||||||
|
ProviderName: f.ProviderName,
|
||||||
|
KeyMasked: f.KeyMasked,
|
||||||
|
Status: StatusError,
|
||||||
|
Error: "pool init: " + err.Error(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
go func() {
|
||||||
|
defer close(out)
|
||||||
|
defer pool.Release()
|
||||||
|
|
||||||
|
for i := range findings {
|
||||||
|
if ctx.Err() != nil {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
f := findings[i]
|
||||||
|
wg.Add(1)
|
||||||
|
submitErr := pool.Submit(func() {
|
||||||
|
defer wg.Done()
|
||||||
|
prov, ok := reg.Get(f.ProviderName)
|
||||||
|
if !ok {
|
||||||
|
out <- Result{
|
||||||
|
ProviderName: f.ProviderName,
|
||||||
|
KeyMasked: f.KeyMasked,
|
||||||
|
Status: StatusUnknown,
|
||||||
|
Error: "provider not found in registry",
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
out <- v.Verify(ctx, f, prov)
|
||||||
|
})
|
||||||
|
if submitErr != nil {
|
||||||
|
wg.Done()
|
||||||
|
out <- Result{
|
||||||
|
ProviderName: f.ProviderName,
|
||||||
|
KeyMasked: f.KeyMasked,
|
||||||
|
Status: StatusError,
|
||||||
|
Error: submitErr.Error(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
}()
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
// substituteKey replaces both {{KEY}} and the legacy {KEY} placeholder.
|
// substituteKey replaces both {{KEY}} and the legacy {KEY} placeholder.
|
||||||
func substituteKey(s, key string) string {
|
func substituteKey(s, key string) string {
|
||||||
s = strings.ReplaceAll(s, "{{KEY}}", key)
|
s = strings.ReplaceAll(s, "{{KEY}}", key)
|
||||||
|
|||||||
Reference in New Issue
Block a user