Merge pull request #3 from parca-dev/headers

Add ability to attach arbitrary headers

Author: brancz

.github/workflows/container.yml

@@ -51,7 +51,7 @@ jobs:
         GORELEASER_CURRENT_TAG: "${{ env.goreleaser_current_tag }}"
     steps:
       - name: Check out the code
-        uses: actions/checkout@2541b1294d2704b0964813337f33b291d3f8596b # tag=v3.0.2
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # tag=v6.0.1
 
       - name: Set up Go
         uses: actions/setup-go@84cbf8094393cdc5fe1fe1671ff2647332956b1a # tag=v3.2.1
@@ -63,7 +63,7 @@ jobs:
         run: goreleaser release --rm-dist --skip-validate --skip-publish --snapshot --debug
 
       - name: Archive generated artifacts
-        uses: actions/upload-artifact@3cea5372237819ed00197afe530f5a7ea3e805c8 # tag=v3.1.0
+        uses: actions/upload-artifact@v5
         with:
           name: parca-push-dist-container
           if-no-files-found: error
@@ -94,15 +94,15 @@ jobs:
         run: dnf install --assumeyes --repo fedora git make jq
 
       - name: Check out code into the Go module directory
-        uses: actions/checkout@2541b1294d2704b0964813337f33b291d3f8596b # tag=v3.0.2
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # tag=v6.0.1
 
       - name: Set up Go
         uses: actions/setup-go@84cbf8094393cdc5fe1fe1671ff2647332956b1a # tag=v3.2.1
         with:
           go-version-file: 'go.mod'
           check-latest: true
 
-      - uses: actions/download-artifact@v3
+      - uses: actions/download-artifact@v4
         with:
           name: parca-push-dist-container
           path: goreleaser/dist

README.md

@@ -37,4 +37,7 @@ Flags:
                                  TLS.
       --remote-store-insecure-skip-verify
                                  Skip TLS certificate verification.
+      --remote-store-grpc-headers=KEY=VALUE;...
+                                 Additional gRPC headers to send with each
+                                 request (key=value pairs).
 ```

cmd/parca-push/main.go

@@ -32,6 +32,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
 )
 
 type flags struct {
@@ -45,11 +46,12 @@ type flags struct {
 
 // FlagsRemoteStore provides remote store configuration flags.
 type FlagsRemoteStore struct {
-	Address            string `kong:"help='gRPC address to send profiles and symbols to.'"`
-	BearerToken        string `kong:"help='Bearer token to authenticate with store.'"`
-	BearerTokenFile    string `kong:"help='File to read bearer token from to authenticate with store.'"`
-	Insecure           bool   `kong:"help='Send gRPC requests via plaintext instead of TLS.'"`
-	InsecureSkipVerify bool   `kong:"help='Skip TLS certificate verification.'"`
+	Address            string            `kong:"help='gRPC address to send profiles and symbols to.'"`
+	BearerToken        string            `kong:"help='Bearer token to authenticate with store.'"`
+	BearerTokenFile    string            `kong:"help='File to read bearer token from to authenticate with store.'"`
+	Insecure           bool              `kong:"help='Send gRPC requests via plaintext instead of TLS.'"`
+	InsecureSkipVerify bool              `kong:"help='Skip TLS certificate verification.'"`
+	GRPCHeaders        map[string]string `kong:"help='Additional gRPC headers to send with each request (key=value pairs).'"`
 }
 
 func main() {
@@ -132,16 +134,33 @@ func run(flags flags) error {
 	return g.Run()
 }
 
+func customHeadersUnaryInterceptor(headers map[string]string) grpc.UnaryClientInterceptor {
+	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+		for key, value := range headers {
+			ctx = metadata.AppendToOutgoingContext(ctx, key, value)
+		}
+		return invoker(ctx, method, req, reply, cc, opts...)
+	}
+}
+
 func grpcConn(reg prometheus.Registerer, flags FlagsRemoteStore) (*grpc.ClientConn, error) {
 	met := grpc_prometheus.NewClientMetrics()
 	met.EnableClientHandlingTimeHistogram()
 	reg.MustRegister(met)
 
-	opts := []grpc.DialOption{
-		grpc.WithUnaryInterceptor(
-			met.UnaryClientInterceptor(),
-		),
+	opts := []grpc.DialOption{}
+
+	// Add custom headers interceptor first if headers are provided
+	if len(flags.GRPCHeaders) > 0 {
+		opts = append(opts, grpc.WithUnaryInterceptor(
+			customHeadersUnaryInterceptor(flags.GRPCHeaders),
+		))
 	}
+
+	// Add metrics interceptor
+	opts = append(opts, grpc.WithUnaryInterceptor(
+		met.UnaryClientInterceptor(),
+	))
 	if flags.Insecure {
 		opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	} else {