Add wasm helper, wasm test utility

.drone.yml

@@ -9,15 +9,33 @@ environment:
   GOARCH: amd64
 
 steps:
-  - name: build-go1.22-alpine3.20
+  - name: build
     image: golang:1.22-alpine3.20
-    environment:
-      GOOS: linux
-      GOARCH: amd64
     commands:
       - apk update
       - apk add --no-cache git
       - go build -v ./cmd/go-away
+
+  - name: test-wasm-success
+    image: golang:1.22-alpine3.20
+    commands:
+      - >
+        go run ./cmd/test-wasm-runtime -wasm ./embed/challenge/js-pow-sha256/runtime/runtime.wasm
+        -make-challenge ./embed/challenge/js-pow-sha256/test/make-challenge.json
+        -make-challenge-out ./embed/challenge/js-pow-sha256/test/make-challenge-out.json
+        -verify-challenge ./embed/challenge/js-pow-sha256/test/verify-challenge.json
+        -verify-challenge-out 0
+
+  - name: test-wasm-fail
+    image: golang:1.22-alpine3.20
+    commands:
+      - >
+        go run ./cmd/test-wasm-runtime -wasm ./embed/challenge/js-pow-sha256/runtime/runtime.wasm
+        -make-challenge ./embed/challenge/js-pow-sha256/test/make-challenge.json
+        -make-challenge-out ./embed/challenge/js-pow-sha256/test/make-challenge-out.json
+        -verify-challenge ./embed/challenge/js-pow-sha256/test/verify-challenge-fail.json
+        -verify-challenge-out 1
+
 ---
 kind: pipeline
 type: docker
@@ -29,7 +47,7 @@ environment:
   GOARCH: amd64
 
 steps:
-  - name: build-go1.24-alpine3.21
+  - name: build
     image: golang:1.24-alpine3.21
     commands:
       - apk update

cmd/test-wasm-runtime/main.go

@@ -0,0 +1,107 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"git.gammaspectra.live/git/go-away/lib/challenge"
+	"github.com/tetratelabs/wazero/api"
+	"os"
+	"reflect"
+	"slices"
+)
+
+func main() {
+
+	pathToTest := flag.String("wasm", "", "Path to test file")
+	makeChallenge := flag.String("make-challenge", "", "Path to contents for MakeChallenge input")
+	makeChallengeOutput := flag.String("make-challenge-out", "", "Path to contents for expected MakeChallenge output")
+	verifyChallenge := flag.String("verify-challenge", "", "Path to contents for VerifyChallenge input")
+	verifyChallengeOutput := flag.Uint64("verify-challenge-out", uint64(challenge.VerifyChallengeOutputOK), "Path to contents for expected VerifyChallenge output")
+
+	flag.Parse()
+
+	if *pathToTest == "" || *makeChallenge == "" || *makeChallengeOutput == "" || *verifyChallenge == "" {
+		flag.PrintDefaults()
+		os.Exit(1)
+	}
+
+	wasmData, err := os.ReadFile(*pathToTest)
+	if err != nil {
+		panic(err)
+	}
+
+	runner := challenge.NewRunner(true)
+	defer runner.Close()
+
+	err = runner.Compile("test", wasmData)
+	if err != nil {
+		panic(err)
+	}
+
+	makeData, err := os.ReadFile(*makeChallenge)
+	if err != nil {
+		panic(err)
+	}
+	var makeIn challenge.MakeChallengeInput
+	err = json.Unmarshal(makeData, &makeIn)
+	if err != nil {
+		panic(err)
+	}
+
+	makeOutData, err := os.ReadFile(*makeChallengeOutput)
+	if err != nil {
+		panic(err)
+	}
+	var makeOut challenge.MakeChallengeOutput
+	err = json.Unmarshal(makeOutData, &makeOut)
+	if err != nil {
+		panic(err)
+	}
+
+	verifyData, err := os.ReadFile(*verifyChallenge)
+	if err != nil {
+		panic(err)
+	}
+	var verifyIn challenge.VerifyChallengeInput
+	err = json.Unmarshal(verifyData, &verifyIn)
+	if err != nil {
+		panic(err)
+	}
+
+	if slices.Compare(makeIn.Key, verifyIn.Key) != 0 {
+		panic("challenge keys do not match")
+	}
+
+	err = runner.Instantiate("test", func(ctx context.Context, mod api.Module) error {
+		out, err := challenge.MakeChallengeCall(ctx, mod, makeIn)
+		if err != nil {
+			return err
+		}
+
+		if !reflect.DeepEqual(*out, makeOut) {
+			return fmt.Errorf("challenge output did not match expected output, got %v, expected %v", *out, makeOut)
+		}
+		return nil
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	err = runner.Instantiate("test", func(ctx context.Context, mod api.Module) error {
+		out, err := challenge.VerifyChallengeCall(ctx, mod, verifyIn)
+		if err != nil {
+			return err
+		}
+
+		if out != challenge.VerifyChallengeOutput(*verifyChallengeOutput) {
+			return fmt.Errorf("verify output did not match expected output, got %d expected %d", out, challenge.VerifyChallengeOutput(*verifyChallengeOutput))
+		}
+		return nil
+	})
+	if err != nil {
+		panic(err)
+	}
+
+}

embed/challenge/js-pow-sha256/test/make-challenge-out.json

@@ -0,0 +1,10 @@
+{
+  "Code": 200,
+  "Data": "eyJjaGFsbGVuZ2UiOiAiMzc5NjRjZGI4NmM4YjRmZTI0YWU5NjU0Y2Q0MGIwZDZkNjgyOTZjZTg0NmY0NTc1YmE0MTRlNDFiMDdmMjg5MyIsICJ0YXJnZXQiOiAiMDAwMDBmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZiIsICJkaWZmaWN1bHR5IjogMjB9",
+  "Error": "",
+  "Headers": {
+    "Content-Type": [
+      "application/json; charset=utf-8"
+    ]
+  }
+}

embed/challenge/js-pow-sha256/test/make-challenge.json

@@ -0,0 +1,60 @@
+{
+  "Data": "",
+  "Headers": {
+    "Accept": [
+      "*/*"
+    ],
+    "Accept-Encoding": [
+      "gzip, deflate, br, zstd"
+    ],
+    "Accept-Language": [
+      "en-US,en;q=0.9"
+    ],
+    "Cache-Control": [
+      "no-cache"
+    ],
+    "Connection": [
+      "keep-alive"
+    ],
+    "Content-Length": [
+      "0"
+    ],
+    "Dnt": [
+      "1"
+    ],
+    "Origin": [
+      "http://127.0.0.1:8787"
+    ],
+    "Pragma": [
+      "no-cache"
+    ],
+    "Sec-Ch-Ua": [
+      "\"Google Chrome\";v=\"135\", \"Not-A.Brand\";v=\"8\", \"Chromium\";v=\"135\""
+    ],
+    "Sec-Ch-Ua-Mobile": [
+      "?0"
+    ],
+    "Sec-Ch-Ua-Platform": [
+      "\"Linux\""
+    ],
+    "Sec-Fetch-Dest": [
+      "empty"
+    ],
+    "Sec-Fetch-Mode": [
+      "cors"
+    ],
+    "Sec-Fetch-Site": [
+      "same-origin"
+    ],
+    "User-Agent": [
+      "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36"
+    ],
+    "X-Away-Id": [
+      "cf33e115f699c50822c4e56ed3c610dc"
+    ]
+  },
+  "Key": "Pl02g55pPapXdVc3SVfMZQGymmyE0dTCpq0qm8ax9ss=",
+  "Parameters": {
+    "difficulty": "20"
+  }
+}

embed/challenge/js-pow-sha256/test/verify-challenge-fail.json

@@ -0,0 +1,7 @@
+{
+  "Key":"Pl02g55pPapXdVc3SVfMZQGymmyE0dTCpq0qm8ax9ss=",
+  "Parameters":{
+    "difficulty":"20"
+  },
+  "Result":"Mzc5NjRjZGI4NmM4YjRmZTI0YWU5NjU0Y2Q0MGIwZDZkNjgyOTZjZTg0NmY0NTc1YmE0MTRlNDFiMDdmMjg5MzY4Y2QxMTAwMDAwMDAwMDE="
+}

embed/challenge/js-pow-sha256/test/verify-challenge.json

@@ -0,0 +1,7 @@
+{
+  "Key":"Pl02g55pPapXdVc3SVfMZQGymmyE0dTCpq0qm8ax9ss=",
+  "Parameters":{
+    "difficulty":"20"
+  },
+  "Result":"Mzc5NjRjZGI4NmM4YjRmZTI0YWU5NjU0Y2Q0MGIwZDZkNjgyOTZjZTg0NmY0NTc1YmE0MTRlNDFiMDdmMjg5MzY4Y2QxMTAwMDAwMDAwMDA="
+}

examples/forgejo.yml

@@ -1,5 +1,5 @@
 # Example cmdline (forward requests from upstream to port :8080)
-# $ go-away --bind :8080 --backend git.example.com:http://forgejo:3000 --policy forgejo.yml --challenge-template forgejo --challenge-template-theme forgejo-dark
+# $ go-away --bind :8080 --backend git.example.com:http://forgejo:3000 --policy examples/forgejo.yml --challenge-template forgejo --challenge-template-theme forgejo-dark
 
 
 

lib/challenge.go

@@ -2,15 +2,12 @@ package lib
 
 import (
 	"bytes"
-	"context"
 	"crypto/sha256"
 	"encoding/binary"
 	"errors"
 	"fmt"
 	"github.com/go-jose/go-jose/v4"
 	"github.com/go-jose/go-jose/v4/jwt"
-	"github.com/tetratelabs/wazero"
-	"github.com/tetratelabs/wazero/api"
 	"math/rand/v2"
 	"net"
 	"net/http"
@@ -171,29 +168,3 @@ func (state *State) VerifyChallengeToken(name string, expectedKey []byte, w http
 
 	return true, nil
 }
-
-func (state *State) ChallengeMod(name string, cb func(ctx context.Context, mod api.Module) error) error {
-	c, ok := state.Challenges[name]
-	if !ok {
-		return errors.New("challenge not found")
-	}
-	if c.RuntimeModule == nil {
-		return errors.New("challenge module is nil")
-	}
-
-	ctx := state.WasmContext
-	mod, err := state.WasmRuntime.InstantiateModule(
-		ctx,
-		c.RuntimeModule,
-		wazero.NewModuleConfig().WithName(name).WithStartFunctions("_initialize"),
-	)
-	if err != nil {
-		return err
-	}
-	defer mod.Close(ctx)
-	err = cb(ctx, mod)
-	if err != nil {
-		return err
-	}
-	return nil
-}

lib/challenge/generic.go

@@ -1,74 +0,0 @@
-//go:build !tinygo
-
-package challenge
-
-import (
-	"context"
-	"encoding/json"
-	"errors"
-	"github.com/tetratelabs/wazero/api"
-)
-
-func MakeChallengeCall(ctx context.Context, mod api.Module, in MakeChallengeInput) (*MakeChallengeOutput, error) {
-	makeChallengeFunc := mod.ExportedFunction("MakeChallenge")
-	malloc := mod.ExportedFunction("malloc")
-	free := mod.ExportedFunction("free")
-
-	inData, err := json.Marshal(in)
-
-	mallocResult, err := malloc.Call(ctx, uint64(len(inData)))
-	if err != nil {
-		return nil, err
-	}
-	defer free.Call(ctx, mallocResult[0])
-	if !mod.Memory().Write(uint32(mallocResult[0]), inData) {
-		return nil, errors.New("could not write memory")
-	}
-	result, err := makeChallengeFunc.Call(ctx, uint64(NewAllocation(uint32(mallocResult[0]), uint32(len(inData)))))
-	if err != nil {
-		return nil, err
-	}
-	resultPtr := Allocation(result[0])
-	outData, ok := mod.Memory().Read(resultPtr.Pointer(), resultPtr.Size())
-	if !ok {
-		return nil, errors.New("could not read result")
-	}
-	defer free.Call(ctx, uint64(resultPtr.Pointer()))
-
-	var out MakeChallengeOutput
-	err = json.Unmarshal(outData, &out)
-	if err != nil {
-		return nil, err
-	}
-	return &out, nil
-}
-
-func VerifyChallengeCall(ctx context.Context, mod api.Module, in VerifyChallengeInput) (VerifyChallengeOutput, error) {
-	verifyChallengeFunc := mod.ExportedFunction("VerifyChallenge")
-	malloc := mod.ExportedFunction("malloc")
-	free := mod.ExportedFunction("free")
-
-	inData, err := json.Marshal(in)
-
-	mallocResult, err := malloc.Call(ctx, uint64(len(inData)))
-	if err != nil {
-		return VerifyChallengeOutputError, err
-	}
-	defer free.Call(ctx, mallocResult[0])
-	if !mod.Memory().Write(uint32(mallocResult[0]), inData) {
-		return VerifyChallengeOutputError, errors.New("could not write memory")
-	}
-	result, err := verifyChallengeFunc.Call(ctx, uint64(NewAllocation(uint32(mallocResult[0]), uint32(len(inData)))))
-	if err != nil {
-		return VerifyChallengeOutputError, err
-	}
-
-	return VerifyChallengeOutput(result[0]), nil
-}
-
-func PtrToBytes(ptr uint32, size uint32) []byte   { panic("not implemented") }
-func BytesToPtr(s []byte) (uint32, uint32)        { panic("not implemented") }
-func BytesToLeakedPtr(s []byte) (uint32, uint32)  { panic("not implemented") }
-func PtrToString(ptr uint32, size uint32) string  { panic("not implemented") }
-func StringToPtr(s string) (uint32, uint32)       { panic("not implemented") }
-func StringToLeakedPtr(s string) (uint32, uint32) { panic("not implemented") }

lib/challenge/interface.go

@@ -5,8 +5,7 @@ import (
 	"git.gammaspectra.live/git/go-away/utils/inline"
 )
 
-type MakeChallenge func(in Allocation) (out Allocation)
+// Allocation is a combination of pointer location in WASM memory and size of it
-
 type Allocation uint64
 
 func NewAllocation(ptr, size uint32) Allocation {

lib/challenge/utils_generic.go

@@ -0,0 +1,195 @@
+//go:build !tinygo || !wasip1
+
+package challenge
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
+	"slices"
+)
+
+type Runner struct {
+	context context.Context
+	runtime wazero.Runtime
+
+	modules map[string]wazero.CompiledModule
+}
+
+func NewRunner(useNativeCompiler bool) *Runner {
+	var r Runner
+	r.context = context.Background()
+	var runtimeConfig wazero.RuntimeConfig
+	if useNativeCompiler {
+		runtimeConfig = wazero.NewRuntimeConfigCompiler()
+	} else {
+		runtimeConfig = wazero.NewRuntimeConfigInterpreter()
+	}
+	r.runtime = wazero.NewRuntimeWithConfig(r.context, runtimeConfig)
+	wasi_snapshot_preview1.MustInstantiate(r.context, r.runtime)
+
+	r.modules = make(map[string]wazero.CompiledModule)
+
+	return &r
+}
+
+func (r *Runner) Compile(key string, binary []byte) error {
+	module, err := r.runtime.CompileModule(r.context, binary)
+	if err != nil {
+		return err
+	}
+
+	// check interface
+	functions := module.ExportedFunctions()
+	if f, ok := functions["MakeChallenge"]; ok {
+		if slices.Compare(f.ParamTypes(), []api.ValueType{api.ValueTypeI64}) != 0 {
+			return fmt.Errorf("MakeChallenge does not follow parameter interface")
+		}
+		if slices.Compare(f.ResultTypes(), []api.ValueType{api.ValueTypeI64}) != 0 {
+			return fmt.Errorf("MakeChallenge does not follow result interface")
+		}
+	} else {
+		module.Close(r.context)
+		return errors.New("no MakeChallenge exported")
+	}
+
+	if f, ok := functions["VerifyChallenge"]; ok {
+		if slices.Compare(f.ParamTypes(), []api.ValueType{api.ValueTypeI64}) != 0 {
+			return fmt.Errorf("VerifyChallenge does not follow parameter interface")
+		}
+		if slices.Compare(f.ResultTypes(), []api.ValueType{api.ValueTypeI64}) != 0 {
+			return fmt.Errorf("VerifyChallenge does not follow result interface")
+		}
+	} else {
+		module.Close(r.context)
+		return errors.New("no VerifyChallenge exported")
+	}
+
+	if f, ok := functions["malloc"]; ok {
+		if slices.Compare(f.ParamTypes(), []api.ValueType{api.ValueTypeI32}) != 0 {
+			return fmt.Errorf("malloc does not follow parameter interface")
+		}
+		if slices.Compare(f.ResultTypes(), []api.ValueType{api.ValueTypeI32}) != 0 {
+			return fmt.Errorf("malloc does not follow result interface")
+		}
+	} else {
+		module.Close(r.context)
+		return errors.New("no malloc exported")
+	}
+
+	if f, ok := functions["free"]; ok {
+		if slices.Compare(f.ParamTypes(), []api.ValueType{api.ValueTypeI32}) != 0 {
+			return fmt.Errorf("free does not follow parameter interface")
+		}
+		if slices.Compare(f.ResultTypes(), []api.ValueType{}) != 0 {
+			return fmt.Errorf("free does not follow result interface")
+		}
+	} else {
+		module.Close(r.context)
+		return errors.New("no free exported")
+	}
+
+	r.modules[key] = module
+	return nil
+}
+
+func (r *Runner) Close() {
+	for _, module := range r.modules {
+		module.Close(r.context)
+	}
+	r.runtime.Close(r.context)
+}
+
+var ErrModuleNotFound = errors.New("module not found")
+
+func (r *Runner) Instantiate(key string, f func(ctx context.Context, mod api.Module) error) (err error) {
+	compiledModule, ok := r.modules[key]
+	if !ok {
+		return ErrModuleNotFound
+	}
+	mod, err := r.runtime.InstantiateModule(
+		r.context,
+		compiledModule,
+		wazero.NewModuleConfig().WithName(key).WithStartFunctions("_initialize"),
+	)
+	if err != nil {
+		return err
+	}
+	defer mod.Close(r.context)
+
+	return f(r.context, mod)
+}
+
+func MakeChallengeCall(ctx context.Context, mod api.Module, in MakeChallengeInput) (*MakeChallengeOutput, error) {
+	makeChallengeFunc := mod.ExportedFunction("MakeChallenge")
+	malloc := mod.ExportedFunction("malloc")
+	free := mod.ExportedFunction("free")
+
+	inData, err := json.Marshal(in)
+	if err != nil {
+		return nil, err
+	}
+
+	mallocResult, err := malloc.Call(ctx, uint64(len(inData)))
+	if err != nil {
+		return nil, err
+	}
+	defer free.Call(ctx, mallocResult[0])
+	if !mod.Memory().Write(uint32(mallocResult[0]), inData) {
+		return nil, errors.New("could not write memory")
+	}
+	result, err := makeChallengeFunc.Call(ctx, uint64(NewAllocation(uint32(mallocResult[0]), uint32(len(inData)))))
+	if err != nil {
+		return nil, err
+	}
+	resultPtr := Allocation(result[0])
+	outData, ok := mod.Memory().Read(resultPtr.Pointer(), resultPtr.Size())
+	if !ok {
+		return nil, errors.New("could not read result")
+	}
+	defer free.Call(ctx, uint64(resultPtr.Pointer()))
+
+	var out MakeChallengeOutput
+	err = json.Unmarshal(outData, &out)
+	if err != nil {
+		return nil, err
+	}
+	return &out, nil
+}
+
+func VerifyChallengeCall(ctx context.Context, mod api.Module, in VerifyChallengeInput) (VerifyChallengeOutput, error) {
+	verifyChallengeFunc := mod.ExportedFunction("VerifyChallenge")
+	malloc := mod.ExportedFunction("malloc")
+	free := mod.ExportedFunction("free")
+
+	inData, err := json.Marshal(in)
+	if err != nil {
+		return VerifyChallengeOutputError, err
+	}
+
+	mallocResult, err := malloc.Call(ctx, uint64(len(inData)))
+	if err != nil {
+		return VerifyChallengeOutputError, err
+	}
+	defer free.Call(ctx, mallocResult[0])
+	if !mod.Memory().Write(uint32(mallocResult[0]), inData) {
+		return VerifyChallengeOutputError, errors.New("could not write memory")
+	}
+	result, err := verifyChallengeFunc.Call(ctx, uint64(NewAllocation(uint32(mallocResult[0]), uint32(len(inData)))))
+	if err != nil {
+		return VerifyChallengeOutputError, err
+	}
+
+	return VerifyChallengeOutput(result[0]), nil
+}
+
+func PtrToBytes(ptr uint32, size uint32) []byte   { panic("not implemented") }
+func BytesToPtr(s []byte) (uint32, uint32)        { panic("not implemented") }
+func BytesToLeakedPtr(s []byte) (uint32, uint32)  { panic("not implemented") }
+func PtrToString(ptr uint32, size uint32) string  { panic("not implemented") }
+func StringToPtr(s string) (uint32, uint32)       { panic("not implemented") }
+func StringToLeakedPtr(s string) (uint32, uint32) { panic("not implemented") }

lib/state.go

@@ -12,16 +12,14 @@ import (
 	"errors"
 	"fmt"
 	"git.gammaspectra.live/git/go-away/embed"
-	challenge2 "git.gammaspectra.live/git/go-away/lib/challenge"
+	"git.gammaspectra.live/git/go-away/lib/challenge"
 	"git.gammaspectra.live/git/go-away/lib/condition"
 	"git.gammaspectra.live/git/go-away/lib/policy"
 	"git.gammaspectra.live/git/go-away/utils/inline"
 	"github.com/google/cel-go/cel"
 	"github.com/google/cel-go/common/types"
 	"github.com/google/cel-go/common/types/ref"
-	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
-	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
 	"github.com/yl2chen/cidranger"
 	"html/template"
 	"io"
@@ -46,8 +44,7 @@ type State struct {
 
 	Networks map[string]cidranger.Ranger
 
-	WasmRuntime wazero.Runtime
+	Wasm *challenge.Runner
-	WasmContext context.Context
 
 	Challenges map[string]ChallengeState
 
@@ -84,8 +81,6 @@ const (
 )
 
 type ChallengeState struct {
-	RuntimeModule wazero.CompiledModule
-
 	Path string
 
 	Static              http.Handler
@@ -191,9 +186,7 @@ func NewState(p policy.Policy, settings StateSettings) (state *State, err error)
 		state.Networks[k] = ranger
 	}
 
-	state.WasmContext = context.Background()
+	state.Wasm = challenge.NewRunner(true)
-	state.WasmRuntime = wazero.NewRuntimeWithConfig(state.WasmContext, wazero.NewRuntimeConfigCompiler())
-	wasi_snapshot_preview1.MustInstantiate(state.WasmContext, state.WasmRuntime)
 
 	state.Challenges = make(map[string]ChallengeState)
 
@@ -475,15 +468,15 @@ func NewState(p policy.Policy, settings StateSettings) (state *State, err error)
 			if err != nil {
 				return nil, fmt.Errorf("c %s: could not load runtime: %w", challengeName, err)
 			}
-			c.RuntimeModule, err = state.WasmRuntime.CompileModule(state.WasmContext, wasmData)
+			err = state.Wasm.Compile(challengeName, wasmData)
 			if err != nil {
 				return nil, fmt.Errorf("c %s: compiling runtime: %w", challengeName, err)
 			}
 
 			c.MakeChallenge = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				err := state.ChallengeMod(challengeName, func(ctx context.Context, mod api.Module) (err error) {
+				err := state.Wasm.Instantiate(challengeName, func(ctx context.Context, mod api.Module) (err error) {
 
-					in := challenge2.MakeChallengeInput{
+					in := challenge.MakeChallengeInput{
 						Key:        state.GetChallengeKeyForRequest(challengeName, time.Now().UTC().Add(DefaultValidity).Round(DefaultValidity), r),
 						Parameters: p.Parameters,
 						Headers:    inline.MIMEHeader(r.Header),
@@ -493,7 +486,7 @@ func NewState(p policy.Policy, settings StateSettings) (state *State, err error)
 						return err
 					}
 
-					out, err := challenge2.MakeChallengeCall(state.WasmContext, mod, in)
+					out, err := challenge.MakeChallengeCall(ctx, mod, in)
 					if err != nil {
 						return err
 					}
@@ -514,22 +507,22 @@ func NewState(p policy.Policy, settings StateSettings) (state *State, err error)
 			})
 
 			c.Verify = func(key []byte, result string) (ok bool, err error) {
-				err = state.ChallengeMod(challengeName, func(ctx context.Context, mod api.Module) (err error) {
+				err = state.Wasm.Instantiate(challengeName, func(ctx context.Context, mod api.Module) (err error) {
-					in := challenge2.VerifyChallengeInput{
+					in := challenge.VerifyChallengeInput{
 						Key:        key,
 						Parameters: p.Parameters,
 						Result:     []byte(result),
 					}
 
-					out, err := challenge2.VerifyChallengeCall(state.WasmContext, mod, in)
+					out, err := challenge.VerifyChallengeCall(ctx, mod, in)
 					if err != nil {
 						return err
 					}
 
-					if out == challenge2.VerifyChallengeOutputError {
+					if out == challenge.VerifyChallengeOutputError {
 						return errors.New("error checking challenge")
 					}
-					ok = out == challenge2.VerifyChallengeOutputOK
+					ok = out == challenge.VerifyChallengeOutputOK
 					return nil
 				})
 				if err != nil {