feat: support sigv4 signing (#169)

* feat: support sigv4 signing

---------

Signed-off-by: obanby <[email protected]>

Author: obanby

Diff for: pkg/remote/client.go

@@ -12,6 +12,8 @@ import (
 	"net/url"
 	"time"
 
+	"github.com/grafana/xk6-output-prometheus-remote/pkg/sigv4"
+
 	prompb "buf.build/gen/go/prometheus/prometheus/protocolbuffers/go"
 	"github.com/klauspost/compress/snappy"
 	"google.golang.org/protobuf/proto"
@@ -22,6 +24,7 @@ type HTTPConfig struct {
 	Timeout   time.Duration
 	TLSConfig *tls.Config
 	BasicAuth *BasicAuth
+	SigV4     *sigv4.Config
 	Headers   http.Header
 }
 
@@ -60,6 +63,13 @@ func NewWriteClient(endpoint string, cfg *HTTPConfig) (*WriteClient, error) {
 			TLSClientConfig: cfg.TLSConfig,
 		}
 	}
+	if cfg.SigV4 != nil {
+		tripper, err := sigv4.NewRoundTripper(cfg.SigV4, wc.hc.Transport)
+		if err != nil {
+			return nil, err
+		}
+		wc.hc.Transport = tripper
+	}
 	return wc, nil
 }
 

Diff for: pkg/remotewrite/config.go

@@ -3,12 +3,15 @@ package remotewrite
 import (
 	"crypto/tls"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net/http"
 	"strconv"
 	"strings"
 	"time"
 
+	"github.com/grafana/xk6-output-prometheus-remote/pkg/sigv4"
+
 	"github.com/grafana/xk6-output-prometheus-remote/pkg/remote"
 	"go.k6.io/k6/lib/types"
 	"gopkg.in/guregu/null.v3"
@@ -68,6 +71,15 @@ type Config struct {
 	TrendStats []string `json:"trendStats"`
 
 	StaleMarkers null.Bool `json:"staleMarkers"`
+
+	// SigV4Region is the AWS region where the workspace is.
+	SigV4Region null.String `json:"sigV4Region"`
+
+	// SigV4AccessKey is the AWS access key.
+	SigV4AccessKey null.String `json:"sigV4AccessKey"`
+
+	// SigV4SecretKey is the AWS secret key.
+	SigV4SecretKey null.String `json:"sigV4SecretKey"`
 }
 
 // NewConfig creates an Output's configuration.
@@ -81,6 +93,9 @@ func NewConfig() Config {
 		Headers:               make(map[string]string),
 		TrendStats:            defaultTrendStats,
 		StaleMarkers:          null.BoolFrom(false),
+		SigV4Region:           null.NewString("", false),
+		SigV4AccessKey:        null.NewString("", false),
+		SigV4SecretKey:        null.NewString("", false),
 	}
 }
 
@@ -110,6 +125,22 @@ func (conf Config) RemoteConfig() (*remote.HTTPConfig, error) {
 		hc.TLSConfig.Certificates = []tls.Certificate{cert}
 	}
 
+	if isSigV4PartiallyConfigured(conf.SigV4Region, conf.SigV4AccessKey, conf.SigV4SecretKey) {
+		return nil, errors.New(
+			"sigv4 seems to be partially configured. All of " +
+				"K6_PROMETHEUS_RW_SIGV4_REGION, K6_PROMETHEUS_RW_SIGV4_ACCESS_KEY, K6_PROMETHEUS_RW_SIGV4_SECRET_KEY " +
+				"must all be set. Unset all to bypass sigv4",
+		)
+	}
+
+	if conf.SigV4Region.Valid && conf.SigV4AccessKey.Valid && conf.SigV4SecretKey.Valid {
+		hc.SigV4 = &sigv4.Config{
+			Region:             conf.SigV4Region.String,
+			AwsAccessKeyID:     conf.SigV4AccessKey.String,
+			AwsSecretAccessKey: conf.SigV4SecretKey.String,
+		}
+	}
+
 	if len(conf.Headers) > 0 {
 		hc.Headers = make(http.Header)
 		for k, v := range conf.Headers {
@@ -149,6 +180,18 @@ func (conf Config) Apply(applied Config) Config {
 		conf.BearerToken = applied.BearerToken
 	}
 
+	if applied.SigV4Region.Valid {
+		conf.SigV4Region = applied.SigV4Region
+	}
+
+	if applied.SigV4AccessKey.Valid {
+		conf.SigV4AccessKey = applied.SigV4AccessKey
+	}
+
+	if applied.SigV4SecretKey.Valid {
+		conf.SigV4SecretKey = applied.SigV4SecretKey
+	}
+
 	if applied.PushInterval.Valid {
 		conf.PushInterval = applied.PushInterval
 	}
@@ -299,6 +342,18 @@ func parseEnvs(env map[string]string) (Config, error) {
 		}
 	}
 
+	if sigV4Region, sigV4RegionDefined := env["K6_PROMETHEUS_RW_SIGV4_REGION"]; sigV4RegionDefined {
+		c.SigV4Region = null.StringFrom(sigV4Region)
+	}
+
+	if sigV4AccessKey, sigV4AccessKeyDefined := env["K6_PROMETHEUS_RW_SIGV4_ACCESS_KEY"]; sigV4AccessKeyDefined {
+		c.SigV4AccessKey = null.StringFrom(sigV4AccessKey)
+	}
+
+	if sigV4SecretKey, sigV4SecretKeyDefined := env["K6_PROMETHEUS_RW_SIGV4_SECRET_KEY"]; sigV4SecretKeyDefined {
+		c.SigV4SecretKey = null.StringFrom(sigV4SecretKey)
+	}
+
 	if b, err := envBool(env, "K6_PROMETHEUS_RW_TREND_AS_NATIVE_HISTOGRAM"); err != nil {
 		return c, err
 	} else if b.Valid {
@@ -384,3 +439,12 @@ func parseArg(text string) (Config, error) {
 
 	return c, nil
 }
+
+func isSigV4PartiallyConfigured(region, accessKey, secretKey null.String) bool {
+	hasRegion := region.Valid && len(strings.TrimSpace(region.String)) != 0
+	hasAccessID := accessKey.Valid && len(strings.TrimSpace(accessKey.String)) != 0
+	hasSecretAccessKey := secretKey.Valid && len(strings.TrimSpace(secretKey.String)) != 0
+	// either they are all set, or all not set. False if partial
+	isComplete := (hasRegion && hasAccessID && hasSecretAccessKey) || (!hasRegion && !hasAccessID && !hasSecretAccessKey)
+	return !isComplete
+}

Diff for: pkg/sigv4/const.go

@@ -0,0 +1,23 @@
+package sigv4
+
+const (
+	// Amazon Managed Service for Prometheus
+	awsServiceName = "aps"
+
+	signingAlgorithm = "AWS4-HMAC-SHA256"
+
+	authorizationHeaderKey = "Authorization"
+	amzDateKey             = "X-Amz-Date"
+
+	// emptyStringSHA256 is the hex encoded sha256 value of an empty string
+	emptyStringSHA256 = `e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855`
+
+	// timeFormat is the time format to be used in the X-Amz-Date header or query parameter
+	timeFormat = "20060102T150405Z"
+
+	// shortTimeFormat is the shorten time format used in the credential scope
+	shortTimeFormat = "20060102"
+
+	// contentSHAKey is the SHA256 of request body
+	contentSHAKey = "X-Amz-Content-Sha256"
+)