Skip to content

Commit 0ae0ba8

Browse files
cx-rogerio-dalotcx-rogerio-dalot
authored
fix: thread exhaustion error (#297)
<!-- Thanks for contributing to 2ms by offering a pull request. --> **Proposed Changes** The binary as is can take the process to resource exhaustion due to going over the limit of 10k goroutines. This is a quickfix/hotfix and in no case sustainable for the future, we need indeed to refactor our current implementation but it should solve the issue for now. <!-- Please describe the big picture of your changes here. If it fixes a bug or resolves a feature request, be sure to link to that issue. --> **Checklist** - [ ] I covered my changes with tests. - [ ] I Updated the documentation that is affected by my changes: - [ ] Change in the CLI arguments - [ ] Change in the configuration file I submit this contribution under the Apache-2.0 license.
1 parent a2e47f1 commit 0ae0ba8

File tree

15 files changed

+196
-231
lines changed

15 files changed

+196
-231
lines changed

‎.2ms.yml‎

Lines changed: 111 additions & 111 deletions
Large diffs are not rendered by default.

‎cmd/workers.go‎

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2,18 +2,17 @@ package cmd
22

33
import (
44
"context"
5-
"sync"
65

76
"github.com/checkmarx/2ms/v3/engine"
87
"github.com/checkmarx/2ms/v3/engine/extra"
9-
"github.com/checkmarx/2ms/v3/lib/secrets"
108
"golang.org/x/sync/errgroup"
119
)
1210

1311
func ProcessItems(engineInstance engine.IEngine, pluginName string) {
1412
defer Channels.WaitGroup.Done()
1513

1614
g, ctx := errgroup.WithContext(context.Background())
15+
g.SetLimit(1000)
1716
for item := range Channels.Items {
1817
Report.TotalItemsScanned++
1918
item := item
@@ -71,37 +70,43 @@ func ProcessSecretsWithValidation() {
7170
func ProcessSecretsExtras() {
7271
defer Channels.WaitGroup.Done()
7372

74-
wgExtras := &sync.WaitGroup{}
73+
g := errgroup.Group{}
74+
g.SetLimit(10)
7575
for secret := range SecretsExtrasChan {
76-
wgExtras.Add(1)
77-
go extra.AddExtraToSecret(secret, wgExtras)
76+
g.Go(func() error {
77+
extra.AddExtraToSecret(secret)
78+
return nil
79+
})
7880
}
79-
wgExtras.Wait()
81+
_ = g.Wait()
8082
}
8183

8284
func ProcessValidationAndScoreWithValidation(engine engine.IEngine) {
8385
defer Channels.WaitGroup.Done()
8486

85-
wgValidation := &sync.WaitGroup{}
87+
g := errgroup.Group{}
88+
g.SetLimit(10)
8689
for secret := range ValidationChan {
87-
wgValidation.Add(2)
88-
go func(secret *secrets.Secret, wg *sync.WaitGroup) {
89-
engine.RegisterForValidation(secret, wg)
90-
engine.Score(secret, true, wg)
91-
}(secret, wgValidation)
90+
g.Go(func() error {
91+
engine.RegisterForValidation(secret)
92+
engine.Score(secret, true)
93+
return nil
94+
})
9295
}
93-
wgValidation.Wait()
94-
96+
_ = g.Wait()
9597
engine.Validate()
9698
}
9799

98100
func ProcessScoreWithoutValidation(engine engine.IEngine) {
99101
defer Channels.WaitGroup.Done()
100102

101-
wgScore := &sync.WaitGroup{}
103+
g := errgroup.Group{}
104+
g.SetLimit(10)
102105
for secret := range CvssScoreWithoutValidationChan {
103-
wgScore.Add(1)
104-
go engine.Score(secret, false, wgScore)
106+
g.Go(func() error {
107+
engine.Score(secret, false)
108+
return nil
109+
})
105110
}
106-
wgScore.Wait()
111+
_ = g.Wait()
107112
}

‎engine/engine.go‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ import (
1111
"os"
1212
"regexp"
1313
"strings"
14-
"sync"
1514
"text/tabwriter"
1615

1716
"github.com/checkmarx/2ms/v3/engine/chunk"
@@ -45,8 +44,8 @@ type IEngine interface {
4544
DetectFragment(item plugins.ISourceItem, secretsChannel chan *secrets.Secret, pluginName string) error
4645
DetectFile(ctx context.Context, item plugins.ISourceItem, secretsChannel chan *secrets.Secret) error
4746
AddRegexRules(patterns []string) error
48-
RegisterForValidation(secret *secrets.Secret, wg *sync.WaitGroup)
49-
Score(secret *secrets.Secret, validateFlag bool, wg *sync.WaitGroup)
47+
RegisterForValidation(secret *secrets.Secret)
48+
Score(secret *secrets.Secret, validateFlag bool)
5049
Validate()
5150
GetRuleBaseRiskScore(ruleId string) float64
5251
}
@@ -252,13 +251,11 @@ func (e *Engine) AddRegexRules(patterns []string) error {
252251
return nil
253252
}
254253

255-
func (e *Engine) RegisterForValidation(secret *secrets.Secret, wg *sync.WaitGroup) {
256-
defer wg.Done()
254+
func (e *Engine) RegisterForValidation(secret *secrets.Secret) {
257255
e.validator.RegisterForValidation(secret)
258256
}
259257

260-
func (e *Engine) Score(secret *secrets.Secret, validateFlag bool, wg *sync.WaitGroup) {
261-
defer wg.Done()
258+
func (e *Engine) Score(secret *secrets.Secret, validateFlag bool) {
262259
validationStatus := secrets.UnknownResult // default validity
263260
if validateFlag {
264261
validationStatus = secret.ValidationStatus

‎engine/engine_mock.go‎

Lines changed: 8 additions & 9 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎engine/engine_test.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,13 @@ import (
44
"bytes"
55
"context"
66
"fmt"
7-
"go.uber.org/mock/gomock"
87
"io"
98
"os"
109
"path/filepath"
1110
"testing"
1211

12+
"go.uber.org/mock/gomock"
13+
1314
"github.com/checkmarx/2ms/v3/engine/chunk"
1415
"github.com/checkmarx/2ms/v3/engine/rules"
1516
"github.com/checkmarx/2ms/v3/engine/semaphore"

‎engine/extra/extra.go‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"encoding/json"
66
"fmt"
77
"strings"
8-
"sync"
98

109
"github.com/checkmarx/2ms/v3/lib/secrets"
1110
)
@@ -16,8 +15,7 @@ var ruleIDToFunction = map[string]addExtraFunc{
1615
"jwt": addExtraJWT,
1716
}
1817

19-
func AddExtraToSecret(secret *secrets.Secret, wg *sync.WaitGroup) {
20-
defer wg.Done()
18+
func AddExtraToSecret(secret *secrets.Secret) {
2119
if addExtra, ok := ruleIDToFunction[secret.RuleID]; ok {
2220
extraData := addExtra(secret)
2321
if extraData != nil && extraData != "" {

‎engine/extra/extra_test.go‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,10 @@ package extra
33
import (
44
"encoding/base64"
55
"fmt"
6+
"testing"
7+
68
"github.com/checkmarx/2ms/v3/lib/secrets"
79
"github.com/stretchr/testify/assert"
8-
"sync"
9-
"testing"
1010
)
1111

1212
func TestAddExtraToSecret(t *testing.T) {
@@ -50,10 +50,7 @@ func TestAddExtraToSecret(t *testing.T) {
5050
ExtraDetails: make(map[string]interface{}),
5151
}
5252

53-
var wg sync.WaitGroup
54-
wg.Add(1)
55-
AddExtraToSecret(secret, &wg)
56-
wg.Wait()
53+
AddExtraToSecret(secret)
5754

5855
assert.Equal(t, tt.expectedOutput, secret.ExtraDetails["secretDetails"])
5956
})

‎engine/score/score_test.go‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
11
package score_test
22

33
import (
4+
"sync"
5+
"testing"
6+
47
. "github.com/checkmarx/2ms/v3/engine"
58
"github.com/checkmarx/2ms/v3/engine/rules"
69
"github.com/checkmarx/2ms/v3/engine/score"
710
"github.com/checkmarx/2ms/v3/lib/secrets"
811
"github.com/stretchr/testify/assert"
912
ruleConfig "github.com/zricethezav/gitleaks/v8/cmd/generate/config/rules"
10-
"sync"
11-
"testing"
1213
)
1314

1415
func TestScore(t *testing.T) {
@@ -216,9 +217,9 @@ func TestScore(t *testing.T) {
216217
expectedRuleScores := expectedCvssScores[secret.RuleID]
217218
validityIndex := getValidityIndex(secret.ValidationStatus)
218219
unknownIndex := getValidityIndex(secrets.UnknownResult)
219-
engine.Score(secret, true, &wg)
220+
engine.Score(secret, true)
220221
assert.Equal(t, expectedRuleScores[validityIndex], secret.CvssScore, "rule: %s", secret.RuleID)
221-
engine.Score(secret, false, &wg)
222+
engine.Score(secret, false)
222223
assert.Equal(t, expectedRuleScores[unknownIndex], secret.CvssScore, "rule: %s", secret.RuleID)
223224
}
224225
}

‎engine/validation/pairs.go‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
package validation
22

33
import (
4-
"sync"
5-
64
"github.com/checkmarx/2ms/v3/lib/secrets"
75
)
86

@@ -38,8 +36,7 @@ func (p *pairsCollector) addIfNeeded(secret *secrets.Secret) bool {
3836
return true
3937
}
4038

41-
func (p *pairsCollector) validate(generalKey string, rulesById pairsByRuleId, wg *sync.WaitGroup) {
42-
defer wg.Done()
39+
func (p *pairsCollector) validate(generalKey string, rulesById pairsByRuleId) {
4340
generalKeyToValidation[generalKey](rulesById)
4441
}
4542

‎engine/validation/validator.go‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
package validation
22

33
import (
4-
"sync"
5-
64
"github.com/checkmarx/2ms/v3/engine/extra"
75
"github.com/checkmarx/2ms/v3/lib/secrets"
86
)
@@ -35,14 +33,11 @@ func (v *Validator) RegisterForValidation(secret *secrets.Secret) {
3533
}
3634

3735
func (v *Validator) Validate() {
38-
wg := &sync.WaitGroup{}
3936
for generalKey, bySource := range v.pairsCollector.pairs {
4037
for _, byRule := range bySource {
41-
wg.Add(1)
42-
v.pairsCollector.validate(generalKey, byRule, wg)
38+
v.pairsCollector.validate(generalKey, byRule)
4339
}
4440
}
45-
wg.Wait()
4641
}
4742

4843
func IsCanValidateRule(ruleID string) bool {

0 commit comments

Comments
 (0)