Exemplo n.º 1
0
// Retrieve generates a new set of temporary credentials using STS.
func (p *AssumeRoleProvider) Retrieve() (credentials.Value, error) {

	// Apply defaults where parameters are not set.
	if p.Client == nil {
		p.Client = sts.New(nil)
	}
	if p.RoleSessionName == "" {
		// Try to work out a role name that will hopefully end up unique.
		p.RoleSessionName = fmt.Sprintf("%d", time.Now().UTC().UnixNano())
	}
	if p.Duration == 0 {
		// Expire as often as AWS permits.
		p.Duration = 15 * time.Minute
	}

	roleOutput, err := p.Client.AssumeRole(&sts.AssumeRoleInput{
		DurationSeconds: aws.Int64(int64(p.Duration / time.Second)),
		RoleArn:         aws.String(p.RoleARN),
		RoleSessionName: aws.String(p.RoleSessionName),
		ExternalId:      p.ExternalID,
	})

	if err != nil {
		return credentials.Value{}, err
	}

	// We will proactively generate new credentials before they expire.
	p.SetExpiration(*roleOutput.Credentials.Expiration, p.ExpiryWindow)

	return credentials.Value{
		AccessKeyID:     *roleOutput.Credentials.AccessKeyId,
		SecretAccessKey: *roleOutput.Credentials.SecretAccessKey,
		SessionToken:    *roleOutput.Credentials.SessionToken,
	}, nil
}
Exemplo n.º 2
0
func TestPutBuffers(t *testing.T) {
	testCases := [][]message{
		{{aws.String("twinkle"), []byte("twinkle")}},
		{{aws.String("hey"), []byte("there")}, {aws.String("big"), []byte("fella")}},
	}

	for _, testCase := range testCases {
		producer := producerWithStubClient(len(testCase) + 1)
		client := producer.client.(*StubClient)

		for _, input := range testCase {
			err := producer.Put(input.PartitionKey, input.Value)
			if err != nil {
				t.Fatalf("expected no Put errors. got '%s'", err)
			}
		}

		if client.puts != 0 {
			t.Errorf("expected 0 sent messages, got %d", client.puts)
		}
		if producer.current != len(testCase) {
			t.Errorf("expected %d buffered messages, found %d", len(testCase), len(producer.messages))
		}
	}
}
Exemplo n.º 3
0
func TestPutRetriesFailedRecords(t *testing.T) {
	testCases := []struct {
		name      string
		messages  []message
		responses []clientResponse
	}{
		{
			"one message retrying once",
			[]message{{aws.String("twinkle"), []byte("twinkle")}},
			[]clientResponse{
				{outputWithErrors("ProvisionedThroughputExceededException"), nil},
			},
		},
		{
			"one message retrying twice",
			[]message{{aws.String("twinkle"), []byte("twinkle")}},
			[]clientResponse{
				{outputWithErrors("ProvisionedThroughputExceededException"), nil},
				{outputWithErrors("ProvisionedThroughputExceededException"), nil},
			},
		},
		{
			"two messages with one retrying once",
			[]message{{aws.String("hey"), []byte("there")}, {aws.String("big"), []byte("fella")}},
			[]clientResponse{
				{outputWithErrors("", "ProvisionedThroughputExceededException"), nil},
			},
		},
	}

	for _, testCase := range testCases {
		actualRetries := 0

		producer := producerRespondingWith(MaxSendSize, testCase.responses...)
		producer.Throttle = func() Throttle {
			actualRetries++
			return &noOpThrottle{}
		}

		client := producer.client.(*StubClient)

		for _, m := range testCase.messages {
			if err := producer.Put(m.PartitionKey, m.Value); err != nil {
				t.Fatalf("unexpected producer error! %s", err)
			}
		}
		producer.Flush()

		expectedRetries := len(testCase.responses)
		if actualRetries != expectedRetries {
			t.Errorf("expected %d retries, got %d", expectedRetries, actualRetries)
		}
		assertSentMessages(t, testCase.name, testCase.messages, client.sent)
	}
}
Exemplo n.º 4
0
func shardsToAws(shards ...shard) []*kinesis.Shard {
	var awsShards []*kinesis.Shard
	for _, s := range shards {
		awsShard := &kinesis.Shard{ShardId: aws.String(s.id)}
		if s.parentOne != "" {
			awsShard.ParentShardId = aws.String(s.parentOne)
		}
		if s.parentTwo != "" {
			awsShard.AdjacentParentShardId = aws.String(s.parentTwo)
		}
		awsShards = append(awsShards, awsShard)
	}
	return awsShards
}
Exemplo n.º 5
0
// A response where every record in the given input is a success.
func successfulResponse(input *kinesis.PutRecordsInput) (*kinesis.PutRecordsOutput, error) {
	records := make([]*kinesis.PutRecordsResultEntry, len(input.Records))
	for i := range input.Records {
		records[i] = &kinesis.PutRecordsResultEntry{
			SequenceNumber: aws.String("sequence_number"),
			ShardId:        aws.String("an_shard"),
		}
	}
	output := &kinesis.PutRecordsOutput{
		FailedRecordCount: aws.Int64(0),
		Records:           records,
	}
	return output, nil
}
Exemplo n.º 6
0
func (p *Producer) send() error {
	if len(p.messages) == 0 {
		return nil
	}

	defer p.reset()

	stream, messages := aws.String(p.StreamName), p.messages[0:p.current]
	for {
		res, err := p.client.PutRecords(putRecordsInput(stream, messages))

		if err != nil {
			return err
		}

		if *res.FailedRecordCount == 0 {
			if p.Debug {
				log.Printf("Put %d message(s).", len(res.Records))
			}
			return nil
		}

		messages = failedMessages(messages, res.Records)
		if p.Debug {
			log.Printf("Put failed for %d message(s). Backing off and trying again.", *res.FailedRecordCount)
		}
		p.Throttle().Await()
	}
	return nil
}
Exemplo n.º 7
0
func consumerWith(descriptions [][]shard, data map[string][]string, processor Processor) *Consumer {
	return &Consumer{
		stream:     aws.String(defaultStream),
		client:     &StubClient{describe: descriptions, records: data},
		complete:   make(chan string),
		processor:  processor,
		waiterFunc: func() waiter { return &stubWaiter{} },
	}
}
Exemplo n.º 8
0
func TestInvalidPut(t *testing.T) {
	invalid := []struct {
		name     string
		key      *string
		value    []byte
		expected error
	}{
		{"key too long", aws.String(randomString(257)), []byte{0x6c, 0x6f, 0x6c}, PartitionKeyTooLong},
		{"invalid unicode", aws.String(string([]byte{0xc1, 0xbf})), []byte{0x6c, 0x6f, 0x6c}, InvalidUnicode},
		{"empty value", aws.String(randomString(123)), []byte{}, EmptyValue},
	}

	producer := producerWithStubClient(123)
	for _, tc := range invalid {
		if err := producer.Put(tc.key, tc.value); !errContains(err, tc.expected) {
			t.Errorf("expected %s to error with '%s'. got '%+v'", tc.name, tc.expected, err)
		}
	}
}
Exemplo n.º 9
0
// Start a consumer at the given Stream's LATEST and process each shard with
// processor.
//
// Each shard will be processed in an individual goroutine.
func Tail(stream string, debug bool, processor Processor) error {
	c := &Consumer{
		stream:    aws.String(stream),
		client:    kinesis.New(nil),
		processor: processor,

		debug: debug,

		complete:   make(chan string),
		waiterFunc: func() waiter { return &realWaiter{} },
	}

	return c.tail()
}
Exemplo n.º 10
0
func TestPutSends(t *testing.T) {
	testCases := [][]message{
		{{aws.String("twinkle"), []byte("twinkle")}},
		{{aws.String("hey"), []byte("there")}, {aws.String("big"), []byte("fella")}},
	}

	for _, testCase := range testCases {
		producer := producerWithStubClient(len(testCase))
		client := producer.client.(*StubClient)

		for _, input := range testCase {
			err := producer.Put(input.PartitionKey, input.Value)
			if err != nil {
				t.Fatalf("expected no Put errors. got '%s'", err)
			}
		}

		if client.puts != 1 {
			t.Errorf("expected a single send to Kinesis")
			return
		}
		assertSentMessages(t, fmt.Sprintf("%d successful records", len(testCase)), testCase, client.sent)
	}
}
Exemplo n.º 11
0
func (c *Consumer) startShardConsumer(shard string, iterType *string, processor Processor) {
	s := &shardConsumer{
		client:    c.client,
		stream:    c.stream,
		shard:     aws.String(shard),
		debug:     c.debug,
		processor: processor,

		waiter:   c.waiterFunc(),
		complete: c.complete,
	}

	go func() {
		s.init(iterType)
		s.consume()
	}()
}
Exemplo n.º 12
0
// Return an output response with the given error codes. Blank strings imply
// that a message was not an error.
func outputWithErrors(codes ...string) *kinesis.PutRecordsOutput {
	resultEntries := make([]*kinesis.PutRecordsResultEntry, len(codes))

	errorCount := 0
	for i, code := range codes {
		resultEntries[i] = &kinesis.PutRecordsResultEntry{
			ErrorCode: aws.String(code),
		}

		if code != "" {
			errorCount++
		}
	}

	return &kinesis.PutRecordsOutput{
		FailedRecordCount: aws.Int64(int64(errorCount)),
		Records:           resultEntries,
	}
}
Exemplo n.º 13
0
// Send the given string to Kinesis. The first 256 bytes of the string will be
// used as the partition key. message must not be a valid Unicode string, and
// must be non-empty.
func (p *Producer) PutString(message string) error {
	return p.Put(aws.String(message[:intMin(len(message), 256)]), []byte(message))
}
Exemplo n.º 14
0
func (c *Consumer) tail() error {
	shards, err := c.listShards()
	if err != nil {
		return err
	}

	go c.monitor()

	for _, id := range withNoChildren(shards) {
		c.startShardConsumer(id, LATEST, c.processor)
	}

	return nil
}

var LATEST = aws.String(kinesis.ShardIteratorTypeLatest)
var TRIM_HORIZON = aws.String(kinesis.ShardIteratorTypeTrimHorizon)

func (c *Consumer) startShardConsumer(shard string, iterType *string, processor Processor) {
	s := &shardConsumer{
		client:    c.client,
		stream:    c.stream,
		shard:     aws.String(shard),
		debug:     c.debug,
		processor: processor,

		waiter:   c.waiterFunc(),
		complete: c.complete,
	}

	go func() {