1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28
29. 29
30. 30
31. 31
32. 32
33. 33
34. 34
35. 35
36. 36
37. 37
38. 38
39. 39
40. 40
41. 41
42. 42
43. 43
44. 44
45. 45
46. 46
47. 47
48. 48
49. 49
50. 50
51. 51
52. 52
53. 53
54. 54
55. 55
56. 56
57. 57
58. 58
59. 59
60. 60
61. 61
62. 62
63. 63
64. 64
65. 65
66. 66
67. 67
68. 68
69. 69
70. 70
71. 71
72. 72
73. 73
74. 74
75. 75
76. 76
77. 77
78. 78
79. 79
80. 80
81. 81
82. 82
83. 83
84. 84
85. 85
86. 86
87. 87
88. 88
89. 89
90. 90
91. 91
92. 92
93. 93
94. 94
95. 95
96. 96
97. 97
98. 98
99. 99
100. 100
101. 101
102. 102
103. 103
104. 104
105. 105
106. 106
107. 107
108. 108
109. 109
110. 110
111. 111
112. 112
113. 113
114. 114
115. 115
116. 116
117. 117
118. 118
119. 119
120. 120
121. 121
122. 122
123. 123
124. 124
125. 125
126. 126
127. 127
128. 128
129. 129
130. 130
131. 131
132. 132
133. 133
134. 134
135. 135
136. 136
137. 137
138. 138
139. 139
140. 140
141. 141
142. 142
143. 143
144. 144
145. 145
146. 146
147. 147
148. 148
149. 149
150. 150
151. 151
152. 152
153. 153
154. 154
155. 155
156. 156
157. 157
158. 158

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

package src

import (
	"bytes"
	"context"
	"fmt"

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

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

	corePtr uint64
}

func Init(ctx context.Context, wasm []byte) (*Core, error) {
	runtime := wazero.NewRuntime(ctx)

	mod, err := runtime.Instantiate(ctx, wasm)
	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
}