Merge pull request #12 from wgpsec/dev

fix: fix a bug
wgpsec · May 13, 2024 · 1f34480 · 1f34480
2 parents b9bbd80 + 23dbe74
commit 1f34480
Show file tree

Hide file tree

Showing 15 changed files with 41 additions and 30 deletions.
diff --git a/cmd/runner.go b/cmd/runner.go
@@ -74,7 +74,7 @@ func (r *Runner) Enumerate() {
 			continue
 		}
 		var Count int
-		for _, instance := range instances.Items {
+		for _, instance := range instances.GetItems() {
 			builder.Reset()
 			if instance.DNSName != "" {
 				Count++

diff --git a/pkg/providers/aliyun/aliyun.go b/pkg/providers/aliyun/aliyun.go
@@ -123,7 +123,7 @@ func (p *Provider) Resources(ctx context.Context) (*schema.Resources, error) {
 	var err error
 	ecsProvider := &instanceProvider{id: p.id, provider: p.provider, ecsRegions: p.ecsRegions, config: p.config}
 	ecsList, err := ecsProvider.GetEcsResource(ctx)
-	gologger.Info().Msgf("获取到 %d 条阿里云 ECS 信息", len(ecsList.Items))
+	gologger.Info().Msgf("获取到 %d 条阿里云 ECS 信息", len(ecsList.GetItems()))
 	if err != nil {
 		return nil, err
 	}
@@ -132,14 +132,14 @@ func (p *Provider) Resources(ctx context.Context) (*schema.Resources, error) {
 	if err != nil {
 		return nil, err
 	}
-	gologger.Info().Msgf("获取到 %d 条阿里云 RDS 信息", len(rdsList.Items))
+	gologger.Info().Msgf("获取到 %d 条阿里云 RDS 信息", len(rdsList.GetItems()))
 
 	ossProvider := &ossProvider{ossClient: p.ossClient, id: p.id, provider: p.provider}
 	buckets, err := ossProvider.GetResource(ctx)
 	if err != nil {
 		return nil, err
 	}
-	gologger.Info().Msgf("获取到 %d 条阿里云 OSS 信息", len(buckets.Items))
+	gologger.Info().Msgf("获取到 %d 条阿里云 OSS 信息", len(buckets.GetItems()))
 
 	finalList := schema.NewResources()
 	finalList.Merge(ecsList)

diff --git a/pkg/providers/aliyun/ecs.go b/pkg/providers/aliyun/ecs.go
@@ -32,7 +32,7 @@ func (d *instanceProvider) GetEcsResource(ctx context.Context) (*schema.Resource
 		regions = append(regions, region.RegionId)
 	}
 
-	taskCh := make(chan string)
+	taskCh := make(chan string, threads)
 	for i := 0; i < threads; i++ {
 		wg.Add(1)
 		go func() {
@@ -51,7 +51,7 @@ func (d *instanceProvider) GetEcsResource(ctx context.Context) (*schema.Resource
 }
 
 func (d *instanceProvider) describeEcsInstances(ch <-chan string, wg *sync.WaitGroup) error {
-	wg.Done()
+	defer wg.Done()
 	var (
 		err       error
 		ecsClient *ecs.Client

diff --git a/pkg/providers/aliyun/rds.go b/pkg/providers/aliyun/rds.go
@@ -40,7 +40,7 @@ func (d *dbInstanceProvider) GetRdsResource(ctx context.Context) (*schema.Resour
 	}
 	regions = utils.RemoveRepeatedElement(regions)
 
-	taskCh := make(chan string)
+	taskCh := make(chan string, threads)
 	for i := 0; i < threads; i++ {
 		wg.Add(1)
 		go func() {
@@ -63,7 +63,7 @@ func (d *dbInstanceProvider) GetRdsResource(ctx context.Context) (*schema.Resour
 }
 
 func (d *dbInstanceProvider) describeRdsInstances(ch <-chan string, wg *sync.WaitGroup) error {
-	wg.Done()
+	defer wg.Done()
 	var (
 		err       error
 		rdsClient *rds.Client

diff --git a/pkg/providers/baidu/baidu.go b/pkg/providers/baidu/baidu.go
@@ -90,13 +90,13 @@ func (p *Provider) Resources(ctx context.Context) (*schema.Resources, error) {
 	if err != nil {
 		return nil, err
 	}
-	gologger.Info().Msgf("获取到 %d 条百度云 BCC 信息", len(lists.Items))
+	gologger.Info().Msgf("获取到 %d 条百度云 BCC 信息", len(lists.GetItems()))
 	bosProvider := &bosProvider{bosClient: p.bosClient, id: p.id, provider: p.provider}
 	buckets, err := bosProvider.GetResource(ctx)
 	if err != nil {
 		return nil, err
 	}
-	gologger.Info().Msgf("获取到 %d 条百度云 BOS 信息", len(buckets.Items))
+	gologger.Info().Msgf("获取到 %d 条百度云 BOS 信息", len(buckets.GetItems()))
 	finalList := schema.NewResources()
 	finalList.Merge(lists)
 	finalList.Merge(buckets)

diff --git a/pkg/providers/baidu/bcc.go b/pkg/providers/baidu/bcc.go
@@ -36,7 +36,7 @@ func (d *instanceProvider) GetResource(ctx context.Context) (*schema.Resources,
 	}
 	threads = schema.GetThreads()
 
-	taskCh := make(chan string)
+	taskCh := make(chan string, threads)
 	for i := 0; i < threads; i++ {
 		wg.Add(1)
 		go func() {
@@ -55,7 +55,7 @@ func (d *instanceProvider) GetResource(ctx context.Context) (*schema.Resources,
 }
 
 func (d *instanceProvider) describeInstances(ch <-chan string, wg *sync.WaitGroup) error {
-	wg.Done()
+	defer wg.Done()
 	var (
 		err       error
 		bccClient *bcc.Client

diff --git a/pkg/providers/huawei/huawei.go b/pkg/providers/huawei/huawei.go
@@ -57,7 +57,7 @@ func (p *Provider) Resources(ctx context.Context) (*schema.Resources, error) {
 	if err != nil {
 		return nil, err
 	}
-	gologger.Info().Msgf("获取到 %d 条华为云 OBS 信息", len(buckets.Items))
+	gologger.Info().Msgf("获取到 %d 条华为云 OBS 信息", len(buckets.GetItems()))
 	finalList := schema.NewResources()
 	finalList.Merge(buckets)
 	return finalList, nil

diff --git a/pkg/providers/liantong/liantong.go b/pkg/providers/liantong/liantong.go
@@ -55,7 +55,7 @@ func (p *Provider) Resources(ctx context.Context) (*schema.Resources, error) {
 	if err != nil {
 		return nil, err
 	}
-	gologger.Info().Msgf("获取到 %d 条联通云 OSS 信息", len(buckets.Items))
+	gologger.Info().Msgf("获取到 %d 条联通云 OSS 信息", len(buckets.GetItems()))
 	finalList := schema.NewResources()
 	finalList.Merge(buckets)
 	return finalList, nil

diff --git a/pkg/providers/liantong/oss.go b/pkg/providers/liantong/oss.go
@@ -49,7 +49,7 @@ func (d *s3Provider) GetResource(ctx context.Context) (*schema.Resources, error)
 	}
 	threads = schema.GetThreads()
 
-	taskCh := make(chan regions)
+	taskCh := make(chan regions, threads)
 	for i := 0; i < threads; i++ {
 		wg.Add(1)
 		go func() {
@@ -69,7 +69,7 @@ func (d *s3Provider) GetResource(ctx context.Context) (*schema.Resources, error)
 }
 
 func (d *s3Provider) listBuckets(ch <-chan regions, wg *sync.WaitGroup) error {
-	wg.Done()
+	defer wg.Done()
 	var err error
 	for region := range ch {
 		config := aws.NewConfig()

diff --git a/pkg/providers/qiniu/qiniu.go b/pkg/providers/qiniu/qiniu.go
@@ -43,7 +43,7 @@ func (p *Provider) Resources(ctx context.Context) (*schema.Resources, error) {
 	if err != nil {
 		return nil, err
 	}
-	gologger.Info().Msgf("获取到 %d 条七牛云 Kodo 对象存储信息", len(buckets.Items))
+	gologger.Info().Msgf("获取到 %d 条七牛云 Kodo 对象存储信息", len(buckets.GetItems()))
 	finalList := schema.NewResources()
 	finalList.Merge(buckets)
 	return finalList, nil

diff --git a/pkg/providers/tencent/cvm.go b/pkg/providers/tencent/cvm.go
@@ -33,7 +33,7 @@ func (d *instanceProvider) GetCVMResource(ctx context.Context) (*schema.Resource
 		regions = append(regions, *region.Region)
 	}
 
-	taskCh := make(chan string)
+	taskCh := make(chan string, threads)
 	for i := 0; i < threads; i++ {
 		wg.Add(1)
 		go func() {
@@ -52,7 +52,7 @@ func (d *instanceProvider) GetCVMResource(ctx context.Context) (*schema.Resource
 }
 
 func (d *instanceProvider) describeCVMInstances(ch <-chan string, wg *sync.WaitGroup) error {
-	wg.Done()
+	defer wg.Done()
 	var (
 		err       error
 		cvmClient *cvm.Client
@@ -63,7 +63,6 @@ func (d *instanceProvider) describeCVMInstances(ch <-chan string, wg *sync.WaitG
 		cpf.HttpProfile.Endpoint = "cvm.tencentcloudapi.com"
 		cvmClient, err = cvm.NewClient(d.credential, region, cpf)
 		if err != nil {
-			wg.Done()
 			continue
 		}
 		request := cvm.NewDescribeInstancesRequest()

diff --git a/pkg/providers/tencent/lh.go b/pkg/providers/tencent/lh.go
@@ -23,7 +23,7 @@ func (d *instanceProvider) GetLHResource(ctx context.Context) (*schema.Resources
 	for _, region := range d.lhRegions {
 		regions = append(regions, *region.Region)
 	}
-	taskCh := make(chan string)
+	taskCh := make(chan string, threads)
 	for i := 0; i < threads; i++ {
 		wg.Add(1)
 		go func() {
@@ -42,7 +42,7 @@ func (d *instanceProvider) GetLHResource(ctx context.Context) (*schema.Resources
 }
 
 func (d *instanceProvider) describeLHInstances(ch <-chan string, wg *sync.WaitGroup) error {
-	wg.Done()
+	defer wg.Done()
 	var (
 		err      error
 		lhClient *lh.Client

diff --git a/pkg/providers/tencent/tencent.go b/pkg/providers/tencent/tencent.go
@@ -101,21 +101,21 @@ func (p *Provider) Resources(ctx context.Context) (*schema.Resources, error) {
 	if err != nil {
 		return nil, err
 	}
-	gologger.Info().Msgf("获取到 %d 条腾讯云 CVM 信息", len(cvmList.Items))
+	gologger.Info().Msgf("获取到 %d 条腾讯云 CVM 信息", len(cvmList.GetItems()))
 
 	lhProvider := &instanceProvider{id: p.id, provider: p.provider, cvmRegions: p.cvmRegions, lhRegions: p.lhRegions, credential: p.credential}
 	lhList, err := lhProvider.GetLHResource(ctx)
 	if err != nil {
 		return nil, err
 	}
-	gologger.Info().Msgf("获取到 %d 条腾讯云 LH 信息", len(lhList.Items))
+	gologger.Info().Msgf("获取到 %d 条腾讯云 LH 信息", len(lhList.GetItems()))
 
 	cosProvider := &cosProvider{provider: p.provider, id: p.id, cosClient: p.cosClient}
 	cosList, err := cosProvider.GetResource(ctx)
 	if err != nil {
 		return nil, err
 	}
-	gologger.Info().Msgf("获取到 %d 条腾讯云 COS 信息", len(cosList.Items))
+	gologger.Info().Msgf("获取到 %d 条腾讯云 COS 信息", len(cosList.GetItems()))
 
 	finalList := schema.NewResources()
 	finalList.Merge(cvmList)

diff --git a/pkg/providers/tianyi/tianyi.go b/pkg/providers/tianyi/tianyi.go
@@ -49,7 +49,7 @@ func (p *Provider) Resources(ctx context.Context) (*schema.Resources, error) {
 	if err != nil {
 		return nil, err
 	}
-	gologger.Info().Msgf("获取到 %d 条天翼云 OOS 对象存储信息", len(buckets.Items))
+	gologger.Info().Msgf("获取到 %d 条天翼云 OOS 对象存储信息", len(buckets.GetItems()))
 	finalList := schema.NewResources()
 	finalList.Merge(buckets)
 	return finalList, nil

diff --git a/pkg/schema/schema.go b/pkg/schema/schema.go
@@ -14,7 +14,19 @@ var validator *validate.Validator
 var Threads int
 
 type Resources struct {
-	Items []*Resource
+	items []*Resource
+	sync.RWMutex
+}
+
+func (r *Resources) AppendItem(item *Resource) {
+	r.Lock()
+	defer r.Unlock()
+	r.items = append(r.items, item)
+}
+func (r *Resources) GetItems() []*Resource {
+	r.RLock()
+	defer r.RUnlock()
+	return r.items
 }
 
 type Provider interface {
@@ -80,7 +92,7 @@ func (r *Resources) appendResourceWithTypeAndMeta(resourceType validate.Resource
 	default:
 		return
 	}
-	r.Items = append(r.Items, resource)
+	r.AppendItem(resource)
 }
 
 func (r *Resources) Append(resource *Resource) {
@@ -92,7 +104,7 @@ func (r *Resources) Merge(resources *Resources) {
 		return
 	}
 	mergeUniqueMap := &sync.Map{}
-	for _, item := range resources.Items {
+	for _, item := range resources.GetItems() {
 		r.appendResource(item, mergeUniqueMap)
 	}
 }
@@ -116,7 +128,7 @@ func (o OptionBlock) GetMetadata(key string) (string, bool) {
 // Other
 
 func NewResources() *Resources {
-	return &Resources{Items: make([]*Resource, 0)}
+	return &Resources{items: make([]*Resource, 0)}
 }
 
 func SetThreads(threads int) {