Repository: codenotary/immudb
File: pkg/auth/clientinterceptors.go (Go)

Code before:
import (
	"context"

	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// ClientStreamInterceptor gRPC client interceptor for streams
func ClientStreamInterceptor(token string) func(context.Context, *grpc.StreamDesc, *grpc.ClientConn, string, grpc.Streamer, ...grpc.CallOption) (grpc.ClientStream, error) {
	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
		ctx = updateAuthHeader(ctx, token)
		return streamer(ctx, desc, cc, method, opts...)
	}
}

// ClientUnaryInterceptor gRPC client interceptor for unary methods
func ClientUnaryInterceptor(token string) func(context.Context, string, interface{}, interface{}, *grpc.ClientConn, grpc.UnaryInvoker, ...grpc.CallOption) error {
	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		ctx = updateAuthHeader(ctx, token)
		return invoker(ctx, method, req, reply, cc, opts...)
	}
}

func updateAuthHeader(ctx context.Context, token string) context.Context {
	if md, ok := metadata.FromOutgoingContext(ctx); ok && len(md.Get("authorization")) > 0 {
		// Token from existing metadata overwrites the one stored in
		token = md.Get("authorization")[0]
	}

	md, ok := metadata.FromOutgoingContext(ctx)
	if !ok {
		md = metadata.New(nil)
	}
	md.Set("authorization", "Bearer "+token)
	return metadata.NewOutgoingContext(ctx, md)
}


Review comment (suggestion):
existent token shouldn't be replaced by the one provided as parameter then? if so, I'd suggest changing the comment from `stored in` by `provided in current call` or something like that. And the method could do a fast return if no change is made into the metadata.

However, I would expect the provided as parameter having more precedence actually...

Code after:
import (
	"context"

	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// ClientStreamInterceptor gRPC client interceptor for streams
func ClientStreamInterceptor(token string) func(context.Context, *grpc.StreamDesc, *grpc.ClientConn, string, grpc.Streamer, ...grpc.CallOption) (grpc.ClientStream, error) {
	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
		ctx = updateAuthHeader(ctx, token)
		return streamer(ctx, desc, cc, method, opts...)
	}
}

// ClientUnaryInterceptor gRPC client interceptor for unary methods
func ClientUnaryInterceptor(token string) func(context.Context, string, interface{}, interface{}, *grpc.ClientConn, grpc.UnaryInvoker, ...grpc.CallOption) error {
	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		ctx = updateAuthHeader(ctx, token)
		return invoker(ctx, method, req, reply, cc, opts...)
	}
}

// updateAuthHeader ensures the grpc metadata in the context contains the correct
// authorization header value. The token may be either taken from the client
// object where it is managed by a token service through the `token` argument,
// or it can be given in the metadata inside the context.
func updateAuthHeader(ctx context.Context, token string) context.Context {
	if md, ok := metadata.FromOutgoingContext(ctx); ok && len(md.Get("authorization")) > 0 {
		// The token provided through the metadata has a higher priority than the
		// one set for the whole client - this allows customization of the token
		// per each API call.
		token = md.Get("authorization")[0]
	}

	md, ok := metadata.FromOutgoingContext(ctx)
	if !ok {
		md = metadata.New(nil)
	}
	// The final token value must be provided with the `Bearer ` prefix.
	md.Set("authorization", "Bearer "+token)
	return metadata.NewOutgoingContext(ctx, md)
}