Writing an Istio WASM Plugin in Go for migrating 100s of services to new auth strategy (Part 5)

Part 5: Testing your WASM plugin

This article will mostly be about the building of the WASM plugin itself and less on the reasons why we needed it, but a 1-liner explanation is “We had to route our service-to-service traffic through an Nginx proxy to acquire an auth JWT from an Auth Service, and wanted to remove that extra network hop by having the WASM plugin call the Auth Service directly instead of Nginx”.

This section is all about adding tests to the plugin. To refer back to earlier parts use the links below:

• Part 1: Getting the basic project bootstrapped and running
• Part 2: Calling out to an external service
• Part 3: Configuring the plugin and accessing local container files
• Part 4: Metrics
• Part 5: Testing your WASM plugin
• Part 6: Gotchas

“In the software engineering world, unit-testing is like a trust fall exercise; you code with your eyes closed and just hope your tests will catch you.” — ChatGTP

Considering that we’re writing a WASM plugin that is intercepting a substantial portion of requests going through our mesh, it’s probably a good idea to write some tests to prevent regressions! 🙈

Unfortunately it’s not as easy as providing our methods with some parameters and verifying the output as all side-effects are sent to the Envoy Host worker, and if we try to call proxywasm functions in our tests, we’ll get a lot of errors or segfaults as it tries to access uninitialized memory.

Thankfully the folks at Tetralabs have provided a proxywasm mocking framework for testing. This avoids having to spin up an actual Envoy Host process to exercise our WASM binary.

There is still a choice in how we want to setup our tests:

• Run our tests with tinygo test ./...in memory, so the test process acts as the Envoy Host itself.
• Run our tests with go test ./...in memory, so the test process acts as the Envoy Host itself.
• Compile our WASM binary with tinygo build and run our tests with tinygo test ./...by loading in the output WASM binary.
• Compile our WASM binary with tinygo build and run our tests with go test ./...by loading in the output WASM binary.

There are pros/cons to each, but I found it’s easier to generate the WASM binary with TinyGo and run the tests with Go to be able to take advantage of full support of type reflection as it makes it easier to validate JSON or use any community libraries like testify. Building with TinyGo still allows validation that bugs haven’t been introduced based on TinyGo's limited reflection or other limitations.

The main downside to testing this way is that I’m constantly forgetting to rebuild the WASM binary between test runs if I’ve changed the code. 🤦‍♂ The TetrateLabs SDK shows in their samples a way to do it all in one which you might prefer, and to be honest I might switch myself if I keep forgetting to recompile many more times. 😅

As ever, you can follow along with the code from the sample repo.

To start with, we have to create a test_helper.go file to enable the loading of a pre-compiled WASM binary with proxy-wasm bootstrapping.

package internal

import (
 "os"
 "testing"

 "github.com/stretchr/testify/require"
 "github.com/tetratelabs/proxy-wasm-go-sdk/proxywasm/proxytest"
 "github.com/tetratelabs/proxy-wasm-go-sdk/proxywasm/types"
)

const DefaultTestConfig = `
{
 "auth_authority": "auth",
 "auth_cluster_name": "auth",
 "auth_timeout_ms": 5
}`

// InitPlugin loads the WASM binary that should have been previously compiled.
func InitPlugin(t *testing.T) proxytest.WasmVMContext {
  wasm, err := os.ReadFile("../main.wasm")
  if err != nil {
    t.Fatalf("wasm not found")
  }

  // Create a new vmContext based on the loaded WASM binary.
  // You can then use this to call NewHostEmulator() on to create an Envoy Worker context
  vmContext, err := proxytest.NewWasmVMContext(wasm)
  require.NoError(t, err)
  return vmContext
}

// NewContextWithConfig creates a new 'Envoy Host' environment based on a previously loaded WasmVMContext
// Returns:
//   HostEmulator handle that allows you to interact with the WASM plugin via Envoy Host function calls
//   uint32 context ID to pass to Envoy Host function to reference any state with this context
//   function reference to call to reset all state for this particular Envoy Host context
func NewContextWithConfig(t *testing.T, vmContext proxytest.WasmVMContext, config string) (proxytest.HostEmulator, uint32, func()) {
  opt := proxytest.
    NewEmulatorOption().
    WithPluginConfiguration([]byte(config)).
    WithVMContext(vmContext)
  host, reset := proxytest.NewHostEmulator(opt)
  
  // Call our WASM's OnVMStart.
  require.Equal(t, types.OnVMStartStatusOK, host.StartVM())
   
  // Call the OnPluginStart callback in our WASM to read config
  require.Equal(t, host.StartPlugin(), types.OnPluginStartStatusOK)
  
  // Initialize http context to allow us to trigger callbacks like OnHttpRequestHeaders.
  return host, host.InitializeHttpContext(), reset
}

// Helper function if we don't want to override a default PluginConfig
func NewContext(t *testing.T, vmContext proxytest.WasmVMContext) (proxytest.HostEmulator, uint32, func()) {
  return NewContextWithConfig(t, vmContext, DefaultTestConfig)
}

Let’s create a simple test to illustrate how this is used:

package internal

import (
 "testing"

 "github.com/stretchr/testify/require"
 "github.com/tetratelabs/proxy-wasm-go-sdk/proxywasm/types"
)

func TestRequestContext_OnHttpRequestHeaders(t *testing.T) {
  tests := []struct {
    name         string
    headers      [][2]string
    wantResponse types.Action
    wantAuthCall bool
  }{
    {
      name:         "Empty headers",
      headers:      [][2]string{},
      wantAuthCall: false,
      wantResponse: types.ActionContinue,
    },
    {
      name:         "No auth headers",
      headers:      [][2]string{{XRequestIdHeader, "abc"}},
      wantAuthCall: false,
      wantResponse: types.ActionContinue,
    },
    {
      name:         "authorization header present",
      headers:      [][2]string{{XRequestIdHeader, "abc"}, {AuthHeader, "MAC <some mac digest>"}},
      wantAuthCall: true,
      wantResponse: types.ActionPause,
    },
  }
  
  // Load the WASM binary and initialize a bare state for all the proxywasm APIs to work
  vmContext := InitPlugin(t)

  for _, tt := range tests {
    t.Run(tt.name, func(t *testing.T) {
      // Initialize new plugin context like Envoy would do before intercepting a request
      host, contextID, reset := NewContext(t, vmContext=)
      
      // We want to reset the state of our 'Envoy Host' after every test
      defer reset()
      
      // Instruct the 'Envoy Host' to call our WASM OnHttpRequestHeaders callback
      action := host.CallOnRequestHeaders(contextID, tt.headers, true)

      // Now we just validate that all the side-effects match expectations
      require.Equal(t, tt.wantResponse, action)
      if tt.wantAuthCall {
        // Verify auth service is called.
        require.Len(t, host.GetCalloutAttributesFromContext(contextID), 1)
      } else {
        require.Empty(t, host.GetCalloutAttributesFromContext(contextID))
      }
    })
  }
}

We only needed to InitPlugin once as plugins are meant to be re-usable. And then for each test we need to initialize a new HttpContext so that the mocked Envoy Host can store state (including headers, body, paused action, …) for each request we want to test. Then, we want to call the returned reset function to free up that http context as we’re done with it.

Just remember we have to compile the WASM binary before running the tests:

$ tinygo build -o main.wasm -scheduler=none -target=wasi ./main.go

$ go test ./internal/...
ok   envoyfilter/internal 0.348s

To get a bit more complicated let’s add a test for our Auth Server:

package internal

import (
 "testing"

 "github.com/stretchr/testify/require"
 "github.com/tetratelabs/proxy-wasm-go-sdk/proxywasm/types"
)

func TestAuthClient_RequestJWT(t *testing.T) {
  // Initialize new plugin with http context
  host, contextID, reset := NewContext(t, InitPlugin(t))
  defer reset()
  
  // Call OnRequestHeaders to initialize context
  _ = host.CallOnRequestHeaders(contextID, [][2]string{{AuthHeader, "MAC ts....."}}, true)
  
  // Verify Auth Service is called with DispatchHttpCall by checking the length of the callout array
  require.Equal(t, 1, len(host.GetCalloutAttributesFromContext(contextID)))
  
  // At this point, none of dispatched callouts received response therefore the current status must be paused.
  require.Equal(t, types.ActionPause, host.GetCurrentHttpStreamAction(contextID))
  
  // Get the handle to the DispatchHttpCall request
  callout := host.GetCalloutAttributesFromContext(contextID)[0]
  
  // Validate we sent the Auth Service the right headers
  require.Equal(t, [][2]string{
    {"accept", "*/*"},
    {":authority", "auth"},
    {":method", "GET"},
    {":path", "/base64/RkFLRV9KV1QK"},
    {AuthHeader, "MAC ts....."},
  }, callout.Headers)
  
  // Now have a pretend Auth Service respond to our callback handler. Note we can specify any response headers or bodies.
  host.CallOnHttpCallResponse(callout.CalloutID, [][2]string{{":status", "200"}}, nil, []byte("test JWT"))
  
  // Verify the JWT from the above request was added to the original request's headers
  require.Equal(t, [][2]string{
    {AuthHeader, "MAC ts....."},
    {"x-auth-jwt", "test JWT"},
  }, host.GetCurrentRequestHeaders(contextID))
  
  // The request should now have been marked as continued after processing Auth response
  require.Equal(t, types.ActionContinue, host.GetCurrentHttpStreamAction(contextID))
}

Hopefully the comments make each line clear, but the gist is:

• Trigger a callout to our auth-service by getting the Envoy Host to call our OnRequestHeaders hook
• Get the associated callout belonging to the call to DispatchHttpCall we triggered.
• Use that callout to verify payload we sent to an Auth Service, and also to respond with arbitrary data using CallOnHttpCallResponse.
• Verify any changes we made to the intercepted request headers with GetCurrentRequestHeaders.
• Ensure we unblocked the intercepted request by ensuring the state returned by GetCurrentHttpStreamAction is ActionContinue.

Have a look through the linked repo’s code for examples on testing the Metrics and Config objects.

If a test fails, we also helpfully get all the proxywasm debug logging output. For example I intentionally broke the TestAuthClient_RequestJWT test:

$ go test ./internal/...
2023/06/05 14:13:39 proxy_info_log: NewPluginContext context:1
2023/06/05 14:13:39 proxy_info_log: Getting WASM plugin config...
2023/06/05 14:13:39 proxy_info_log: Config loaded: &{AuthClusterName:auth AuthAuthority:auth AuthTimeout:5 Namespace:example-service}
2023/06/05 14:13:39 proxy_info_log: WASM plugin Handling request
2023/06/05 14:13:39 proxy_info_log:   : request header --> authorization: MAC ts.....
2023/06/05 14:13:39 proxy_info_log: : Requesting JWT from Auth Service
2023/06/05 14:13:39 [http callout to auth] timeout: 5
2023/06/05 14:13:39 [http callout to auth] headers: [[accept */*] [:authority auth] [:method GET] [:path /base64/RkFLRV9KV1QK] [authorization MAC ts.....]]
2023/06/05 14:13:39 [http callout to auth] body:
2023/06/05 14:13:39 [http callout to auth] trailers: []
--- FAIL: TestAuthClient_RequestJWT (0.06s)
    auth_client_test.go:19:
         Error Trace: /Users/shender/Code/writing-an-envoy-wasm-plugin/internal/auth_client_test.go:19
         Error:       Not equal:
                      expected: 0
                      actual  : 1
         Test:        TestAuthClient_RequestJWT

In the next post we’ll see some of the gotchas I ran into when developing the WASM plugin using Go.