/packages/backend/server.go at 75d721c8a12c754cdea30fcfedca8eaad288af6a

// SPDX-FileCopyrightText: 2025 Shota FUJI <pockawoooh@gmail.com>
// SPDX-License-Identifier: AGPL-3.0-only

package main

import (
	"bytes"
	"context"
	_ "embed"
	"fmt"
	"log"

	"github.com/tetratelabs/wazero"
	"github.com/tetratelabs/wazero/api"

	_ "pocka.jp/x/yamori/proto/go/meta/v1"
)

//go:embed zig-out/bin/yamori_backend.wasm
var coreWasm []byte

func main() {
	ctx := context.Background()

	core, err := initCore(ctx)
	if err != nil {
		log.Panicln(err)
	}
	defer core.Close(ctx)

	resp, err := core.Handle(ctx, "yamori.meta.v1.MetaService", "Ping", []byte{})
	if err != nil {
		log.Panicln(err)
	}

	log.Printf("%v\n", resp)
}

// Core は WASM で実装されており、特にメモリアクセスでかなり煩雑になる。
// この構造体はそれらの操作を抽象化する。
// 内部実装は基本的に JavaScript 実装 (worker.ts) のポートとなっている。
type Core struct {
	wasmRuntime wazero.Runtime
	mod         api.Module

	corePtr uint64
}

func initCore(ctx context.Context) (*Core, error) {
	runtime := wazero.NewRuntime(ctx)

	mod, err := runtime.Instantiate(ctx, coreWasm)
	if err != nil {
		return nil, err
	}

	initCore := mod.ExportedFunction("init_core")

	initCoreResults, err := initCore.Call(ctx)
	if err != nil {
		return nil, err
	}

	return &Core{
		wasmRuntime: runtime,
		mod:         mod,
		corePtr:     initCoreResults[0],
	}, nil
}

// allocate は WASM の "allocate_bytes" を呼び、確保されたアドレスを返す。
func (core Core) allocate(ctx context.Context, length uint32) (uint32, error) {
	fn := core.mod.ExportedFunction("allocate_bytes")

	results, err := fn.Call(ctx, uint64(length))
	if err != nil {
		return 0, err
	}

	return uint32(results[0]), nil
}

// copy は受け取ったバイト列を WASM のメモリ上にコピーして WASM から読めるようにする。
// 確保したアドレスとバイト長を返す。
func (core Core) copy(ctx context.Context, bytes []byte) (uint32, uint32, error) {
	length := uint32(len(bytes))

	ptr, err := core.allocate(ctx, length)
	if err != nil {
		return 0, 0, err
	}

	if length == 0 {
		return ptr, length, nil
	}

	if !core.mod.Memory().Write(ptr, bytes) {
		return 0, 0, fmt.Errorf("Failed to write to memory (%d bytes)", length)
	}

	return ptr, length, nil
}

// Handle は RPC のリクエストを受けてレスポンスを返す関数。
// サービス名もメソッド名も変更しないため、レスポンスのバイナリデータのみを返す。
func (core Core) Handle(ctx context.Context, service string, method string, data []byte) ([]byte, error) {
	freeBytes := core.mod.ExportedFunction("free_bytes")

	servicePtr, serviceLen, err := core.copy(ctx, []byte(service))
	if err != nil {
		return nil, err
	}
	defer freeBytes.Call(ctx, uint64(servicePtr), uint64(serviceLen))

	methodPtr, methodLen, err := core.copy(ctx, []byte(method))
	if err != nil {
		return nil, err
	}
	defer freeBytes.Call(ctx, uint64(methodPtr), uint64(methodLen))

	dataPtr, dataLen, err := core.copy(ctx, data)
	if err != nil {
		return nil, err
	}
	defer freeBytes.Call(ctx, uint64(dataPtr), uint64(dataLen))

	destroyResponse := core.mod.ExportedFunction("destroy_response")

	createResponse := core.mod.ExportedFunction("create_response")
	createResponseResult, err := createResponse.Call(ctx)
	if err != nil {
		return nil, err
	}
	defer destroyResponse.Call(ctx, createResponseResult[0])

	handleFn := core.mod.ExportedFunction("handle")
	handleResults, err := handleFn.Call(
		ctx, core.corePtr, uint64(servicePtr), uint64(serviceLen),
		uint64(methodPtr), uint64(methodLen), uint64(dataPtr), uint64(dataLen),
		createResponseResult[0],
	)
	if err != nil {
		return nil, err
	}
	if handleResults[0] != 0 {
		return nil, fmt.Errorf("handler returned non-zero result: %d", handleResults[0])
	}

	getResponsePtr := core.mod.ExportedFunction("get_response_ptr")
	getResponseLen := core.mod.ExportedFunction("get_response_len")

	responsePtr, err := getResponsePtr.Call(ctx, createResponseResult[0])
	if err != nil {
		return nil, err
	}

	responseLen, err := getResponseLen.Call(ctx, createResponseResult[0])
	if err != nil {
		return nil, err
	}

	if responseLen[0] == 0 {
		return []byte{}, nil
	}

	responseData, isRangeOK := core.mod.Memory().Read(uint32(responsePtr[0]), uint32(responseLen[0]))
	if !isRangeOK {
		return nil, fmt.Errorf("get_response_ptr/len returned out-of-range address or length")
	}

	// destroy_response によってメモリが解放されるため、返す前にクローンする
	// 必要がある。しないと read after free になってしまう。
	return bytes.Clone(responseData), nil
}

// Close を呼ぶと以降は Core を使った操作はできなくなる。
// WASM ランタイムのリソースが解放される。
func (core *Core) Close(ctx context.Context) {
	core.mod.Close(ctx)
	core.wasmRuntime.Close(ctx)
	core = nil
}