Any way to add a custom header to a presigned sts:GetCallerIdentity request?

[nklaassen]

I'm working on an application where we currently add a custom header to an sts:GetCallerIdentity request before signing. We do this to implement a challenge-response mechanism where the server issues a random challenge string and the client must respond with a signed sts:GetCallerIdentity request including the challenge string in a signed header. The server is the one which actually sends the request to the sts endpoint to confirm the identity of the client.

You can do this fairly easily in the v1 sdk with the following snippet:

stsService := sts.New(sess)
req, _ := stsService.GetCallerIdentityRequest(&sts.GetCallerIdentityInput{})
// set challenge header
req.HTTPRequest.Header.Set(challengeHeaderKey, challenge)
// request json for simpler parsing
req.HTTPRequest.Header.Set("Accept", "application/json")
// sign the request, including headers
if err := req.Sign(); err != nil {
    return nil, err
}
// write the signed HTTP request to a buffer
var signedRequest bytes.Buffer
if err := req.HTTPRequest.Write(&signedRequest); err != nil {
    return nil, err
}

Is there any nice way to do this with the v2 sdk? I'm currently looking at mucking about with the smithy middleware.

[nickzelei]

You should be able to do this by overriding the sts.HTTPPresignerV4 I think.

package main

import (
	"context"
	"net/http"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/sts"
)

func main() {
	ctx := context.Background()
	cfg, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		panic(err)
	}
	stsclient := sts.NewFromConfig(cfg)
	presignclient := sts.NewPresignClient(stsclient)

	_, err = presignclient.PresignGetCallerIdentity(ctx, &sts.GetCallerIdentityInput{}, func(opt *sts.PresignOptions) {
		opt.Presigner = NewCustomHTTPPresignerV4(opt.Presigner, map[string]string{
			"Foo": "bar",
		})
	})
	if err != nil {
		panic(err)
	}
}

type customHTTPPresignerV4 struct {
	client  sts.HTTPPresignerV4
	headers map[string]string
}

func NewCustomHTTPPresignerV4(client sts.HTTPPresignerV4, headers map[string]string) sts.HTTPPresignerV4 {
	return &customHTTPPresignerV4{
		client:  client,
		headers: headers,
	}
}

func (p *customHTTPPresignerV4) PresignHTTP(
	ctx context.Context, credentials aws.Credentials, r *http.Request,
	payloadHash string, service string, region string, signingTime time.Time,
	optFns ...func(*v4.SignerOptions),
) (url string, signedHeader http.Header, err error) {
	for key, val := range p.headers {
		r.Header.Add(key, val)
	}
	return p.client.PresignHTTP(ctx, credentials, r, payloadHash, service, region, signingTime, optFns...)
}

[nklaassen]

I'm finally getting back to this. I'm pretty sure the header needs to be added before the payloadHash is computed, or at least it may need to be recomputed

[nklaassen]

Oh it turns out you can use a custom EndpointResolverV2, or add a Serializer to the middleware stack. I need a custom endpoint resolver anyway, so doing it like this.

type customResolver struct {
	defaultResolver sts.EndpointResolverV2
	challenge       string
}

func newCustomResolver(challenge string) *customResolver {
	return &customResolver{
		defaultResolver: sts.NewDefaultEndpointResolverV2(),
		challenge:       challenge,
	}
}

func (r customResolver) ResolveEndpoint(ctx context.Context, params sts.EndpointParameters) (smithyendpoints.Endpoint, error) {
	endpoint, err := r.defaultResolver.ResolveEndpoint(ctx, params)
	if err != nil {
		return endpoint, err
	}
	endpoint.Headers.Add(challengeHeaderKey, r.challenge)
	return endpoint, nil
}