Add sigv4 module

sigv4/go.mod

@@ -0,0 +1,11 @@
+module github.com/prometheus/common/sigv4
+
+go 1.16
+
+require (
+	github.com/aws/aws-sdk-go v1.38.35
+	github.com/prometheus/client_golang v1.11.0
+	github.com/prometheus/common v0.29.0
+	github.com/stretchr/testify v1.7.0
+	gopkg.in/yaml.v2 v2.4.0
+)

sigv4/sigv4.go

@@ -0,0 +1,137 @@
+// Copyright 2021 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sigv4
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"net/textproto"
+	"sync"
+	"time"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
+	"github.com/aws/aws-sdk-go/aws/session"
+	signer "github.com/aws/aws-sdk-go/aws/signer/v4"
+)
+
+var sigv4HeaderDenylist = []string{
+	"uber-trace-id",
+}
+
+type sigV4RoundTripper struct {
+	region string
+	next   http.RoundTripper
+	pool   sync.Pool
+
+	signer *signer.Signer
+}
+
+// NewSigV4RoundTripper returns a new http.RoundTripper that will sign requests
+// using Amazon's Signature Verification V4 signing procedure. The request will
+// then be handed off to the next RoundTripper provided by next. If next is nil,
+// http.DefaultTransport will be used.
+//
+// Credentials for signing are retrieved using the the default AWS credential
+// chain. If credentials cannot be found, an error will be returned.
+func NewSigV4RoundTripper(cfg *SigV4Config, next http.RoundTripper) (http.RoundTripper, error) {
+	if next == nil {
+		next = http.DefaultTransport
+	}
+
+	creds := credentials.NewStaticCredentials(cfg.AccessKey, string(cfg.SecretKey), "")
+	if cfg.AccessKey == "" && cfg.SecretKey == "" {
+		creds = nil
+	}
+
+	sess, err := session.NewSessionWithOptions(session.Options{
+		Config: aws.Config{
+			Region:      aws.String(cfg.Region),
+			Credentials: creds,
+		},
+		Profile: cfg.Profile,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("could not create new AWS session: %w", err)
+	}
+	if _, err := sess.Config.Credentials.Get(); err != nil {
+		return nil, fmt.Errorf("could not get SigV4 credentials: %w", err)
+	}
+	if aws.StringValue(sess.Config.Region) == "" {
+		return nil, fmt.Errorf("region not configured in sigv4 or in default credentials chain")
+	}
+
+	signerCreds := sess.Config.Credentials
+	if cfg.RoleARN != "" {
+		signerCreds = stscreds.NewCredentials(sess, cfg.RoleARN)
+	}
+
+	rt := &sigV4RoundTripper{
+		region: cfg.Region,
+		next:   next,
+		signer: signer.NewSigner(signerCreds),
+	}
+	rt.pool.New = rt.newBuf
+	return rt, nil
+}
+
+func (rt *sigV4RoundTripper) newBuf() interface{} {
+	return bytes.NewBuffer(make([]byte, 0, 1024))
+}
+
+func (rt *sigV4RoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	// rt.signer.Sign needs a seekable body, so we replace the body with a
+	// buffered reader filled with the contents of original body.
+	buf := rt.pool.Get().(*bytes.Buffer)
+	defer func() {
+		buf.Reset()
+		rt.pool.Put(buf)
+	}()
+	if _, err := io.Copy(buf, req.Body); err != nil {
+		return nil, err
+	}
+	// Close the original body since we don't need it anymore.
+	_ = req.Body.Close()
+
+	// Ensure our seeker is back at the start of the buffer once we return.
+	var seeker io.ReadSeeker = bytes.NewReader(buf.Bytes())
+	defer func() {
+		_, _ = seeker.Seek(0, io.SeekStart)
+	}()
+	req.Body = ioutil.NopCloser(seeker)
+
+	// Clone the request and trim out headers that we don't want to sign.
+	signReq := req.Clone(req.Context())
+	for _, header := range sigv4HeaderDenylist {
+		signReq.Header.Del(header)
+	}
+
+	headers, err := rt.signer.Sign(signReq, seeker, "aps", rt.region, time.Now().UTC())
+	if err != nil {
+		return nil, fmt.Errorf("failed to sign request: %w", err)
+	}
+
+	// Copy over signed headers. Authorization header is not returned by
+	// rt.signer.Sign and needs to be copied separately.
+	for k, v := range headers {
+		req.Header[textproto.CanonicalMIMEHeaderKey(k)] = v
+	}
+	req.Header.Set("Authorization", signReq.Header.Get("Authorization"))
+
+	return rt.next.RoundTrip(req)
+}

sigv4/sigv4_config.go

@@ -0,0 +1,46 @@
+// Copyright 2021 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sigv4
+
+import (
+	"fmt"
+
+	"github.com/prometheus/common/config"
+)
+
+// SigV4Config is the configuration for signing remote write requests with
+// AWS's SigV4 verification process. Empty values will be retrieved using the
+// AWS default credentials chain.
+type SigV4Config struct {
+	Region    string        `yaml:"region,omitempty"`
+	AccessKey string        `yaml:"access_key,omitempty"`
+	SecretKey config.Secret `yaml:"secret_key,omitempty"`
+	Profile   string        `yaml:"profile,omitempty"`
+	RoleARN   string        `yaml:"role_arn,omitempty"`
+}
+
+func (c *SigV4Config) Validate() error {
+	if (c.AccessKey == "") != (c.SecretKey == "") {
+		return fmt.Errorf("must provide a AWS SigV4 Access key and Secret Key if credentials are specified in the SigV4 config")
+	}
+	return nil
+}
+func (c *SigV4Config) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	type plain SigV4Config
+	*c = SigV4Config{}
+	if err := unmarshal((*plain)(c)); err != nil {
+		return err
+	}
+	return c.Validate()
+}

sigv4/sigv4_config_test.go

@@ -0,0 +1,59 @@
+// Copyright 2021 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sigv4
+
+import (
+	"io/ioutil"
+	"strings"
+	"testing"
+
+	"gopkg.in/yaml.v2"
+)
+
+func loadSigv4Config(filename string) (*SigV4Config, error) {
+	content, err := ioutil.ReadFile(filename)
+	if err != nil {
+		return nil, err
+	}
+	cfg := SigV4Config{}
+	if err = yaml.UnmarshalStrict(content, &cfg); err != nil {
+		return nil, err
+	}
+	return &cfg, nil
+}
+
+func testGoodConfig(t *testing.T, filename string) {
+	_, err := loadSigv4Config(filename)
+	if err != nil {
+		t.Fatalf("Unexpected error parsing %s: %s", filename, err)
+	}
+}
+
+func TestGoodSigV4Configs(t *testing.T) {
+	filesToTest := []string{"testdata/sigv4_good.yaml", "testdata/sigv4_good.yaml"}
+	for _, filename := range filesToTest {
+		testGoodConfig(t, filename)
+	}
+}
+
+func TestBadSigV4Config(t *testing.T) {
+	filename := "testdata/sigv4_bad.yaml"
+	_, err := loadSigv4Config(filename)
+	if err == nil {
+		t.Fatalf("Did not receive expected error unmarshaling bad sigv4 config")
+	}
+	if !strings.Contains(err.Error(), "must provide a AWS SigV4 Access key and Secret Key") {
+		t.Errorf("Received unexpected error from unmarshal of %s: %s", filename, err.Error())
+	}
+}

sigv4/sigv4_test.go

@@ -0,0 +1,92 @@
+// Copyright 2021 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sigv4
+
+import (
+	"net/http"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/session"
+	signer "github.com/aws/aws-sdk-go/aws/signer/v4"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSigV4_Inferred_Region(t *testing.T) {
+	os.Setenv("AWS_ACCESS_KEY_ID", "secret")
+	os.Setenv("AWS_SECRET_ACCESS_KEY", "token")
+	os.Setenv("AWS_REGION", "us-west-2")
+
+	sess, err := session.NewSession(&aws.Config{
+		// Setting to an empty string to demostrate the default value from the yaml
+		// won't override the environment's region.
+		Region: aws.String(""),
+	})
+	require.NoError(t, err)
+	_, err = sess.Config.Credentials.Get()
+	require.NoError(t, err)
+
+	require.NotNil(t, sess.Config.Region)
+	require.Equal(t, "us-west-2", *sess.Config.Region)
+}
+
+func TestSigV4RoundTripper(t *testing.T) {
+	var gotReq *http.Request
+
+	rt := &sigV4RoundTripper{
+		region: "us-east-2",
+		next: promhttp.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
+			gotReq = req
+			return &http.Response{StatusCode: http.StatusOK}, nil
+		}),
+		signer: signer.NewSigner(credentials.NewStaticCredentials(
+			"test-id",
+			"secret",
+			"token",
+		)),
+	}
+	rt.pool.New = rt.newBuf
+
+	cli := &http.Client{Transport: rt}
+
+	req, err := http.NewRequest(http.MethodPost, "google.com", strings.NewReader("Hello, world!"))
+	require.NoError(t, err)
+
+	_, err = cli.Do(req)
+	require.NoError(t, err)
+	require.NotNil(t, gotReq)
+
+	origReq := gotReq
+	require.NotEmpty(t, origReq.Header.Get("Authorization"))
+	require.NotEmpty(t, origReq.Header.Get("X-Amz-Date"))
+
+	// Perform the same request but with a header that shouldn't included in the
+	// signature; validate that the Authorization signature matches.
+	t.Run("Ignored Headers", func(t *testing.T) {
+		req, err := http.NewRequest(http.MethodPost, "google.com", strings.NewReader("Hello, world!"))
+		require.NoError(t, err)
+
+		req.Header.Add("Uber-Trace-Id", "some-trace-id")
+
+		_, err = cli.Do(req)
+		require.NoError(t, err)
+		require.NotNil(t, gotReq)
+
+		require.Equal(t, origReq.Header.Get("Authorization"), gotReq.Header.Get("Authorization"))
+	})
+}

sigv4/testdata/sigv4_bad.yaml

@@ -0,0 +1,4 @@
+region: us-east-2
+access_key: AccessKey
+profile: profile
+role_arn: blah:role/arn

sigv4/testdata/sigv4_good.yaml

@@ -0,0 +1,5 @@
+region: us-east-2
+access_key: AccessKey
+secret_key: SecretKey
+profile: profile
+role_arn: blah:role/arn

sigv4/testdata/sigv4_good_empty_keys.yaml

@@ -0,0 +1,3 @@
+region: us-east-2
+profile: profile
+role_arn: blah:role/arn