Skip to content

refactor(go-forwarder): add client package #1134

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ndakkoune merged 1 commit into from
May 22, 2026
Merged

refactor(go-forwarder): add client package #1134

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 72 additions & 0 deletions aws/logs_monitoring_go/internal/client/kms.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2026-Present Datadog, Inc.

package client

//go:generate go tool mockgen -source=kms.go -package=client -destination=kms_mockgen.go

import (
"context"
"encoding/base64"
"fmt"
"log/slog"
"strings"

"github.com/aws/aws-sdk-go-v2/aws"
awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http"
awsconfig "github.com/aws/aws-sdk-go-v2/config"
"github.com/aws/aws-sdk-go-v2/service/kms"
)

type KMS interface {
Decrypt(ctx context.Context, params *kms.DecryptInput, optFns ...func(*kms.Options)) (*kms.DecryptOutput, error)
}

func NewKMS(ctx context.Context, useFIPS bool) (KMS, error) {
cfg, err := awsconfig.LoadDefaultConfig(ctx, awsconfig.WithHTTPClient(awshttp.NewBuildableClient().WithTimeout(timeout)))
if err != nil {
return nil, err
}

resolver := kms.NewDefaultEndpointResolverV2()
params := kms.EndpointParameters{
Region: aws.String(cfg.Region),
UseFIPS: aws.Bool(useFIPS),
}

endpoint, err := resolver.ResolveEndpoint(ctx, params)
if err != nil && useFIPS {
slog.Warn("FIPS endpoint not available, falling back to standard endpoint", slog.String("service", "kms"), slog.String("region", cfg.Region))
params.UseFIPS = aws.Bool(false)
endpoint, err = resolver.ResolveEndpoint(ctx, params)
}
if err != nil {
return nil, fmt.Errorf("resolve endpoint: %w", err)
}

return kms.NewFromConfig(cfg, func(o *kms.Options) {
o.BaseEndpoint = aws.String(endpoint.URI.String())
}), nil
}

func DecryptKMSCiphertext(ctx context.Context, kmsClient KMS, ciphertext string) (string, error) {
decoded, err := base64.StdEncoding.DecodeString(ciphertext)
if err != nil {
return "", fmt.Errorf("base64-decoding ciphertext: %w", err)
}

result, err := kmsClient.Decrypt(ctx, &kms.DecryptInput{
CiphertextBlob: decoded,
})
if err != nil {
return "", fmt.Errorf("decrypting KMS ciphertext: %w", err)
}

if result.Plaintext == nil {
return "", fmt.Errorf("KMS decryption returned no plaintext")
}

return strings.TrimSpace(string(result.Plaintext)), nil
}
62 changes: 62 additions & 0 deletions aws/logs_monitoring_go/internal/client/kms_mockgen.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

20 changes: 10 additions & 10 deletions ...monitoring_go/internal/config/kms_test.go → ...monitoring_go/internal/client/kms_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2026-Present Datadog, Inc.

package config
package client

import (
"encoding/base64"
Expand All @@ -16,17 +16,17 @@ import (
"go.uber.org/mock/gomock"
)

func TestResolveFromKMS(t *testing.T) {
func TestDecryptKMSCiphertext(t *testing.T) {
validCiphertext := base64.StdEncoding.EncodeToString([]byte("encrypted-blob"))

tests := map[string]struct {
mockSetup func(m *MockKMSAPIClient)
mockSetup func(m *MockKMS)
ciphertext string
wantKey string
wantErr bool
}{
"success": {
mockSetup: func(m *MockKMSAPIClient) {
mockSetup: func(m *MockKMS) {
m.EXPECT().
Decrypt(gomock.Any(), gomock.Any()).
Return(&kms.DecryptOutput{
Expand All @@ -37,7 +37,7 @@ func TestResolveFromKMS(t *testing.T) {
wantKey: "abcdef1234567890abcdef1234567890",
},
"whitespace_trimmed": {
mockSetup: func(m *MockKMSAPIClient) {
mockSetup: func(m *MockKMS) {
m.EXPECT().
Decrypt(gomock.Any(), gomock.Any()).
Return(&kms.DecryptOutput{
Expand All @@ -48,12 +48,12 @@ func TestResolveFromKMS(t *testing.T) {
wantKey: "abcdef1234567890abcdef1234567890",
},
"invalid_base64": {
mockSetup: func(m *MockKMSAPIClient) {},
mockSetup: func(m *MockKMS) {},
ciphertext: "not-valid-base64!@#$",
wantErr: true,
},
"aws_error": {
mockSetup: func(m *MockKMSAPIClient) {
mockSetup: func(m *MockKMS) {
m.EXPECT().
Decrypt(gomock.Any(), gomock.Any()).
Return(nil, errors.New("InvalidCiphertextException: invalid ciphertext"))
Expand All @@ -62,7 +62,7 @@ func TestResolveFromKMS(t *testing.T) {
wantErr: true,
},
"nil_plaintext": {
mockSetup: func(m *MockKMSAPIClient) {
mockSetup: func(m *MockKMS) {
m.EXPECT().
Decrypt(gomock.Any(), gomock.Any()).
Return(&kms.DecryptOutput{
Expand All @@ -77,10 +77,10 @@ func TestResolveFromKMS(t *testing.T) {
for name, tc := range tests {
t.Run(name, func(t *testing.T) {
ctrl := gomock.NewController(t)
mock := NewMockKMSAPIClient(ctrl)
mock := NewMockKMS(ctrl)
tc.mockSetup(mock)

got, err := decryptKMSCiphertext(t.Context(), mock, tc.ciphertext)
got, err := DecryptKMSCiphertext(t.Context(), mock, tc.ciphertext)
if tc.wantErr {
require.Error(t, err)
return
Expand Down
19 changes: 11 additions & 8 deletions ...itoring_go/internal/handling/s3_client.go → aws/logs_monitoring_go/internal/client/s3.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2026-Present Datadog, Inc.

package handling
package client

//go:generate go tool mockgen -source=s3_client.go -package=handling -destination=s3_client_mockgen.go
//go:generate go tool mockgen -source=s3.go -package=client -destination=s3_mockgen.go

import (
"context"
Expand All @@ -25,22 +25,25 @@ const timeout = 10 * time.Second

var (
s3ClientOnce sync.Once
s3Client S3APIClient
s3Client S3
s3ClientErr error
)

type S3APIClient interface {
type S3 interface {
GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error)
PutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error)
DeleteObject(ctx context.Context, params *s3.DeleteObjectInput, optFns ...func(*s3.Options)) (*s3.DeleteObjectOutput, error)
ListObjectsV2(ctx context.Context, params *s3.ListObjectsV2Input, optFns ...func(*s3.Options)) (*s3.ListObjectsV2Output, error)
}

func getS3APIClient(ctx context.Context, useFIPS bool) (S3APIClient, error) {
func GetS3(ctx context.Context, useFIPS bool) (S3, error) {
s3ClientOnce.Do(func() {
s3Client, s3ClientErr = createS3APIClient(ctx, useFIPS)
s3Client, s3ClientErr = newS3(ctx, useFIPS)
})
return s3Client, s3ClientErr
}

func createS3APIClient(ctx context.Context, useFIPS bool) (S3APIClient, error) {
func newS3(ctx context.Context, useFIPS bool) (S3, error) {
cfg, err := awsconfig.LoadDefaultConfig(ctx, awsconfig.WithHTTPClient(awshttp.NewBuildableClient().WithTimeout(timeout)))
if err != nil {
return nil, err
Expand All @@ -67,7 +70,7 @@ func createS3APIClient(ctx context.Context, useFIPS bool) (S3APIClient, error) {
}), nil
}

func getS3Object(ctx context.Context, client S3APIClient, bucket, key string) (io.ReadCloser, error) {
func GetS3Object(ctx context.Context, client S3, bucket, key string) (io.ReadCloser, error) {
result, err := client.GetObject(ctx, &s3.GetObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(key),
Expand Down
122 changes: 122 additions & 0 deletions aws/logs_monitoring_go/internal/client/s3_mockgen.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading