Go GRPC Refresh token for a bidirectional stream

TLDR: I am looking for a way to update headers on an open stream for each call to stream.Send(msg) without closing the stream and opening a new one.

Summary

I have a GRPC client and server built to handle bidirectional streams. To authenticate with the server the client must send a JWT in the request headers, set as "authorization". The token is valid for 30 minutes. After the token has expired, the server will terminate the connection.

I am looking for a way to refresh my authorization token from the client, and keep the stream open. The client should run in a loop executing a new request every 30 minutes with the updated token, and the updated payload. I have not seen a way to update a header from the client side for an already opened stream.

Let's look at some code to get an idea of what the client side looks like. The code below has a function to create a new instance of the client, and another function to establish the connection to the GRPC server.

func NewWatchClient(config *Config, logger *logrus.Logger) (*WatchClient, error) {
	cc, err := newConnection(config, logger)
	if err != nil {
		return nil, err
	}

	service := proto.NewWatchServiceClient(cc)

	return &WatchClient{
		config:  config,
		conn:    cc,
		logger:  entry,
		service: service,
	}, nil
}

func newConnection(config *Config, logger *logrus.Logger) (*grpc.ClientConn, error) {
	address := fmt.Sprintf("%s:%d", config.Host, config.Port)

    // rpcCredential implements credentials.PerRPCCredentials
	rpcCredential := newTokenAuth(config.Auth, config.TenantID)

	return grpc.Dial(
		address,
        grpc.WithPerRPCCredentials(rpcCredential),
	)
}

Looking at the newConnection function above I can see that there is a call to another function, newTokenAuth, to create an auth token. This func returns a struct that implements the PerRPCCredentials interface.

There are two ways to set the authorization for a request.

1. Use grpc.WithPerRPCCredentials to add the authorization at the time of creating the connection to the server.

2. Use grpc.PerRPCCredentials to add the authorization to each stream opened on the connection to the server.

In this case, I am using grpc.WithPerRPCCredentials to attach the token at the time of creating the connection to the server.

Now, let's take a look at the definition of PerRPCCredentials.

type PerRPCCredentials interface {
	// GetRequestMetadata gets the current request metadata, refreshing
	// tokens if required. This should be called by the transport layer on
	// each request, and the data should be populated in headers or other
	// context. If a status code is returned, it will be used as the status
	// for the RPC. uri is the URI of the entry point for the request.
	// When supported by the underlying implementation, ctx can be used for
	// timeout and cancellation. Additionally, RequestInfo data will be
	// available via ctx to this call.
	// TODO(zhaoq): Define the set of the qualified keys instead of leaving
	// it as an arbitrary string.
	GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error)
	// RequireTransportSecurity indicates whether the credentials requires
	// transport security.
	RequireTransportSecurity() bool
}

The interface requires that you define two methods. The documentation of GetRequestMetadata says

> GetRequestMetadata gets the current request metadata, refreshing tokens if required

So, it looks like my implementation of PerRPCCredentials should be able to handle a token refresh for my stream or connection. Let's take a look at my implementation of PerRPCCredentials.

// tokenAuth implements the PerRPCCredentials interface
type tokenAuth struct {
	tenantID       string
	tokenRequester auth.PlatformTokenGetter
    token          string
}

// RequireTransportSecurity leave as false for now
func (tokenAuth) RequireTransportSecurity() bool {
	return false
}

// GetRequestMetadata sets the http header prior to transport
func (t tokenAuth) GetRequestMetadata(_ context.Context, _ ...string) (map[string]string, error) {
	token, err := t.tokenRequester.GetToken()
	if err != nil {
		return nil, err
	}
	t.token = token

	go func() {
		time.Sleep(25 * time.Minute)
		token, _ := t.tokenRequester.GetToken()
		t.token = token
	}()

	return map[string]string{
		"tenant-id": t.tenantID,
		"authorization":     "Bearer " + t.token,
	}, nil
}

As you can see, the call to GetRequestMetadata will establish a go routine that will attempt to refresh a token every 25 minutes. Adding a go routine right here is probably not the right way to do it. It was an attempt to get the auth header to refresh, which doesn't work.

Let's take a look at the stream.

func (w WatchClient) CreateWatch() error {
	topic := &proto.Request{SelfLink: w.config.TopicSelfLink}

	stream, err := w.service.CreateWatch(context.Background())
	if err != nil {
		return err
	}

	for {
		err = stream.Send(topic)
		if err != nil {
			return err
		}
		time.Sleep(25 * time.Minute)
	}
}

The client sends a message on the stream every 25 minutes. All I'm looking to get here is that when stream.Send is called, the updated token is also sent.

This function, GetRequestMetadata only gets called once, regardless if I am setting the auth through grpc.WithPerRPCCredentials or grpc.PerRPCCredsCallOption so there appears to be no way to update the authorization header.

If you have any idea what I have missed in my attempt to utilize the PerRPCCredentials for token refresh then please let me know.

Thank you.

Headers are sent at the beginning of an RPC, and cannot be updated during the RPC. If you need to send data during the life of a stream, it needs to be part of the request message in your proto definition.