Skip to content

Commit 54c5450

Browse files
Ash-expAsh-exp
committed
fix: streamline AWS session creation and improve error handling
1 parent dcbc830 commit 54c5450

7 files changed

Lines changed: 120 additions & 98 deletions

File tree

‎ci-runner/executor/util/envUtils.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,14 @@ package util
1919
import (
2020
"encoding/json"
2121
"fmt"
22+
"os"
23+
"strconv"
24+
"strings"
25+
2226
"github.com/caarlos0/env"
2327
"github.com/devtron-labs/ci-runner/helper"
2428
"github.com/devtron-labs/ci-runner/pubsub"
2529
"github.com/devtron-labs/ci-runner/util"
26-
"os"
27-
"strconv"
28-
"strings"
2930
)
3031

3132
type ScriptEnvVariables struct {
@@ -96,6 +97,7 @@ func GetGlobalEnvVariables(ciCdRequest *helper.CiCdTriggerEvent) (*ScriptEnvVari
9697
envs["ACCESS_KEY"] = ciCdRequest.CommonWorkflowRequest.AccessKey
9798
envs["SECRET_KEY"] = ciCdRequest.CommonWorkflowRequest.SecretKey
9899
envs["AWS_REGION"] = ciCdRequest.CommonWorkflowRequest.AwsRegion
100+
envs["ASSUME_ROLE_ARN"] = ciCdRequest.CommonWorkflowRequest.AssumeRoleArn
99101
envs["LAST_FETCHED_TIME"] = ciCdRequest.CommonWorkflowRequest.CiArtifactLastFetch.String()
100102

101103
//adding some envs for Image scanning plugin

‎ci-runner/helper/DockerHelper.go‎

Lines changed: 24 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,21 @@ import (
2323
"encoding/json"
2424
"errors"
2525
"fmt"
26+
"io"
27+
"io/ioutil"
28+
"log"
29+
"os"
30+
"os/exec"
31+
"path"
32+
"path/filepath"
33+
"strconv"
34+
"strings"
35+
"sync"
36+
"syscall"
37+
"time"
38+
2639
"github.com/aws/aws-sdk-go/aws"
2740
"github.com/aws/aws-sdk-go/aws/credentials"
28-
"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
2941
"github.com/aws/aws-sdk-go/aws/session"
3042
"github.com/aws/aws-sdk-go/service/ecr"
3143
"github.com/aws/aws-sdk-go/service/sts"
@@ -39,18 +51,6 @@ import (
3951
"github.com/devtron-labs/common-lib/utils/dockerOperations"
4052
"github.com/devtron-labs/common-lib/utils/retryFunc"
4153
"golang.org/x/sync/errgroup"
42-
"io"
43-
"io/ioutil"
44-
"log"
45-
"os"
46-
"os/exec"
47-
"path"
48-
"path/filepath"
49-
"strconv"
50-
"strings"
51-
"sync"
52-
"syscall"
53-
"time"
5454
)
5555

5656
const (
@@ -236,25 +236,22 @@ func (impl *DockerHelperImpl) DockerLogin(ciContext cicxt.CiContext, dockerCrede
236236
accessKey, secretKey := dockerCredentials.AccessKey, dockerCredentials.SecretKey
237237
//fmt.Printf("accessKey %s, secretKey %s\n", accessKey, secretKey)
238238

239-
var creds *credentials.Credentials
239+
var sess *session.Session
240+
var err error
240241

241-
if len(dockerCredentials.AccessKey) == 0 || len(dockerCredentials.SecretKey) == 0 {
242-
//fmt.Println("empty accessKey or secretKey")
243-
sess, err := session.NewSession(&aws.Config{
242+
if len(accessKey) == 0 || len(secretKey) == 0 {
243+
// Case 1: IAM role — use default credential chain (IRSA, instance profile, task role, env vars)
244+
sess, err = session.NewSession(&aws.Config{
244245
Region: &dockerCredentials.AwsRegion,
245246
})
246-
if err != nil {
247-
log.Println(err)
248-
return err
249-
}
250-
creds = ec2rolecreds.NewCredentials(sess)
251247
} else {
252-
creds = credentials.NewStaticCredentials(accessKey, secretKey, "")
248+
// Case 2: Static credentials
249+
creds := credentials.NewStaticCredentials(accessKey, secretKey, "")
250+
sess, err = session.NewSession(&aws.Config{
251+
Region: &dockerCredentials.AwsRegion,
252+
Credentials: creds,
253+
})
253254
}
254-
sess, err := session.NewSession(&aws.Config{
255-
Region: &dockerCredentials.AwsRegion,
256-
Credentials: creds,
257-
})
258255
if err != nil {
259256
log.Println(err)
260257
return err

‎common-lib/helmLib/registry/common.go‎

Lines changed: 21 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -4,20 +4,20 @@ import (
44
"crypto/tls"
55
"encoding/base64"
66
"fmt"
7+
"log"
8+
"math/rand"
9+
"net/http"
10+
"net/url"
11+
"os"
12+
"strings"
13+
714
"github.com/aws/aws-sdk-go/aws"
815
"github.com/aws/aws-sdk-go/aws/credentials"
9-
"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
1016
"github.com/aws/aws-sdk-go/aws/session"
1117
"github.com/aws/aws-sdk-go/service/ecr"
1218
"github.com/aws/aws-sdk-go/service/sts"
1319
http2 "github.com/devtron-labs/common-lib/utils/http"
1420
"helm.sh/helm/v3/pkg/registry"
15-
"log"
16-
"math/rand"
17-
"net/http"
18-
"net/url"
19-
"os"
20-
"strings"
2121
)
2222

2323
func GetLoggedInClient(client *registry.Client, config *Configuration) (*registry.Client, error) {
@@ -96,26 +96,24 @@ func extractCredentialsForRegistry(config *Configuration) (string, string, error
9696
}
9797
if config.RegistryType == REGISTRY_TYPE_ECR {
9898
accessKey, secretKey := config.AwsAccessKey, config.AwsSecretKey
99-
var creds *credentials.Credentials
99+
var sess *session.Session
100+
var err error
100101

101-
if len(config.AwsAccessKey) == 0 || len(config.AwsSecretKey) == 0 {
102-
sess, err := session.NewSession(&aws.Config{
102+
if len(accessKey) == 0 || len(secretKey) == 0 {
103+
// Case 1: IAM role — use default credential chain (IRSA, instance profile, task role, env vars)
104+
sess, err = session.NewSession(&aws.Config{
103105
Region: &config.AwsRegion,
104106
})
105-
if err != nil {
106-
log.Printf("error in creating AWS client %w ", err)
107-
return "", "", err
108-
}
109-
creds = ec2rolecreds.NewCredentials(sess)
110107
} else {
111-
creds = credentials.NewStaticCredentials(accessKey, secretKey, "")
108+
// Case 2: Static credentials
109+
creds := credentials.NewStaticCredentials(accessKey, secretKey, "")
110+
sess, err = session.NewSession(&aws.Config{
111+
Region: &config.AwsRegion,
112+
Credentials: creds,
113+
})
112114
}
113-
sess, err := session.NewSession(&aws.Config{
114-
Region: &config.AwsRegion,
115-
Credentials: creds,
116-
})
117115
if err != nil {
118-
log.Printf("error in creating AWS client %w ", err)
116+
log.Printf("error in creating AWS client %v ", err)
119117
return "", "", err
120118
}
121119

@@ -149,14 +147,14 @@ func extractCredentialsForRegistry(config *Configuration) (string, string, error
149147
input := &ecr.GetAuthorizationTokenInput{}
150148
authData, err := svc.GetAuthorizationToken(input)
151149
if err != nil {
152-
log.Printf("error in creating AWS client %w ", err)
150+
log.Printf("error in creating AWS client %v ", err)
153151
return "", "", err
154152
}
155153
// decode token
156154
token := authData.AuthorizationData[0].AuthorizationToken
157155
decodedToken, err := base64.StdEncoding.DecodeString(*token)
158156
if err != nil {
159-
log.Printf("error in creating AWS client %w ", err)
157+
log.Printf("error in creating AWS client %v ", err)
160158
return "", "", err
161159
}
162160
credsSlice := strings.Split(string(decodedToken), ":")

‎common-lib/utils/registry/bean.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,4 +25,5 @@ type RegistryCredential struct {
2525
AWSAccessKeyId string `json:"awsAccessKeyId,omitempty"`
2626
AWSSecretAccessKey string `json:"awsSecretAccessKey,omitempty"`
2727
AWSRegion string `json:"awsRegion,omitempty"`
28+
AssumeRoleArn string `json:"assumeRoleArn,omitempty"`
2829
}

‎common-lib/utils/registry/extractCredentials.go‎

Lines changed: 41 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,13 @@ package registry
33
import (
44
"encoding/base64"
55
"fmt"
6+
"strings"
7+
68
"github.com/aws/aws-sdk-go/aws"
79
"github.com/aws/aws-sdk-go/aws/credentials"
8-
"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
910
"github.com/aws/aws-sdk-go/aws/session"
1011
"github.com/aws/aws-sdk-go/service/ecr"
11-
"strings"
12+
"github.com/aws/aws-sdk-go/service/sts"
1213
)
1314

1415
func ExtractCredentialsForRegistry(registryCredential *RegistryCredential) (string, string, error) {
@@ -24,28 +25,53 @@ func ExtractCredentialsForRegistry(registryCredential *RegistryCredential) (stri
2425
}
2526
if registryCredential.RegistryType == DOCKER_REGISTRY_TYPE_ECR {
2627
accessKey, secretKey := registryCredential.AWSAccessKeyId, registryCredential.AWSSecretAccessKey
27-
var creds *credentials.Credentials
28+
var sess *session.Session
29+
var err error
2830

29-
if len(registryCredential.AWSAccessKeyId) == 0 || len(registryCredential.AWSSecretAccessKey) == 0 {
30-
sess, err := session.NewSession(&aws.Config{
31+
if len(accessKey) == 0 || len(secretKey) == 0 {
32+
// Case 1: IAM role — use default credential chain (IRSA, instance profile, task role, env vars)
33+
sess, err = session.NewSession(&aws.Config{
3134
Region: &registryCredential.AWSRegion,
3235
})
33-
if err != nil {
34-
fmt.Printf("Error in creating AWS client", "err", err)
35-
return "", "", err
36-
}
37-
creds = ec2rolecreds.NewCredentials(sess)
3836
} else {
39-
creds = credentials.NewStaticCredentials(accessKey, secretKey, "")
37+
// Case 2: Static credentials
38+
creds := credentials.NewStaticCredentials(accessKey, secretKey, "")
39+
sess, err = session.NewSession(&aws.Config{
40+
Region: &registryCredential.AWSRegion,
41+
Credentials: creds,
42+
})
4043
}
41-
sess, err := session.NewSession(&aws.Config{
42-
Region: &registryCredential.AWSRegion,
43-
Credentials: creds,
44-
})
4544
if err != nil {
4645
fmt.Println("Error in creating AWS client session", "err", err)
4746
return "", "", err
4847
}
48+
49+
// Case 3: AssumeRole (cross-account) — layered on top of Case 1 or 2
50+
if len(registryCredential.AssumeRoleArn) > 0 {
51+
stsClient := sts.New(sess)
52+
assumeOutput, err := stsClient.AssumeRole(&sts.AssumeRoleInput{
53+
RoleArn: aws.String(registryCredential.AssumeRoleArn),
54+
RoleSessionName: aws.String("devtron-ecr-cross-account"),
55+
})
56+
if err != nil {
57+
fmt.Printf("Error in assuming role %s: %v", registryCredential.AssumeRoleArn, err)
58+
return "", "", err
59+
}
60+
assumedCreds := credentials.NewStaticCredentials(
61+
*assumeOutput.Credentials.AccessKeyId,
62+
*assumeOutput.Credentials.SecretAccessKey,
63+
*assumeOutput.Credentials.SessionToken,
64+
)
65+
sess, err = session.NewSession(&aws.Config{
66+
Region: &registryCredential.AWSRegion,
67+
Credentials: assumedCreds,
68+
})
69+
if err != nil {
70+
fmt.Println("Error in creating AWS session with assumed role credentials", "err", err)
71+
return "", "", err
72+
}
73+
}
74+
4975
svc := ecr.New(sess)
5076
input := &ecr.GetAuthorizationTokenInput{}
5177
authData, err := svc.GetAuthorizationToken(input)

‎image-scanner/pkg/klarService/KlarService.go‎

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,10 @@ package klarService
1818

1919
import (
2020
"fmt"
21+
"strings"
22+
2123
"github.com/aws/aws-sdk-go/aws"
2224
"github.com/aws/aws-sdk-go/aws/credentials"
23-
"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
2425
"github.com/aws/aws-sdk-go/aws/session"
2526
"github.com/aws/aws-sdk-go/service/ecr"
2627
bean2 "github.com/devtron-labs/common-lib/imageScan/bean"
@@ -29,18 +30,20 @@ import (
2930
"github.com/devtron-labs/image-scanner/pkg/security"
3031
"github.com/devtron-labs/image-scanner/pkg/sql/bean"
3132
"github.com/devtron-labs/image-scanner/pkg/sql/repository"
32-
"strings"
3333

3434
"errors"
35+
36+
"time"
37+
3538
"github.com/caarlos0/env/v6"
39+
3640
/*"github.com/devtron-labs/image-scanner/client"*/
3741
/*"github.com/devtron-labs/image-scanner/client"*/
3842
"github.com/devtron-labs/image-scanner/pkg/grafeasService"
3943
"github.com/optiopay/klar/clair"
4044
"github.com/optiopay/klar/docker"
4145
"go.uber.org/zap"
4246
"golang.org/x/oauth2/google"
43-
"time"
4447
)
4548

4649
type KlarConfig struct {
@@ -114,23 +117,21 @@ func (impl *KlarServiceImpl) Process(scanEvent *bean2.ImageScanEvent, executionH
114117
tokenAddr := &tokenData
115118
if dockerRegistry.RegistryType == repository.REGISTRYTYPE_ECR {
116119
accessKey, secretKey := dockerRegistry.AWSAccessKeyId, dockerRegistry.AWSSecretAccessKey.String()
117-
var creds *credentials.Credentials
118-
if len(dockerRegistry.AWSAccessKeyId) == 0 || len(dockerRegistry.AWSSecretAccessKey) == 0 {
119-
sess, err := session.NewSession(&aws.Config{
120+
var sess *session.Session
121+
var err error
122+
if len(accessKey) == 0 || len(secretKey) == 0 {
123+
// Case 1: IAM role — use default credential chain (IRSA, instance profile, task role, env vars)
124+
sess, err = session.NewSession(&aws.Config{
120125
Region: &dockerRegistry.AWSRegion,
121126
})
122-
if err != nil {
123-
impl.logger.Errorw("error in starting aws new session", "err", err)
124-
return nil, err
125-
}
126-
creds = ec2rolecreds.NewCredentials(sess)
127127
} else {
128-
creds = credentials.NewStaticCredentials(accessKey, secretKey, "")
128+
// Case 2: Static credentials
129+
creds := credentials.NewStaticCredentials(accessKey, secretKey, "")
130+
sess, err = session.NewSession(&aws.Config{
131+
Region: &dockerRegistry.AWSRegion,
132+
Credentials: creds,
133+
})
129134
}
130-
sess, err := session.NewSession(&aws.Config{
131-
Region: &dockerRegistry.AWSRegion,
132-
Credentials: creds,
133-
})
134135
if err != nil {
135136
impl.logger.Errorw("error in starting aws new session", "err", err)
136137
return nil, err

‎image-scanner/pkg/roundTripper/RoundTripperService.go‎

Lines changed: 11 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ import (
2020
"context"
2121
"github.com/aws/aws-sdk-go/aws"
2222
"github.com/aws/aws-sdk-go/aws/credentials"
23-
"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
2423
"github.com/aws/aws-sdk-go/aws/session"
2524
"github.com/aws/aws-sdk-go/service/ecr"
2625
"github.com/aws/aws-sdk-go/service/sts"
@@ -116,23 +115,21 @@ func (impl *RoundTripperServiceImpl) GetAuthenticatorByDockerRegistryId(dockerRe
116115
}
117116
if dockerRegistry.RegistryType == repository.REGISTRYTYPE_ECR {
118117
accessKey, secretKey := dockerRegistry.AWSAccessKeyId, dockerRegistry.AWSSecretAccessKey.String()
119-
var creds *credentials.Credentials
120-
if len(dockerRegistry.AWSAccessKeyId) == 0 || len(dockerRegistry.AWSSecretAccessKey) == 0 {
121-
sess, err := session.NewSession(&aws.Config{
118+
var sess *session.Session
119+
var err error
120+
if len(accessKey) == 0 || len(secretKey) == 0 {
121+
// Case 1: IAM role — use default credential chain (IRSA, instance profile, task role, env vars)
122+
sess, err = session.NewSession(&aws.Config{
122123
Region: &dockerRegistry.AWSRegion,
123124
})
124-
if err != nil {
125-
impl.Logger.Errorw("error in starting aws new session", "err", err)
126-
return nil, nil, err
127-
}
128-
creds = ec2rolecreds.NewCredentials(sess)
129125
} else {
130-
creds = credentials.NewStaticCredentials(accessKey, secretKey, "")
126+
// Case 2: Static credentials
127+
creds := credentials.NewStaticCredentials(accessKey, secretKey, "")
128+
sess, err = session.NewSession(&aws.Config{
129+
Region: &dockerRegistry.AWSRegion,
130+
Credentials: creds,
131+
})
131132
}
132-
sess, err := session.NewSession(&aws.Config{
133-
Region: &dockerRegistry.AWSRegion,
134-
Credentials: creds,
135-
})
136133
if err != nil {
137134
impl.Logger.Errorw("error in starting aws new session", "err", err)
138135
return nil, nil, err

0 commit comments

Comments
 (0)