Skip to content

Commit bda5b2f

Browse files
committed
add lambdaName and lambdaIAMRoleName as app flags
1 parent a3793a4 commit bda5b2f

File tree

6 files changed

+64
-33
lines changed

6 files changed

+64
-33
lines changed

Dockerfile

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,17 @@ COPY --from=build-env /src/artifacts/server/linux/awslambdaproxy /app/
1010

1111
ENV AWS_ACCESS_KEY_ID=
1212
ENV AWS_SECRET_ACCESS_KEY=
13+
ENV LAMBDA_NAME=
14+
ENV LAMBDA_IAM_ROLE_NAME=
1315
ENV REGIONS=
1416
ENV FREQUENCY=
1517
ENV MEMORY=
1618
ENV SSH_USER=
1719
ENV SSH_PORT=2222
1820
ENV LISTENERS=
21+
ENV DEBUG=
1922
ENV DEBUG_PROXY=
23+
ENV BYPASS=
2024

2125
WORKDIR /app
2226

cmd/awslambdaproxy/run.go

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,10 @@ import (
1313
)
1414

1515
var (
16-
frequency time.Duration
17-
memory int
18-
debug, debugProxy bool
19-
sshUser, sshPort, regions, listeners, bypass string
16+
frequency time.Duration
17+
memory int
18+
debug, debugProxy bool
19+
lambdaName, lambdaIamRole, sshUser, sshPort, regions, listeners, bypass string
2020
)
2121

2222
// runCmd represents the run command
@@ -50,12 +50,16 @@ var runCmd = &cobra.Command{
5050
aFrequency := viper.GetDuration("frequency")
5151
aListeners := strings.Split(viper.GetString("listeners"), ",")
5252
aBypass := viper.GetString("bypass")
53+
aLambdaName := viper.GetString("lambda-name")
54+
aLambdaIamRoleName := viper.GetString("lambda-iam-role-name")
5355

5456
if _, err := server.GetSessionAWS(); err != nil {
5557
log.Fatal("unable to find valid aws credentials")
5658
}
5759

5860
s, err := server.New(server.Config{
61+
LambdaName: aLambdaName,
62+
LambdaIamRoleName: aLambdaIamRoleName,
5963
LambdaRegions: aRegions,
6064
LambdaMemory: aMemory,
6165
LambdaExecutionFrequency: aFrequency,
@@ -84,6 +88,10 @@ func getCurrentUserName() string {
8488
func init() {
8589
RootCmd.AddCommand(runCmd)
8690

91+
runCmd.Flags().StringVarP(&lambdaName, "lambda-name", "n", "awslambdaproxy",
92+
fmt.Sprintf("name of lambda function"))
93+
runCmd.Flags().StringVarP(&lambdaIamRole, "lambda-iam-role-name", "i", "awslambdaproxy-role",
94+
fmt.Sprintf("name of lambda function"))
8795
runCmd.Flags().StringVarP(&regions, "regions", "r", "us-west-2",
8896
fmt.Sprintf("comma separted list of regions to run proxy (e.g. us-west-2,us-west-1,us-east-1). "+
8997
"valid regions include %v", server.GetValidLambdaRegions()))
@@ -111,6 +119,8 @@ func init() {
111119
"comma separated list of domains/ips to bypass lambda proxy (e.g. *.websocket.org,*.youtube.com). "+
112120
"note that when using sock5 proxy mode you'll need to be remotely resolving dns for this to work.")
113121

122+
viper.BindPFlag("lambda-name", runCmd.Flags().Lookup("lambda-name"))
123+
viper.BindPFlag("lambda-iam-role-name", runCmd.Flags().Lookup("lambda-iam-role-name"))
114124
viper.BindPFlag("regions", runCmd.Flags().Lookup("regions"))
115125
viper.BindPFlag("frequency", runCmd.Flags().Lookup("frequency"))
116126
viper.BindPFlag("memory", runCmd.Flags().Lookup("memory"))

cmd/awslambdaproxy/setup.go

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55

66
"github.com/dan-v/awslambdaproxy/pkg/server"
77
"github.com/spf13/cobra"
8+
"github.com/spf13/viper"
89
)
910

1011
// setupCmd represents the setup command
@@ -13,11 +14,14 @@ var setupCmd = &cobra.Command{
1314
Short: "setup awslambdaproxy aws infrastructure",
1415
Long: `this will setup all required aws infrastructure to run awslambdaproxy.`,
1516
Run: func(cmd *cobra.Command, args []string) {
17+
18+
aLambdaIamRoleName := viper.GetString("lambda-iam-role-name")
19+
1620
if _, err := server.GetSessionAWS(); err != nil {
1721
log.Fatal("unable to find valid aws credentials")
1822
}
1923

20-
err := server.SetupLambdaInfrastructure()
24+
err := server.SetupLambdaInfrastructure(aLambdaIamRoleName)
2125
if err != nil {
2226
log.Fatal("failed to run setup for awslambdaproxy: ", err)
2327
}

pkg/server/infrastructure.go

Lines changed: 24 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -15,31 +15,30 @@ import (
1515
)
1616

1717
const (
18-
lambdaFunctionName = "awslambdaproxy"
19-
lambdaFunctionHandler = "main"
20-
lambdaFunctionRuntime = "go1.x"
21-
lambdaFunctionIamRole = "awslambdaproxy-role"
22-
lambdaFunctionIamRolePolicyName = "awslambdaproxy-role-policy"
23-
lambdaFunctionZipLocation = "artifacts/lambda.zip"
18+
lambdaFunctionHandler = "main"
19+
lambdaFunctionRuntime = "go1.x"
20+
lambdaFunctionZipLocation = "artifacts/lambda.zip"
2421
)
2522

2623
type lambdaInfrastructure struct {
2724
config *aws.Config
25+
name string
26+
iamRole string
2827
regions []string
2928
lambdaTimeout int64
3029
lambdaMemorySize int64
3130
}
3231

3332
// SetupLambdaInfrastructure sets up IAM role needed to run awslambdaproxy
34-
func SetupLambdaInfrastructure() error {
33+
func SetupLambdaInfrastructure(lambdaIamRole string) error {
3534
sess, err := GetSessionAWS()
3635
if err != nil {
3736
return err
3837
}
3938

4039
svc := iam.New(sess, &aws.Config{})
4140
_, err = svc.GetRole(&iam.GetRoleInput{
42-
RoleName: aws.String(lambdaFunctionIamRole),
41+
RoleName: aws.String(lambdaIamRole),
4342
})
4443
if err != nil {
4544
if awsErr, ok := err.(awserr.Error); ok {
@@ -53,7 +52,7 @@ func SetupLambdaInfrastructure() error {
5352
"Action": "sts:AssumeRole"
5453
}
5554
}`),
56-
RoleName: aws.String(lambdaFunctionIamRole),
55+
RoleName: aws.String(lambdaIamRole),
5756
Path: aws.String("/"),
5857
})
5958
if err != nil {
@@ -74,8 +73,8 @@ func SetupLambdaInfrastructure() error {
7473
}
7574
]
7675
}`),
77-
PolicyName: aws.String(lambdaFunctionIamRolePolicyName),
78-
RoleName: aws.String(lambdaFunctionIamRole),
76+
PolicyName: aws.String(lambdaIamRole + "-policy"),
77+
RoleName: aws.String(lambdaIamRole),
7978
})
8079
if err != nil {
8180
return err
@@ -102,10 +101,10 @@ func (infra *lambdaInfrastructure) setup() error {
102101

103102
svc := iam.New(sess, infra.config)
104103
resp, err := svc.GetRole(&iam.GetRoleInput{
105-
RoleName: aws.String(lambdaFunctionIamRole),
104+
RoleName: aws.String(infra.iamRole),
106105
})
107106
if err != nil {
108-
return errors.Wrap(err, "Could not find IAM role "+lambdaFunctionIamRole+". Probably need to run setup.")
107+
return errors.Wrap(err, "Could not find IAM role "+infra.iamRole+". Probably need to run setup.")
109108
}
110109
roleArn := *resp.Role.Arn
111110
zip, err := Asset(lambdaFunctionZipLocation)
@@ -114,16 +113,18 @@ func (infra *lambdaInfrastructure) setup() error {
114113
}
115114
for _, region := range infra.regions {
116115
log.Println("Setting up Lambda function in region: " + region)
117-
err = infra.createOrUpdateLambdaFunction(sess, region, roleArn, zip)
116+
err = infra.createOrUpdateLambdaFunction(sess, infra.name, region, roleArn, zip)
118117
if err != nil {
119118
return errors.Wrap(err, "Could not create Lambda function in region "+region)
120119
}
121120
}
122121
return nil
123122
}
124123

125-
func setupLambdaInfrastructure(regions []string, memorySize int64, timeout int64) error {
124+
func setupLambdaInfrastructure(name string, iamRole string, regions []string, memorySize int64, timeout int64) error {
126125
infra := lambdaInfrastructure{
126+
name: name,
127+
iamRole: iamRole,
127128
regions: regions,
128129
config: &aws.Config{},
129130
lambdaTimeout: timeout,
@@ -135,11 +136,11 @@ func setupLambdaInfrastructure(regions []string, memorySize int64, timeout int64
135136
return nil
136137
}
137138

138-
func (infra *lambdaInfrastructure) createOrUpdateLambdaFunction(sess *session.Session, region, roleArn string, payload []byte) error {
139+
func (infra *lambdaInfrastructure) createOrUpdateLambdaFunction(sess *session.Session, name, region, roleArn string, payload []byte) error {
139140
config := infra.config.WithRegion(region)
140141

141142
svc := lambda.New(sess, config)
142-
exists, err := lambdaExists(svc)
143+
exists, err := lambdaExists(svc, name)
143144
if err != nil {
144145
return err
145146
}
@@ -156,7 +157,7 @@ func (infra *lambdaInfrastructure) createOrUpdateLambdaFunction(sess *session.Se
156157

157158
func (infra *lambdaInfrastructure) deleteLambdaFunction(svc *lambda.Lambda) error {
158159
_, err := svc.DeleteFunction(&lambda.DeleteFunctionInput{
159-
FunctionName: aws.String(lambdaFunctionName),
160+
FunctionName: aws.String(infra.name),
160161
})
161162
if err != nil {
162163
return err
@@ -169,7 +170,7 @@ func (infra *lambdaInfrastructure) createLambdaFunction(svc *lambda.Lambda, role
169170
Code: &lambda.FunctionCode{
170171
ZipFile: payload,
171172
},
172-
FunctionName: aws.String(lambdaFunctionName),
173+
FunctionName: aws.String(infra.name),
173174
Handler: aws.String(lambdaFunctionHandler),
174175
Role: aws.String(roleArn),
175176
Runtime: aws.String(lambdaFunctionRuntime),
@@ -189,9 +190,9 @@ func (infra *lambdaInfrastructure) createLambdaFunction(svc *lambda.Lambda, role
189190
return nil
190191
}
191192

192-
func lambdaExists(svc *lambda.Lambda) (bool, error) {
193+
func lambdaExists(svc *lambda.Lambda, name string) (bool, error) {
193194
_, err := svc.GetFunction(&lambda.GetFunctionInput{
194-
FunctionName: aws.String(lambdaFunctionName),
195+
FunctionName: aws.String(name),
195196
})
196197

197198
if err != nil {
@@ -258,8 +259,8 @@ func (infra *lambdaInfrastructure) createIAMLambdaRolePolicy(sess *session.Sessi
258259
}
259260
]
260261
}`),
261-
PolicyName: aws.String(lambdaFunctionIamRolePolicyName),
262-
RoleName: aws.String(roleName),
262+
PolicyName: aws.String(infra.iamRole + "-policy"),
263+
RoleName: aws.String(infra.iamRole),
263264
})
264265
return err
265266
}

pkg/server/lambdaexecution.go

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
)
1414

1515
type lambdaExecutionManager struct {
16+
name string
1617
regions []string
1718
frequency time.Duration
1819
publicIP string
@@ -64,7 +65,7 @@ func (l *lambdaExecutionManager) executeFunction(region int, setInvokeConfig boo
6465
log.Printf("Setting invoke configuration maximumRetryAttempts=%v maximumEventAgeInSeconds=%v\n",
6566
maximumRetryAttempts, maximumEventAgeInSeconds)
6667
_, err = svc.PutFunctionEventInvokeConfig(&lambda.PutFunctionEventInvokeConfigInput{
67-
FunctionName: aws.String(lambdaFunctionName),
68+
FunctionName: aws.String(l.name),
6869
MaximumEventAgeInSeconds: aws.Int64(maximumEventAgeInSeconds),
6970
MaximumRetryAttempts: aws.Int64(maximumRetryAttempts),
7071
})
@@ -86,7 +87,7 @@ func (l *lambdaExecutionManager) executeFunction(region int, setInvokeConfig boo
8687
}
8788
payload, _ := json.Marshal(lambdaPayload)
8889
params := &lambda.InvokeInput{
89-
FunctionName: aws.String(lambdaFunctionName),
90+
FunctionName: aws.String(l.name),
9091
InvocationType: aws.String(lambda.InvocationTypeEvent),
9192
Payload: payload,
9293
}
@@ -99,9 +100,10 @@ func (l *lambdaExecutionManager) executeFunction(region int, setInvokeConfig boo
99100
return nil
100101
}
101102

102-
func newLambdaExecutionManager(publicIP string, regions []string, frequency time.Duration, sshUser string, sshPort string,
103+
func newLambdaExecutionManager(name string, publicIP string, regions []string, frequency time.Duration, sshUser string, sshPort string,
103104
privateKey []byte, onDemandExecution chan bool) (*lambdaExecutionManager, error) {
104105
executionManager := &lambdaExecutionManager{
106+
name: name,
105107
regions: regions,
106108
frequency: frequency,
107109
publicIP: publicIP,

pkg/server/server.go

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,10 @@ const (
3434

3535
// Config is used to define the configuration for Server
3636
type Config struct {
37+
// LambdaName is a name of Lambda function
38+
LambdaName string
39+
// LambdaIamRoleName is a name of Lambda function IAM role
40+
LambdaIamRoleName string
3741
// LambdaRegions is all regions to execute Lambda functions in
3842
LambdaRegions []string
3943
// LambdaMemory is the size of memory to assign Lambda function
@@ -60,6 +64,8 @@ type Config struct {
6064
// Server is the long running server component of awslambdaproxy
6165
type Server struct {
6266
publicIPClient publicip.Client
67+
lambdaName string
68+
lambdaIamRole string
6369
lambdaRegions []string
6470
lambdaMemory int64
6571
lambdaExecutionFrequency time.Duration
@@ -87,6 +93,8 @@ func New(config Config) (*Server, error) {
8793
functionTimeout := int(config.LambdaExecutionFrequency.Seconds()) + int(LambdaExecutionTimeoutBuffer.Seconds())
8894
s := &Server{
8995
publicIPClient: awspublicip.New(),
96+
lambdaName: config.LambdaName,
97+
lambdaIamRole: config.LambdaIamRoleName,
9098
lambdaRegions: config.LambdaRegions,
9199
lambdaMemory: int64(config.LambdaMemory),
92100
lambdaExecutionFrequency: config.LambdaExecutionFrequency,
@@ -102,6 +110,8 @@ func New(config Config) (*Server, error) {
102110

103111
logger.WithFields(logrus.Fields{
104112
"publicIPClient": s.publicIPClient.ProviderURL(),
113+
"lambdaName": s.lambdaName,
114+
"lambdaIamRole": s.lambdaIamRole,
105115
"lambdaRegions": s.lambdaRegions,
106116
"lambdaMemory": s.lambdaMemory,
107117
"lambdaExecutionFrequency": s.lambdaExecutionFrequency,
@@ -124,7 +134,7 @@ func (s *Server) Run() {
124134
}
125135

126136
s.logger.Infof("setting up lambda infrastructure")
127-
err = setupLambdaInfrastructure(s.lambdaRegions, s.lambdaMemory, s.lambdaTimeoutSeconds)
137+
err = setupLambdaInfrastructure(s.lambdaName, s.lambdaIamRole, s.lambdaRegions, s.lambdaMemory, s.lambdaTimeoutSeconds)
128138
if err != nil {
129139
s.logger.WithError(err).Fatalf("failed to setup lambda infrastructure")
130140
}
@@ -148,7 +158,7 @@ func (s *Server) Run() {
148158
}
149159

150160
s.logger.Println("starting lambda execution manager")
151-
_, err = newLambdaExecutionManager(publicIP, s.lambdaRegions, s.lambdaExecutionFrequency,
161+
_, err = newLambdaExecutionManager(s.lambdaName, publicIP, s.lambdaRegions, s.lambdaExecutionFrequency,
152162
s.reverseTunnelSSHUser, s.reverseTunnelSSHPort, privateKey, tunnelConnectionManager.tunnelRedeployNeeded)
153163
if err != nil {
154164
s.logger.WithError(err).Fatalf("failed to setup lambda execution manager")

0 commit comments

Comments
 (0)