feat(grpc): allow files to be read by add_metadata_jmespath_expression

This allows add_metadata_jmespath_expression to refer to the contents of
certain pre-declared files which are reloaded on a user-specified
schedule. For example, if you want to use k8s bound service
account tokens you could add the following to your Pod:

  volumes:
  - name: token-vol
    projected:
      sources:
      - serviceAccountToken:
          audience: my-buildbarn-instance
          expirationSeconds: 3600
          path: buildbarn

Assuming this is mounted at '/tokens', you could then specify in your
buildbarn config:

  addMetadataJmespathExpression: |||
    {
      "authorization": [std.format('bearer %s', files.token)]
    }
  |||,
  addMetadataJmespathFiles: [
    {
      key: "token",
      path: "/tokens/buildbarn",
      refreshInterval: "1800s",
    }
  ]

This is quite useful for k8s service account tokens, as the maximum
validity is often capped. Likewise this can also be used for Google
service account id tokens, which also have a relatively short maximum
validity.

Author: tomgr

pkg/grpc/BUILD.bazel

@@ -143,6 +143,7 @@ go_test(
         "@org_golang_google_grpc//peer",
         "@org_golang_google_grpc//status",
         "@org_golang_google_protobuf//proto",
+        "@org_golang_google_protobuf//types/known/durationpb",
         "@org_golang_google_protobuf//types/known/emptypb",
         "@org_golang_google_protobuf//types/known/structpb",
         "@org_uber_go_mock//gomock",

pkg/grpc/base_client_factory.go

@@ -179,11 +179,15 @@ func (cf baseClientFactory) NewClientFromConfiguration(config *configuration.Cli
 
 	// Optional: metadata extraction.
 	if jmesExpression := config.AddMetadataJmespathExpression; jmesExpression != "" {
+		filesProvider, err := NewJMESPathMetadataFileProvider(context.Background(), config.AddMetadataJmespathFiles)
+		if err != nil {
+			return nil, util.StatusWrap(err, "Failed to create JMESPath metadata file provider")
+		}
 		expr, err := jmespath.Compile(jmesExpression)
 		if err != nil {
 			return nil, util.StatusWrap(err, "Failed to compile JMESPath expression")
 		}
-		extractor, err := NewJMESPathMetadataExtractor(expr)
+		extractor, err := NewJMESPathMetadataExtractor(expr, filesProvider)
 		if err != nil {
 			return nil, util.StatusWrap(err, "Failed to create JMESPath extractor")
 		}

pkg/grpc/jmespath_extractor.go

@@ -2,8 +2,13 @@ package grpc
 
 import (
 	"context"
+	"log"
+	"os"
+	"sync"
+	"time"
 
 	"github.com/buildbarn/bb-storage/pkg/auth"
+	pb "github.com/buildbarn/bb-storage/pkg/proto/configuration/grpc"
 	"github.com/buildbarn/bb-storage/pkg/util"
 	"github.com/jmespath/go-jmespath"
 
@@ -22,19 +27,24 @@ import (
 //
 //	{
 //	    "authenticationMetadata": value,
+//		"files": map<string, string>,
 //	    "incomingGRPCMetadata": map<string, repeated string>
 //	}
-func NewJMESPathMetadataExtractor(expression *jmespath.JMESPath) (MetadataExtractor, error) {
+func NewJMESPathMetadataExtractor(expression *jmespath.JMESPath, files *MetadataFileProvider) (MetadataExtractor, error) {
 	return func(ctx context.Context) (MetadataHeaderValues, error) {
-		searchContext := make(map[string]interface{}, 2)
+		searchContext := make(map[string]interface{}, 3)
 		if authenticationMetadata := auth.AuthenticationMetadataFromContext(ctx); authenticationMetadata != nil {
 			searchContext["authenticationMetadata"] = authenticationMetadata.GetRaw()
 		}
 
+		if files != nil {
+			searchContext["files"] = files.getMetadata()
+		}
+
 		if md, ok := metadata.FromIncomingContext(ctx); ok {
 			// JMESPath only treats map[string]interface{}, struct, or *struct as map types,
 			// so we need to copy from the map[string][]string.
-			incomingGRPCMetadata := make(map[string]interface{}, len(md))
+			incomingGRPCMetadata := make(map[string]any, len(md))
 			for k, rawVs := range md {
 				vs := make([]interface{}, 0, len(rawVs))
 				for _, rawV := range rawVs {
@@ -84,3 +94,67 @@ func matchToHeaders(rawMatch interface{}) (MetadataHeaderValues, error) {
 	}
 	return headers, nil
 }
+
+// MetadataFileProvider makes the contents of files available as
+// metadata to JMESPath expressions.
+type MetadataFileProvider struct {
+	lock            sync.RWMutex
+	files           []*pb.ClientConfiguration_RefreshedFile
+	currentContents map[string]string
+}
+
+// NewJMESPathMetadataFileProvider creates a MetadataFileProvider that
+// reads files from the filesystem and makes their contents available
+// as metadata. The contents of the files are reloaded periodically.
+func NewJMESPathMetadataFileProvider(ctx context.Context, files []*pb.ClientConfiguration_RefreshedFile) (*MetadataFileProvider, error) {
+	provider := &MetadataFileProvider{
+		files:           files,
+		currentContents: make(map[string]string),
+	}
+	for _, file := range files {
+		contents, err := readFile(file.Path)
+		if err != nil {
+			return nil, err
+		}
+		provider.currentContents[file.Key] = contents
+		go func() {
+			t := time.NewTicker(file.RefreshInterval.AsDuration())
+			for {
+				select {
+				case <-t.C:
+				case <-ctx.Done():
+					t.Stop()
+					return
+				}
+				contents, err := readFile(file.Path)
+				if err != nil {
+					log.Printf("Failed to reload %s file: %v", file.Path, err)
+				} else {
+					provider.lock.Lock()
+					provider.currentContents[file.Key] = contents
+					provider.lock.Unlock()
+				}
+			}
+		}()
+	}
+
+	return provider, nil
+}
+
+func readFile(path string) (string, error) {
+	content, err := os.ReadFile(path)
+	if err != nil {
+		return "", util.StatusWrapf(err, "Failed to read %q", path)
+	}
+	return string(content), nil
+}
+
+func (p *MetadataFileProvider) getMetadata() map[string]any {
+	p.lock.RLock()
+	defer p.lock.RUnlock()
+	metadata := make(map[string]any, len(p.currentContents))
+	for k, v := range p.currentContents {
+		metadata[k] = v
+	}
+	return metadata
+}

pkg/grpc/jmespath_extractor_test.go

@@ -2,11 +2,15 @@ package grpc_test
 
 import (
 	"context"
+	"os"
+	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/buildbarn/bb-storage/pkg/auth"
 	"github.com/buildbarn/bb-storage/pkg/grpc"
 	auth_pb "github.com/buildbarn/bb-storage/pkg/proto/auth"
+	pb "github.com/buildbarn/bb-storage/pkg/proto/configuration/grpc"
 	"github.com/buildbarn/bb-storage/pkg/testutil"
 	"github.com/buildbarn/bb-storage/pkg/util"
 	"github.com/jmespath/go-jmespath"
@@ -15,6 +19,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/durationpb"
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
@@ -25,7 +30,7 @@ func TestJMESPathMetadataExtractorSimple(t *testing.T) {
 		"this-is-static": ['and great'],
 		"hdr-from-both": [incomingGRPCMetadata.whiz[0], authenticationMetadata.public],
 		"optional-hdr": incomingGRPCMetadata.missing
-	}`))
+	}`), nil)
 	require.NoError(t, err)
 
 	// We compare with metadata.Pairs because JMESPath evaluation traverses maps
@@ -52,7 +57,7 @@ func TestJMESPathMetadataExtractorSimple(t *testing.T) {
 }
 
 func TestJMESPathMetadataExtractorAuthMatchToString(t *testing.T) {
-	extractor, err := grpc.NewJMESPathMetadataExtractor(jmespath.MustCompile(`{"hdr": authenticationMetadata.public}`))
+	extractor, err := grpc.NewJMESPathMetadataExtractor(jmespath.MustCompile(`{"hdr": authenticationMetadata.public}`), nil)
 	require.NoError(t, err)
 
 	// The resulting header value must be a list. Yielding a string
@@ -66,7 +71,7 @@ func TestJMESPathMetadataExtractorAuthMatchToString(t *testing.T) {
 }
 
 func TestJMESPathMetadataExtractorAuthMatchToHeterogenousSlice(t *testing.T) {
-	extractor, err := grpc.NewJMESPathMetadataExtractor(jmespath.MustCompile(`{"hdr": authenticationMetadata.public}`))
+	extractor, err := grpc.NewJMESPathMetadataExtractor(jmespath.MustCompile(`{"hdr": authenticationMetadata.public}`), nil)
 	require.NoError(t, err)
 
 	// Each of the header values should be a valid string. Integer
@@ -83,3 +88,54 @@ func TestJMESPathMetadataExtractorAuthMatchToHeterogenousSlice(t *testing.T) {
 	_, err = extractor(ctx)
 	testutil.RequireEqualStatus(t, status.Errorf(codes.InvalidArgument, "Failed to extract JMESPath result: Non-string metadata value"), err)
 }
+
+func TestNewJMESPathMetadataFileProvider(t *testing.T) {
+	// Create a temporary file with test content.
+	tempDir := t.TempDir()
+	filePath := filepath.Join(tempDir, "test-token")
+	err := os.WriteFile(filePath, []byte("token-value1"), 0o644)
+	require.NoError(t, err)
+
+	// Build the extractor.
+	context, cancel := context.WithCancel(context.Background())
+	provider, err := grpc.NewJMESPathMetadataFileProvider(context, []*pb.ClientConfiguration_RefreshedFile{
+		{
+			Key:             "token",
+			Path:            filePath,
+			RefreshInterval: durationpb.New(time.Millisecond),
+		},
+	})
+	require.NoError(t, err)
+	extractor, err := grpc.NewJMESPathMetadataExtractor(
+		jmespath.MustCompile(`{"authorization": [files.token]}`),
+		provider,
+	)
+	require.NoError(t, err)
+
+	// Validate the initial contents are correct.
+	headers, err := extractor(context)
+	require.NoError(t, err)
+	want := grpc.MetadataHeaderValues([]string{
+		"authorization", "token-value1",
+	})
+	require.Equal(t, want, headers)
+
+	// Modify the file.
+	err = os.WriteFile(filePath, []byte("token-value2"), 0o644)
+	require.NoError(t, err)
+
+	// Wait for the file to be reloaded. This is potentially fragile.
+	time.Sleep(time.Second)
+
+	// Validate the updated contents are correct.
+	headers, err = extractor(context)
+	require.NoError(t, err)
+	want = grpc.MetadataHeaderValues([]string{
+		"authorization", "token-value2",
+	})
+	require.Equal(t, want, headers)
+
+	// Cancel the context to stop the file reloading.
+	cancel()
+	time.Sleep(time.Second)
+}