Commit 5809f290 authored by nagayama15's avatar nagayama15

Implement the watermarker using function reordering method

parent 8a10725e
(module
(export "add2" (func $add2))
(export "add3" (func $add3))
(export "add2" (func $add2))
(func $add2 (param $a i32) (param $b i32) (result i32)
(i32.add
(local.get $a)
(local.get $b)
)
)
(import "env" "importfunc" (func $importfunc))
(func $inc (param $a i32) (result i32)
(i32.add
(local.get $a)
(i32.const 1)
)
)
(func $add3 (param $a i32) (param $b i32) (param $c i32) (result i32)
(i32.add
(i32.add
......@@ -16,10 +23,4 @@
(local.get $c)
)
)
(func $inc (param $a i32) (result i32)
(i32.add
(local.get $a)
(i32.const 1)
)
)
)
#include "FunctionOrderingWatermarker.hpp"
#include <algorithm>
#include <array>
#include <cassert>
#include <iostream>
namespace kyut::watermarker {
void embedFunctionOrdering(wasm::Module &module, BitStreamReader &stream) {
(void)module;
(void)stream;
namespace {
/**
* Move the `element` in [begin, end) to the beginning of the range.
* The order of elements other than `element` is maintained.
*/
template <typename RandomAccessIterator>
void moveToFront(RandomAccessIterator begin, RandomAccessIterator end, RandomAccessIterator element) {
assert(begin <= element && element < end);
auto tmp = std::move(*element);
std::move_backward(begin, element, std::next(element));
*begin = std::move(tmp);
}
} // namespace
std::size_t embedFunctionOrdering(wasm::Module &module, BitStreamReader &stream, std::size_t divisions) {
assert(2 <= divisions && divisions < 21 && "because 21! > 2^64");
// Number of bits embedded in the module
std::size_t numBits = 0;
// Split according to the function in the module has body or not
// [begin, start) has no body, and [start, end) has
const auto start = std::partition(
std::begin(module.functions), std::end(module.functions), [](const auto &f) { return f->body == nullptr; });
const size_t count = std::distance(start, std::end(module.functions));
for (size_t i = 0; i < count; i += divisions) {
const auto chunkSize = (std::min)(divisions, count - i);
const auto chunkBegin = start + i;
const auto chunkEnd = chunkBegin + chunkSize;
// Number of bits that can be embedded in the chunk
// {floor(log2(x!)) | x in [0, 21)}
constexpr auto bitLengthTable = std::array<std::size_t, 21>{
0, 0, 1, 2, 4, 6, 9, 12, 15, 18, 21, 25, 28, 32, 36, 40, 44, 48, 52, 56, 61,
};
const auto numBitsEmbeddedInChunk = bitLengthTable[chunkSize];
// Sort functions in the chunk
std::sort(chunkBegin, chunkEnd, [](const auto &a, const auto &b) { return a->name < b->name; });
// A watermark to embed in the chunk
auto watermark = stream.read(numBitsEmbeddedInChunk);
// Reorder the functions
for (auto it = chunkBegin; it != chunkEnd; ++it) {
const auto distance = std::distance(it, chunkEnd);
const auto n = watermark % distance;
watermark /= distance;
moveToFront(it, chunkEnd, it + n);
}
numBits += numBitsEmbeddedInChunk;
}
return numBits;
}
} // namespace kyut::watermarker
......@@ -6,7 +6,7 @@
#include "../BitStreamReader.hpp"
namespace kyut::watermarker {
void embedFunctionOrdering(wasm::Module &module, BitStreamReader &stream);
std::size_t embedFunctionOrdering(wasm::Module &module, BitStreamReader &stream, std::size_t divisions);
}
#endif // INCLUDE_kyut_watermark_FunctionOrderingWatermarker_hpp
......@@ -23,7 +23,7 @@ int main(int argc, char *argv[]) {
// Embed watermarks
const auto stream = kyut::BitStreamReader::fromString(watermark);
kyut::watermarker::embedFunctionOrdering(module, *stream);
kyut::watermarker::embedFunctionOrdering(module, *stream, 10);
// Output the result
wasm::ModuleWriter{}.writeText(module, "");
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment