Modules: new module to filter DNS tunneling queries

meson.build

@@ -167,6 +167,9 @@ malloc = meson.get_compiler('c').find_library(
 ### filter module
 libidn2 = dependency('libidn2', version: '>=2.0.0', required: false)
 libpcre2 = dependency('libpcre2-8', version: '>=10.00', required: false)
+
+### tunnel_filter module
+libtorch = dependency('Torch', version: '>=2.4.1', required: false)
 message('---------------------------')
 
 ## Compiler args

modules/meson.build

@@ -60,6 +60,7 @@ subdir('stats')
 subdir('ta_update')
 subdir('view')
 subdir('filter')
+subdir('tunnel_filter')
 
 # install lua modules
 foreach mod : lua_mod_src

modules/tunnel_filter/README.rst

@@ -0,0 +1,6 @@
+.. SPDX-License-Identifier: GPL-3.0-or-later
+
+.. _mod-tunnel_filter:
+
+Tunnel Filter
+======

modules/tunnel_filter/libblcnn.cpp

@@ -0,0 +1,56 @@
+#include <torch/script.h>
+#include <iostream>
+#include <vector>
+#include <cstring>
+#include <random>
+#include <filesystem>
+#include "libblcnn.h"
+
+constexpr int MAX_PACKET_SIZE = 300;
+constexpr int ONE_HOT_SIZE = 257;
+
+struct TorchModuleWrapper {
+	torch::jit::script::Module model;
+};
+
+TorchModule load_model(const char *model_path) {
+	try {
+		namespace fs = std::filesystem;
+		auto *wrapper = new TorchModuleWrapper();
+
+		fs::path file_path = fs::relative(__FILE__, "../");
+		fs::path absolute_path = fs::absolute(file_path.parent_path()) / model_path;
+		wrapper->model = torch::jit::load(absolute_path);
+		wrapper->model.to(torch::kCPU);
+		wrapper->model.eval();
+
+		return static_cast<TorchModule>(wrapper);
+	} catch (const c10::Error &e) {
+		std::cerr << "Error loading model: " << e.what() << std::endl;
+
+		return nullptr;
+	}
+}
+
+int predict_packet(TorchModule module, const unsigned char *data, size_t size) {
+	if (!module) return -1;
+	auto* wrapper = reinterpret_cast<TorchModuleWrapper*>(module);
+
+	torch::Tensor one_hot = torch::zeros({1, MAX_PACKET_SIZE}, torch::kLong);
+
+	for (size_t i = 0; i < size && i < MAX_PACKET_SIZE; i++) {
+		one_hot[0][i] = data[i];
+	}
+	for (size_t i = size; i < MAX_PACKET_SIZE; i++) {
+		one_hot[0][i] = 256;
+	}
+
+	torch::Tensor output = wrapper->model.forward({one_hot}).toTensor();
+
+	int predicted_class = std::get<1>(torch::max(output, 1)).item<int>();
+	return predicted_class;
+}
+
+void free_model(TorchModule model) {
+	delete static_cast<TorchModuleWrapper*>(model);
+}

modules/tunnel_filter/libblcnn.h

@@ -0,0 +1,18 @@
+#ifndef TORCH_WRAPPER_H
+#define TORCH_WRAPPER_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef void* TorchModule;
+
+TorchModule load_model(const char *model_path);
+int predict_packet(TorchModule model, const unsigned char *data, size_t size);
+void free_model(TorchModule model);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

modules/tunnel_filter/meson.build

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: GPL-3.0-or-later
+# C module: tunnel_filter
+
+tunnel_filter_src = files([
+  'tunnel_filter.c',
+])
+c_src_lint += tunnel_filter_src
+
+if libtorch.found()
+  libtorch_dep = declare_dependency(
+    link_args: ['-ltorch', '-ltorch_cpu'],
+    dependencies: [libtorch]
+  )
+
+
+  blcnn_lib = static_library('blcnn',
+    sources: ['libblcnn.cpp'],
+    dependencies: [libtorch_dep]
+  )
+
+
+  tunnel_filter_mod = shared_module(
+    'tunnel_filter',
+    tunnel_filter_src,
+    dependencies: mod_deps + [
+      libtorch_dep
+    ],
+    include_directories: mod_inc_dir,
+    name_prefix: '',
+    install: true,
+    install_dir: modules_dir,
+    link_with: [blcnn_lib, kresd],
+    link_args: ['-lstdc++', '-lc10']
+  )
+endif

modules/tunnel_filter/tunnel_filter.c

@@ -0,0 +1,84 @@
+/*  Copyright (C) CZ.NIC, z.s.p.o. <knot-resolver@labs.nic.cz>
+ *  SPDX-License-Identifier: GPL-3.0-or-later
+ */
+
+/**
+ * @file tunnel_filter.c
+ * @brief blocks queries that are evaluated as DNS tunneling exfiltration
+ */
+
+#include <stdio.h>
+#include <time.h>
+#include <stdlib.h>
+#include "libblcnn.h"
+#include "lib/layer.h"
+#include "lib/resolve.h"
+
+#define MAX_PACKET_SIZE 300
+
+static int create_exfiltration_answer(kr_layer_t *ctx)
+{
+	struct kr_request *req = ctx->req;
+	knot_pkt_t *answer = kr_request_ensure_answer(req);
+	if (!answer)
+		return ctx->state;
+
+	knot_wire_set_rcode(answer->wire, KNOT_RCODE_NXDOMAIN);
+	knot_wire_clear_ad(answer->wire);
+
+	kr_request_set_extended_error(req, KNOT_EDNS_EDE_BLOCKED,
+			"A4CP: Potential DNS tunnelling exfiltration query");
+	ctx->state = KR_STATE_DONE;
+	return ctx->state;
+}
+
+static int infer(kr_layer_t *ctx)
+{
+	struct kr_module *module = ctx->api->data;
+	TorchModule net = module->data;
+	struct kr_request *req = ctx->req;
+	uint8_t *packet = req->qsource.packet->wire;
+	size_t packet_size = req->qsource.size;
+
+	int result = predict_packet(net, packet, packet_size);
+	if (result)
+		return create_exfiltration_answer(ctx);
+
+	return ctx->state;
+}
+
+KR_EXPORT
+int tunnel_filter_init(struct kr_module *module)
+{
+	static kr_layer_api_t layer = {
+		.begin = &infer,
+	};
+
+	layer.data = module;
+	module->layer = &layer;
+
+	static const struct kr_prop props[] = {
+		{ NULL, NULL, NULL }
+	};
+	module->props = props;
+
+	TorchModule net = load_model("blcnn.pt");
+	if (!net) 
+		return kr_error(ENOMEM);
+
+	module->data = net;
+	return kr_ok();
+}
+
+KR_EXPORT
+int tunnel_filter_deinit(struct kr_module *module)
+{
+	TorchModule net = module->data;
+	if (net) {
+		free_model(net);
+		module->data = NULL;
+	}
+	return kr_ok();
+}
+
+KR_MODULE_EXPORT(tunnel_filter)