Skip to content

Commit

Permalink
feat: added ability to provide assume role env
Browse files Browse the repository at this point in the history
Signed-off-by: Jiri Sveceny <jiri.sveceny@icloud.com>
  • Loading branch information
sitole committed Jul 19, 2024
1 parent 0b3e73b commit dba68f9
Showing 5 changed files with 135 additions and 26 deletions.
6 changes: 2 additions & 4 deletions docker-compose.yaml
Original file line number Diff line number Diff line change
@@ -1,14 +1,12 @@
version: '3.8'

services:
service-discovery:
build:
dockerfile: ./Dockerfile
target: runtime
ports:
- "3000:3000"
- "8080:8080"
environment:
PORT: 3000
PORT: 8080

AWS_ACCESS_KEY_ID: xxx
AWS_SECRET_ACCESS_KEY: xxx
9 changes: 8 additions & 1 deletion internal/service-discovery/server.go
Original file line number Diff line number Diff line change
@@ -26,7 +26,14 @@ func (s *ServiceDiscovery) Serve(_ context.Context, errors chan error) {
http.HandleFunc("/.health", s.serverGetHealth)

log.Println("RDS Instances service discovery running on ", s.serverAddress)
log.Println("Flowing cluster ready for scraping:", s.scraper.Targets)
log.Println("Flowing cluster ready for scraping")
for _, target := range s.scraper.Targets {
if target.ClusterAssumeRoleArn == nil {
log.Println("Cluster: ", target.ClusterArn, " and number: ", target.ClusterNumber)
} else {
log.Println("Cluster: ", target.ClusterArn, " with role: ", *target.ClusterAssumeRoleArn, " and number: ", target.ClusterNumber)
}
}

errors <- http.ListenAndServe(s.serverAddress, nil)
}
74 changes: 55 additions & 19 deletions internal/service-scraper/scraper.go
Original file line number Diff line number Diff line change
@@ -5,14 +5,18 @@ import (
"fmt"
"github.com/aws/aws-sdk-go-v2/aws"
awsConfig "github.com/aws/aws-sdk-go-v2/config"
"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
"github.com/aws/aws-sdk-go-v2/service/rds"
"github.com/aws/aws-sdk-go-v2/service/rds/types"
"github.com/aws/aws-sdk-go-v2/service/sts"
"github.com/aws/aws-sdk-go/aws/arn"
log "github.com/sirupsen/logrus"
)

type ScrapingTarget struct {
ClusterArn string
ClusterNumber int
ClusterArn string
ClusterAssumeRoleArn *string
ClusterNumber int
}

type Scraper struct {
@@ -42,27 +46,27 @@ func NewScraper(targets []ScrapingTarget) *Scraper {
}

func (s *Scraper) GetInstances(ctx context.Context) ([]ScraperDiscoveredTarget, error) {
awsRdsInstances := make([]ScraperDiscoveredTarget, 0)
awsSession, err := awsConfig.LoadDefaultConfig(ctx)
if err != nil {
return nil, err
}

awsRds := rds.NewFromConfig(awsSession)
instances := make([]ScraperDiscoveredTarget, 0)

// iterate over all registered clusters
for _, cluster := range s.Targets {
// filters cluster that contains cluster id that we want
filters := []types.Filter{
{
Name: aws.String("db-cluster-id"),
Values: []string{cluster.ClusterArn},
},
awsRdsClient, err := s.getClusterAwsClient(ctx, cluster)
if err != nil {
log.Errorf("Error creating AWS client for cluster %s %v", cluster.ClusterArn, err)
continue
}

resp, err := awsRds.DescribeDBInstances(ctx, &rds.DescribeDBInstancesInput{Filters: filters})
// describe all instances in the cluster with expected cluster ARN
resp, err := awsRdsClient.DescribeDBInstances(
ctx,
&rds.DescribeDBInstancesInput{
Filters: []types.Filter{
{Name: aws.String("db-cluster-id"), Values: []string{cluster.ClusterArn}},
},
},
)
if err != nil {
log.Errorf("Error describing RDS cluster <%s> %v", cluster.ClusterArn, err)
log.Errorf("Error describing RDS cluster %s %v", cluster.ClusterArn, err)
continue
}

@@ -79,9 +83,41 @@ func (s *Scraper) GetInstances(ctx context.Context) ([]ScraperDiscoveredTarget,
},
}

awsRdsInstances = append(awsRdsInstances, instanceTarget)
instances = append(instances, instanceTarget)
}
}

return awsRdsInstances, nil
return instances, nil
}

func (s *Scraper) getClusterAwsClient(ctx context.Context, target ScrapingTarget) (*rds.Client, error) {
conf, err := awsConfig.LoadDefaultConfig(ctx)
if err != nil {
return nil, err
}

// basic AWS client without assuming role
if target.ClusterAssumeRoleArn == nil {
return rds.NewFromConfig(conf), nil
}

targetArn, err := arn.Parse(target.ClusterArn)
if err != nil {
return nil, fmt.Errorf("invalid RDS cluster ARN: %s", target.ClusterArn)
}

// create a new AWS client with assumed role
stsClient := sts.NewFromConfig(conf)
assumedRoleCred := stscreds.NewAssumeRoleProvider(stsClient, *target.ClusterAssumeRoleArn)

assumedRoleConfig := conf
assumedRoleConfig.Credentials = aws.NewCredentialsCache(assumedRoleCred)

// this must be here because in AWS RDS API there is error,
// you must explicitly set region where cluster is located when assuming role,
// or you will receive weird "The parameter Filter: db-cluster-id is not a valid identifier." error message
// idk why, but it is what it is
assumedRoleConfig.Region = targetArn.Region

return rds.NewFromConfig(assumedRoleConfig), nil
}
16 changes: 14 additions & 2 deletions internal/service-scraper/scraper_targets_parser.go
Original file line number Diff line number Diff line change
@@ -25,9 +25,21 @@ func scraperTargetsEnvParser() ([]ScrapingTarget, error) {
return nil, fmt.Errorf("invalid key number: %s", keyNumberStr)
}

var assumeRoleArn *string

// try to find "SCRAPER_X_ASSUME_ROLE_ARN" environment variable
assumeRoleKey := fmt.Sprintf("SCRAPER_%d_ASSUME_ROLE_ARN", keyNumber)
assumeRoleKeyVal := os.Getenv(assumeRoleKey)
if assumeRoleKeyVal == "" {
assumeRoleArn = nil
} else {
assumeRoleArn = &assumeRoleKeyVal
}

target := ScrapingTarget{
ClusterArn: value,
ClusterNumber: keyNumber,
ClusterArn: value,
ClusterAssumeRoleArn: assumeRoleArn,
ClusterNumber: keyNumber,
}

scrapingTargets = append(scrapingTargets, target)
56 changes: 56 additions & 0 deletions internal/service-scraper/scraper_targets_parser_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
package service_scraper

import (
"os"
"testing"
)

func TestScrapeWithoutAssumeRoles(t *testing.T) {
arnA := "arn:aws:rds:eu-central-1:123456789012:cluster:first-cluster"
arnB := "arn:aws:rds:eu-central-1:123456789012:cluster:second-cluster"

_ = os.Setenv("SCRAPER_0_CLUSTER_ARN", arnA)
_ = os.Setenv("SCRAPER_1_CLUSTER_ARN", arnB)

targets, err := scraperTargetsEnvParser()
if err != nil {
t.Fatalf("Error parsing targets: %v", err)
}

if len(targets) != 2 {
t.Fatalf("Expected 2 targets, got %d", len(targets))
}

if targets[0].ClusterArn != arnA {
t.Fatalf("Expected first target ARN to be %s, got %s", arnA, targets[0].ClusterArn)
}

if targets[1].ClusterArn != arnB {
t.Fatalf("Expected first target ARN to be %s, got %s", arnB, targets[1].ClusterArn)
}
}

func TestScrapeWithAssumeRoles(t *testing.T) {
arn := "arn:aws:rds:eu-central-1:123456789012:cluster:first-cluster"
arnAssume := "arn:aws:iam::123456789013:role/some-fancy-role"

_ = os.Setenv("SCRAPER_0_CLUSTER_ARN", arn)
_ = os.Setenv("SCRAPER_0_ASSUME_ROLE_ARN", arnAssume)

targets, err := scraperTargetsEnvParser()
if err != nil {
t.Fatalf("Error parsing targets: %v", err)
}

if len(targets) != 1 {
t.Fatalf("Expected 1 targets, got %d", len(targets))
}

if targets[0].ClusterArn != arn {
t.Fatalf("Expected first target ARN to be %s, got %s", arn, targets[0].ClusterArn)
}

if targets[0].ClusterAssumeRoleArn == nil || *targets[0].ClusterAssumeRoleArn != arnAssume {
t.Fatalf("Expected first target ARN to be %s", arnAssume)
}
}

0 comments on commit dba68f9

Please sign in to comment.