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"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/panjf2000/ants/v2"
|
||||
"github.com/salvacybersec/keyhunter/pkg/engine"
|
||||
"github.com/salvacybersec/keyhunter/pkg/providers"
|
||||
"github.com/tidwall/gjson"
|
||||
@@ -18,6 +20,9 @@ import (
|
||||
// DefaultTimeout is the per-call verification timeout when none is configured.
|
||||
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.
|
||||
const maxMetadataBody = 1 << 20 // 1 MiB
|
||||
|
||||
@@ -137,6 +142,75 @@ func (v *HTTPVerifier) Verify(ctx context.Context, f engine.Finding, p providers
|
||||
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.
|
||||
func substituteKey(s, key string) string {
|
||||
s = strings.ReplaceAll(s, "{{KEY}}", key)
|
||||
|
||||
Reference in New Issue
Block a user