LuaSecurity.cpp

#include <fstream>
#include <scripting/ModException.hpp>
#include <sanity.hpp>
#include "LuaSecurity.hpp"

using namespace scripting;

namespace {

void copyAll(sol::environment &env, const sol::global_table &globals,
		const std::vector<std::string> &names) {
	for (const auto &name : names) {
		env[name] = globals[name];
	}
}

sol::table deepCopy(sol::state &lua, const sol::table &table) {
	sol::table table2(lua, sol::create);
	for (auto pair : table) {
		table2[pair.first] = pair.second;
	}
	return table2;
}

void copyTables(sol::environment &env, const sol::global_table &globals,
		sol::state &lua, const std::vector<std::string> &names) {
	for (const auto &name : names) {
		env[name] = deepCopy(lua, globals[name]);
	}
}

} // namespace

void LuaSecurity::buildEnvironment() {
	env = sol::environment(lua, sol::create);
	env["_G"] = env;

	const std::vector<std::string> whitelisted = {
			"assert",
			"error",
			"getmetatable", //< Used to extend string class
			"ipairs",
			"next",
			"pairs",
			"pcall",
			"print",

			// TODO: remove these
			"package",
			"require",

			// Required for implementing classes
			"rawequal",
			"rawget",
			"rawset",

			"select",
			"setmetatable", //< Required for implementing classes
			"tonumber",
			"tostring",
			"type",
			"unpack",
			"_VERSION",
			"xpcall",
	};

	std::vector<std::string> safeLibraries = {
			"coroutine", "string", "table", "math"};

	copyAll(env, lua.globals(), whitelisted);
	copyTables(env, lua.globals(), lua, safeLibraries);

	env.set_function("loadstring", &LuaSecurity::loadstring, this);
	env.set_function("loadfile", &LuaSecurity::loadfile, this);
	env.set_function("dofile", &LuaSecurity::dofile, this);

	sol::table os(lua, sol::create);
	os["clock"] = lua["os"]["clock"];
	os["date"] = lua["os"]["date"];
	os["difftime"] = lua["os"]["difftime"];
	os["time"] = lua["os"]["time"];
	env["os"] = os;

#if LUA_VERSION_NUM >= 502
	lua_rawgeti(lua, LUA_REGISTRYINDEX, env.registry_index());
	lua_rawseti(lua, LUA_REGISTRYINDEX, LUA_RIDX_GLOBALS);
#else
	int is_main = lua_pushthread(lua);
	SanityCheck(is_main);
	int thread = lua_gettop(lua);

	lua_rawgeti(lua, LUA_REGISTRYINDEX, env.registry_index());

	if (!lua_setfenv(lua, thread)) {
		throw ModException(
				"Security: Unable to set environment of the main Lua thread!");
	};
	lua_pop(lua, 1); // Pop thread
#endif
}

std::tuple<sol::object, sol::object> LuaSecurity::loadstring(
		const std::string &str, const std::string &chunkname) {
	if (!str.empty() && str[0] == LUA_SIGNATURE[0]) {
		return std::make_tuple(sol::nil,
				sol::make_object(lua, "Bytecode prohibited by Lua sandbox"));
	}

	sol::load_result result = lua.load(str, chunkname, sol::load_mode::text);
	if (result.valid()) {
		sol::function func = result;
		env.set_on(func);
		return std::make_tuple(func, sol::nil);
	} else {
		return std::make_tuple(
				sol::nil, sol::make_object(lua, ((sol::error)result).what()));
	}
}

std::tuple<sol::object, sol::object> LuaSecurity::loadfile(
		const std::string &path) {
	if (!checkPath(path, false)) {
		return std::make_tuple(sol::nil,
				sol::make_object(
						lua, "Path is not allowed by the Lua sandbox"));
	}

	std::ifstream t(path);
	std::string str((std::istreambuf_iterator<char>(t)),
			std::istreambuf_iterator<char>());
	return loadstring(str, "@" + path);
}

sol::object LuaSecurity::dofile(const std::string &path) {
	std::tuple<sol::object, sol::object> ret = loadfile(path);
	if (std::get<0>(ret) == sol::nil) {
		throw sol::error(std::get<1>(ret).as<std::string>());
	}

	sol::unsafe_function func = std::get<0>(ret);
	return func();
}

LuaSecurity.hpp

#pragma once

#include <sol/sol.hpp>

#include <tuple>

namespace scripting {

class LuaSecurity {
	sol::state &lua;
	sol::environment env;

public:
	explicit LuaSecurity(sol::state &lua) : lua(lua) { buildEnvironment(); }

	sol::environment &getEnvironment() { return env; }

	// Checks whether path is allowed
	bool checkPath(const std::string &path, bool write) { return true; }

private:
	void buildEnvironment();

	/// Secure loadstring. Prohibits bytecode, applies environment.
	///
	/// @param str Source code
	/// @param chunkname Chunk name
	/// @return Either (func, nil) or (nil, error-str)
	std::tuple<sol::object, sol::object> loadstring(const std::string &str,
			const std::string &chunkname = sol::detail::default_chunk_name());

	/// Secure loadfile. Checks path, then calls secure loadstring.
	///
	/// @param path Path to file
	/// @return Either (func, nil) or (nil, error-str)
	std::tuple<sol::object, sol::object> loadfile(const std::string &path);

	/// Secure dofile
	/// @param path Path to file
	/// @return Return value of function
	sol::object dofile(const std::string &path);
};

} // namespace scripting