<commit_before>package proxy

import (
	"context"
	"fmt"
	"regexp"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/aws/aws-sdk-go/service/sts/stsiface"
	"github.com/pkg/errors"
)

const (
	maxSessionNameLen int = 32
)

var (
	// matches char that is not valid in a STS role session name
	invalidSessionNameRegexp = regexp.MustCompile(`[^\w+=,.@-]`)

	sessionExpiration = 5 * time.Minute
)

type credentials struct {
	AccessKey   string
	Expiration  time.Time
	GeneratedAt time.Time
	RoleArn     roleArn
	SecretKey   string
	Token       string
}

func (c credentials) ExpiredNow() bool {
	return c.ExpiredAt(time.Now())
}

func (c credentials) ExpiredAt(at time.Time) bool {
	return at.After(c.Expiration)
}

func (c credentials) ExpiresIn(d time.Duration) bool {
	return c.ExpiredAt(time.Now().Add(d))
}

type containerCredentials struct {
	containerInfo
	credentials
}

func (c containerCredentials) IsValid(container containerInfo) bool {
	return c.containerInfo.IamRole.Equals(container.IamRole) &&
		c.containerInfo.ID == container.ID &&
		!c.credentials.ExpiresIn(sessionExpiration)
}

type credentialsProvider struct {
	container            containerService
	awsSts               stsiface.STSAPI
	defaultIamRoleArn    roleArn
	defaultIamPolicy     string
	containerCredentials map[string]containerCredentials
	lock                 sync.Mutex
}

func newCredentialsProvider(stsSvc stsiface.STSAPI, container containerService, defaultIamRoleArn roleArn, defaultIamPolicy string) *credentialsProvider {
	return &credentialsProvider{
		container:            container,
		awsSts:               stsSvc,
		defaultIamRoleArn:    defaultIamRoleArn,
		defaultIamPolicy:     defaultIamPolicy,
		containerCredentials: make(map[string]containerCredentials),
	}
}

// CredentialsForIP resolves the IP to a specific container, then attempts to assume the role
// specified in the container's metadata. Role specific credentials are returned.
//
// If the cache contains no fresh and valid role credentials, a fresh set is requested from
// AWS and cached.
func (c *credentialsProvider) CredentialsForIP(ctx context.Context, containerIP string) (credentials, error) {
	c.lock.Lock()
	defer c.lock.Unlock()

	container, err := c.container.ContainerForIP(ctx, containerIP)
	if err != nil {
		return credentials{}, errors.Wrapf(err, "Error finding container with IP [%s]", containerIP)
	}

	oldCredentials, found := c.containerCredentials[containerIP]

	if !found || !oldCredentials.IsValid(container) {
		arn := container.IamRole
		iamPolicy := container.IamPolicy

		if arn.Empty() {
			arn = c.defaultIamRoleArn

			if len(iamPolicy) == 0 {
				iamPolicy = c.defaultIamPolicy
			}
		}

		role, err := c.AssumeRole(arn, iamPolicy, generateSessionName(c.container.TypeName(), container.ID))

		if err != nil {
			return credentials{}, errors.Wrapf(err, "Error assuming role [%s] for container [%s] at IP {%s]", arn, container.Name, containerIP)
		}

		oldCredentials = containerCredentials{container, role}
		c.containerCredentials[containerIP] = oldCredentials
	}

	return oldCredentials.credentials, nil
}

func (c *credentialsProvider) AssumeRole(role roleArn, iamPolicy, sessionName string) (credentials, error) {
	var policy *string

	if len(iamPolicy) > 0 {
		policy = aws.String(iamPolicy)
	}

	resp, err := c.awsSts.AssumeRole(&sts.AssumeRoleInput{
		DurationSeconds: aws.Int64(3600), // Max is 1 hour
		Policy:          policy,
		RoleArn:         aws.String(role.String()),
		RoleSessionName: aws.String(sessionName),
	})

	if err != nil {
		return credentials{}, errors.Wrapf(err, "Error assuming role [%s] with policy [%s] and session name [%s]]", role, iamPolicy, sessionName)<commit_msg>fix: default policy only applied w/ default role<commit_after>package proxy

import (
	"context"
	"fmt"
	"regexp"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/aws/aws-sdk-go/service/sts/stsiface"
	"github.com/pkg/errors"
)

const (
	maxSessionNameLen int = 32
)

var (
	// matches char that is not valid in a STS role session name
	invalidSessionNameRegexp = regexp.MustCompile(`[^\w+=,.@-]`)

	sessionExpiration = 5 * time.Minute
)

type credentials struct {
	AccessKey   string
	Expiration  time.Time
	GeneratedAt time.Time
	RoleArn     roleArn
	SecretKey   string
	Token       string
}

func (c credentials) ExpiredNow() bool {
	return c.ExpiredAt(time.Now())
}

func (c credentials) ExpiredAt(at time.Time) bool {
	return at.After(c.Expiration)
}

func (c credentials) ExpiresIn(d time.Duration) bool {
	return c.ExpiredAt(time.Now().Add(d))
}

type containerCredentials struct {
	containerInfo
	credentials
}

func (c containerCredentials) IsValid(container containerInfo) bool {
	return c.containerInfo.IamRole.Equals(container.IamRole) &&
		c.containerInfo.ID == container.ID &&
		!c.credentials.ExpiresIn(sessionExpiration)
}

type credentialsProvider struct {
	container            containerService
	awsSts               stsiface.STSAPI
	defaultIamRoleArn    roleArn
	defaultIamPolicy     string
	containerCredentials map[string]containerCredentials
	lock                 sync.Mutex
}

func newCredentialsProvider(stsSvc stsiface.STSAPI, container containerService, defaultIamRoleArn roleArn, defaultIamPolicy string) *credentialsProvider {
	return &credentialsProvider{
		container:            container,
		awsSts:               stsSvc,
		defaultIamRoleArn:    defaultIamRoleArn,
		defaultIamPolicy:     defaultIamPolicy,
		containerCredentials: make(map[string]containerCredentials),
	}
}

// CredentialsForIP resolves the IP to a specific container, then attempts to assume the role
// specified in the container's metadata. Role specific credentials are returned.
//
// If the cache contains no fresh and valid role credentials, a fresh set is requested from
// AWS and cached.
func (c *credentialsProvider) CredentialsForIP(ctx context.Context, containerIP string) (credentials, error) {
	c.lock.Lock()
	defer c.lock.Unlock()

	container, err := c.container.ContainerForIP(ctx, containerIP)
	if err != nil {
		return credentials{}, errors.Wrapf(err, "Error finding container with IP [%s]", containerIP)
	}

	oldCredentials, found := c.containerCredentials[containerIP]

	if !found || !oldCredentials.IsValid(container) {
		arn := container.IamRole
		iamPolicy := container.IamPolicy

		if arn.Empty() {
			arn = c.defaultIamRoleArn
		}

		if len(iamPolicy) == 0 {
			iamPolicy = c.defaultIamPolicy
		}

		role, err := c.AssumeRole(arn, iamPolicy, generateSessionName(c.container.TypeName(), container.ID))

		if err != nil {
			return credentials{}, errors.Wrapf(err, "Error assuming role [%s] for container [%s] at IP {%s]", arn, container.Name, containerIP)
		}

		oldCredentials = containerCredentials{container, role}
		c.containerCredentials[containerIP] = oldCredentials
	}

	return oldCredentials.credentials, nil
}

func (c *credentialsProvider) AssumeRole(role roleArn, iamPolicy, sessionName string) (credentials, error) {
	var policy *string

	if len(iamPolicy) > 0 {
		policy = aws.String(iamPolicy)
	}

	resp, err := c.awsSts.AssumeRole(&sts.AssumeRoleInput{
		DurationSeconds: aws.Int64(3600), // Max is 1 hour
		Policy:          policy,
		RoleArn:         aws.String(role.String()),
		RoleSessionName: aws.String(sessionName),
	})

	if err != nil {
		return credentials{}, errors.Wrapf(err, "Error assuming role [%s] with policy [%s] and session name [%s]]", role, iamPolicy, sessionName)
{'fraction_non_alphanumeric': 0.1077127659574468, 'fraction_numerical': 0.0031914893617021275, 'mean_word_length': 4.186896551724138, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2601179', 'n_tokens_mistral': 2491, 'n_tokens_neox': 2376, 'n_words': 610}