Using Embedding to Disambiguate Types in Go

A few weeks ago, I stumbled upon a trick to use Go’s type system to disambiguate between variants of credentials on a RPC service I work with.

The premise of the example is that we’re writing an RPC request handler that needs to perform two outbound RPCs – one to an external service that uses user credentials, and another to an internal service that uses a set of service-wide credentials:

func (*h RPCHandler) CreateWidget(req *WidgetRequest, resp *WidgetResponse) error {
    userCreds, err := h.CredentialProvider.GetUserCredentials()
    if err != nil {
        return err
    }
    h.ExternalServiceClient.CallPublicAPI(&PublicAPIReq{...}, userCreds)

    ...

    serviceCreds, err := h.CredentialProvider.GetServiceCredentials()
    if err != nil {
        return err
    }
    h.InternalServiceClient.CallPrivateAPI(&PrivateAPIReq{...}, serviceCreds)

    ...
}

This works as expected. For each outbound RPC, we get the relevant credentials for that request and forward them along.

Now, suppose you have this code running for a while, and various developers make changes to it. The CreateWidget handler eventually becomes quite large, and the credentials may be reused in multiple outbound RPC requests as time goes on. Perhaps folks copy/paste code between RPC handlers, and drop the user* and service* naming convention.

This can eventually lead to the “oops” condition below, where you inadvertently call a service with the wrong type of credentials:

func (*h RPCHandler) CreateWidget(req *WidgetRequest, resp *WidgetResponse) error {
    creds, err := h.CredentialProvider.GetUserCredentials()
    if err != nil {
        return err
    }
    h.ExternalServiceClient.CallPublicAPI(&PublicAPIReq{...}, creds)

    ...

    // Oops! Permission denied.
    h.InternalServiceClient.CallPrivateAPI(&PrivateAPIReq{...}, creds)
    ...
}

Whoops! Now the internal (private) API is being called with the external user credentials. However, the type system didn’t catch this. Depending on how thorough your testing, you may only catch this once you run your integration or end-to-end tests (or, in a deployed environment if you lack E2E tests).

Why didn’t the compiler catch this? Our credential structs are the same type, so even though there is a semantic difference between the user and service credentials, this difference isn’t type checkable:

type CredentialProvider interface {
    GetUserCredentials() (*auth.Credentials, error)
    GetServiceCredentials() (*auth.Credentials, error)
}

What could we do to make this semantic difference something that the compiler could check? The solution I went with is to make a struct type for each semantic type of credentials, using type embedding:

type UserCredentials struct {
    *auth.Creds
}

type ServiceCredentials struct {
    *auth.Creds
}

These credential types can more-or-less act the same, but they’re wrapped in an outer type so that we can guard usage in downstream functions.

Now we can update our credential provider interface, and the credential generator itself to return these new types:

type CredentialProvider interface {
    GetUserCredentials() (*UserCredentials, error)
    GetServiceCredentials() *ServiceCredentials, error)
}

func (g *RealCredentialGenerator) UserCredentials() (*UserCredentials, error) {
    creds, err := g.userCredentials()
    if err != nil {
        return nil, err
    }
    return &UserCredentials{creds}, nil
}

func (g *RealCredentialGenerator) ServiceCredentials() (*ServiceCredentials, error) {
    creds, err := g.serviceCredentials()
    if err != nil {
        return nil, err
    }
    return &ServiceCredentials{creds}, nil
}

And in the downstream users of the credentials, you update the function signatures accordingly:

func (c *InternalServiceClient) CallPublicAPI(req *PublicAPIRequest, creds *UserCredentials) error { ... }
func (c *ExternalServiceClient) CallPrivateAPI(req *PrivateAPIRequest, creds *ServiceCredentials) error { ... }

How does this help? Going back to the example, now if we mess up and try providing the wrong type of credential to one of our RPC clients, we get a compile time failure:

func (*h RPCHandler) CreateWidget(req *WidgetRequest, resp *WidgetResponse) error {
    creds, err := h.CredentialProvider.GetUserCredentials()
    if err != nil {
        return err
    }
    h.ExternalServiceClient.CallPublicAPI(&PublicAPIReq{...}, creds)

    ...

    // Compiler error:
    // cannot use creds (type *UserCredentials) as type *ServiceCredentials
    // in argument to CallPrivateAPI
    h.InternalServiceClient.CallPrivateAPI(&PrivateAPIReq{...}, creds)
    ...
}

Type embedding isn’t new to me, but I thought that this was a pretty neat use of it. It had the added benefit of forcing us to audit all the downstream consumers of these credentials to make sure that they were all correct. We discovered a few additional bugs this way.