signedrequests: Pipe context in signing requests (#6344)

Without this, our traces are often disconnected for requests that use
request signing.

Author: brancz

go.mod

@@ -3,6 +3,7 @@ module github.com/parca-dev/parca
 go 1.25.0
 
 require (
+	cloud.google.com/go/compute/metadata v0.9.0
 	cloud.google.com/go/storage v1.61.3
 	github.com/ClickHouse/clickhouse-go/v2 v2.43.0
 	github.com/alecthomas/kong v0.9.0
@@ -74,7 +75,6 @@ require (
 	cloud.google.com/go v0.123.0 // indirect
 	cloud.google.com/go/auth v0.20.0 // indirect
 	cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect
-	cloud.google.com/go/compute/metadata v0.9.0 // indirect
 	cloud.google.com/go/iam v1.5.3 // indirect
 	cloud.google.com/go/monitoring v1.24.3 // indirect
 	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0 // indirect

pkg/signedrequests/gcs.go

@@ -15,22 +15,41 @@ package signedrequests
 
 import (
 	"context"
+	"encoding/base64"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
 	"strconv"
 	"time"
 
+	"cloud.google.com/go/compute/metadata"
 	"cloud.google.com/go/storage"
 	"github.com/thanos-io/objstore/providers/gcs"
 	"golang.org/x/oauth2/google"
+	iamcredentials "google.golang.org/api/iamcredentials/v1"
 	"google.golang.org/api/option"
 	"gopkg.in/yaml.v3"
 )
 
 type GCSClient struct {
 	bucket *storage.BucketHandle
 	closer io.Closer
+
+	// googleAccessID is the service account email used to sign URLs. Detected
+	// once at construction time; passing it explicitly into SignedURLOptions
+	// lets us also override SignBytes, which is what enables trace-context
+	// propagation into the IAM SignBlob call.
+	googleAccessID string
+
+	// privateKey is set when the service account JSON contains one, enabling
+	// local signing without a round-trip to IAM.
+	privateKey []byte
+
+	// iamService is used for remote signing via IAM SignBlob when no private
+	// key is available. It is context-aware, so spans emitted by its HTTP
+	// transport are parented to the caller's trace.
+	iamService *iamcredentials.Service
 }
 
 func NewGCSClient(ctx context.Context, conf []byte) (*GCSClient, error) {
@@ -47,28 +66,107 @@ func NewGCSBucketWithConfig(ctx context.Context, gc gcs.Config) (*GCSClient, err
 		return nil, errors.New("missing Google Cloud Storage bucket name for stored blocks")
 	}
 
-	var opts []option.ClientOption
-
-	// If ServiceAccount is provided, use them in GCS client, otherwise fallback to Google default logic.
+	var (
+		creds *google.Credentials
+		err   error
+	)
 	if gc.ServiceAccount != "" {
-		credentials, err := google.CredentialsFromJSONWithType(ctx, []byte(gc.ServiceAccount), google.ServiceAccount, storage.ScopeFullControl)
-		if err != nil {
-			return nil, fmt.Errorf("create credentials from JSON: %w", err)
-		}
-		opts = append(opts, option.WithCredentials(credentials))
+		creds, err = google.CredentialsFromJSONWithType(ctx, []byte(gc.ServiceAccount), google.ServiceAccount, storage.ScopeFullControl)
+	} else {
+		creds, err = google.FindDefaultCredentials(ctx, storage.ScopeFullControl)
+	}
+	if err != nil {
+		return nil, fmt.Errorf("load credentials: %w", err)
 	}
 
-	opts = append(opts, option.WithUserAgent("parca"))
+	opts := []option.ClientOption{
+		option.WithCredentials(creds),
+		option.WithUserAgent("parca"),
+	}
 
 	gcsClient, err := storage.NewClient(ctx, opts...)
 	if err != nil {
 		return nil, err
 	}
 
-	return &GCSClient{
-		bucket: gcsClient.Bucket(gc.Bucket),
-		closer: gcsClient,
-	}, nil
+	googleAccessID, privateKey, err := signingIdentity(ctx, creds)
+	if err != nil {
+		return nil, fmt.Errorf("resolve signing identity: %w", err)
+	}
+
+	c := &GCSClient{
+		bucket:         gcsClient.Bucket(gc.Bucket),
+		closer:         gcsClient,
+		googleAccessID: googleAccessID,
+		privateKey:     privateKey,
+	}
+
+	if len(privateKey) == 0 {
+		iamService, err := iamcredentials.NewService(ctx, opts...)
+		if err != nil {
+			return nil, fmt.Errorf("create iamcredentials service: %w", err)
+		}
+		c.iamService = iamService
+	}
+
+	return c, nil
+}
+
+// signingIdentity returns the service account email and (optionally) a private
+// key for URL signing. Mirrors storage.BucketHandle.detectDefaultGoogleAccessID,
+// which is unexported — we need this surfaced because overriding SignBytes (to
+// thread context through to IAM) means the library no longer fills in
+// GoogleAccessID for our closure.
+func signingIdentity(ctx context.Context, creds *google.Credentials) (string, []byte, error) {
+	if len(creds.JSON) > 0 {
+		var sa struct {
+			ClientEmail string `json:"client_email"`
+			PrivateKey  string `json:"private_key"`
+		}
+		if err := json.Unmarshal(creds.JSON, &sa); err == nil && sa.ClientEmail != "" {
+			return sa.ClientEmail, []byte(sa.PrivateKey), nil
+		}
+	}
+
+	if !metadata.OnGCE() {
+		return "", nil, errors.New("could not resolve service account email from credentials JSON and not running on GCE")
+	}
+	email, err := metadata.EmailWithContext(ctx, "default")
+	if err != nil {
+		return "", nil, fmt.Errorf("read service account email from GCE metadata: %w", err)
+	}
+	if email == "" {
+		return "", nil, errors.New("empty service account email from GCE metadata")
+	}
+	return email, nil, nil
+}
+
+// signBytesWithContext returns a SignBytes function that calls IAM SignBlob
+// with the given context. Using our context here means the HTTP span emitted
+// by the IAM client inherits the caller's trace.
+func (c *GCSClient) signBytesWithContext(ctx context.Context) func([]byte) ([]byte, error) {
+	return func(in []byte) ([]byte, error) {
+		resp, err := c.iamService.Projects.ServiceAccounts.SignBlob(
+			fmt.Sprintf("projects/-/serviceAccounts/%s", c.googleAccessID),
+			&iamcredentials.SignBlobRequest{
+				Payload: base64.StdEncoding.EncodeToString(in),
+			},
+		).Context(ctx).Do()
+		if err != nil {
+			return nil, fmt.Errorf("iam sign blob: %w", err)
+		}
+		return base64.StdEncoding.DecodeString(resp.SignedBlob)
+	}
+}
+
+func (c *GCSClient) signedURLOptions(ctx context.Context, base *storage.SignedURLOptions) *storage.SignedURLOptions {
+	base.GoogleAccessID = c.googleAccessID
+	if len(c.privateKey) > 0 {
+		base.PrivateKey = c.privateKey
+	} else {
+		base.SignBytes = c.signBytesWithContext(ctx)
+	}
+	return base
 }
 
 func (c *GCSClient) Close() error {
@@ -81,22 +179,22 @@ func (c *GCSClient) SignedPUT(
 	size int64,
 	expiry time.Time,
 ) (string, error) {
-	return c.bucket.SignedURL(objectKey, &storage.SignedURLOptions{
+	return c.bucket.SignedURL(objectKey, c.signedURLOptions(ctx, &storage.SignedURLOptions{
 		Method:  "PUT",
 		Expires: expiry,
 		Headers: []string{
 			"X-Upload-Content-Length:" + strconv.FormatInt(size, 10),
 		},
-	})
+	}))
 }
 
 func (c *GCSClient) SignedGET(
 	ctx context.Context,
 	objectKey string,
 	expiry time.Time,
 ) (string, error) {
-	return c.bucket.SignedURL(objectKey, &storage.SignedURLOptions{
+	return c.bucket.SignedURL(objectKey, c.signedURLOptions(ctx, &storage.SignedURLOptions{
 		Method:  "GET",
 		Expires: expiry,
-	})
+	}))
 }