61 changes: 53 additions & 8 deletions crates/lovely-core/src/lib.rs
@@ -1,10 +1,13 @@
#![allow(non_upper_case_globals)]

use std::collections::{HashMap, HashSet};
use std::fs;
use std::path::Path;
use std::{env, fs};
use std::path::{Path, PathBuf};

use manifest::Patch;
use log::*;

use getargs::{Arg, Options};
use manifest::{Patch, PatchArgs};
use sha2::{Digest, Sha256};
use sys::LuaState;

Expand All @@ -13,14 +16,19 @@ use crate::manifest::PatchManifest;
pub mod sys;
pub mod manifest;
pub mod patch;
pub mod hud;
pub mod log;

type LoadBuffer = dyn Fn(*mut LuaState, *const u8, isize, *const u8) -> u32 + Sync + Send;

pub struct PatchTable {
mod_dir: PathBuf,
loadbuffer: Option<Box<LoadBuffer>>,
targets: HashSet<String>,
vars: HashMap<String, String>,
patches: Vec<Patch>,
vanilla: bool,
vars: HashMap<String, String>,
args: HashMap<String, String>,
}

impl PatchTable {
Expand All @@ -29,7 +37,22 @@ impl PatchTable {
/// - MOD_DIR/lovely.toml
/// - MOD_DIR/lovely/*.toml
pub fn load(mod_dir: &Path) -> PatchTable {
let mod_dirs = fs::read_dir(mod_dir)
// Begin by parsing provided command line arguments. We'll parse patch args later.
let args = std::env::args().skip(1).collect::<Vec<_>>();
let mut opts = Options::new(args.iter().map(String::as_str));

let mut mod_dir = dirs::config_dir().unwrap().join("Balatro\\Mods");
let mut vanilla = false;

while let Some(opt) = opts.next_arg().expect("Failed to parse argument.") {
match opt {
Arg::Long("mod-dir") => mod_dir = opts.value().map(PathBuf::from).unwrap_or(mod_dir),
Arg::Long("vanilla") => vanilla = true,
_ => (),
}
}

let mod_dirs = fs::read_dir(&mod_dir)
.unwrap_or_else(|e| panic!("Failed to read from mod directory within {mod_dir:?}:\n{e:?}"))
.filter_map(|x| x.ok())
.filter(|x| x.path().is_dir())
Expand Down Expand Up @@ -107,10 +130,13 @@ impl PatchTable {
}

PatchTable {
mod_dir: mod_dir.to_path_buf(),
loadbuffer: None,
targets,
vars: var_table,
args: HashMap::new(),
patches,
vanilla,
}
}

Expand All @@ -128,6 +154,24 @@ impl PatchTable {
self.targets.contains(target)
}

/// Inject lovely metadata into the game.
/// # Safety
/// Unsafe due to internal unchecked usages of raw lua state.
pub unsafe fn inject_metadata(&self, state: *mut LuaState) {
let table = vec![
("mod_dir", self.mod_dir.to_str().unwrap().replace('\\', "/")),
("version", env!("CARGO_PKG_VERSION").to_string()),
];

let mut code = include_str!("../lovely.lua").to_string();
for (field, value) in table {
let field = format!("lovely_template:{field}");
code = code.replace(&field, &value);
}

sys::load_module(state, "lovely", &code, self.loadbuffer.as_ref().unwrap())
}

/// Apply one or more patches onto the target's buffer.
/// # Safety
/// Unsafe due to internal unchecked usages of raw lua state.
Expand Down Expand Up @@ -216,14 +260,15 @@ impl PatchTable {
}

let patched = new_buffer.join("\n");
println!("[LOVELY] Applied {patch_count} patches to '{target}'");
info!("[LOVELY] Applied {patch_count} patches to '{target}'");

// Compute the integrity hash of the patched file.
let mut hasher = Sha256::new();
hasher.update(patched.as_bytes());
let hash = format!("{:x}", hasher.finalize());

// Return the patched file with the prepended integrity hash.
format!("LOVELY_INTEGRITY = '{hash}'\n\n{patched}")
format!(
"LOVELY_INTEGRITY = '{hash}'\n\n{patched}"
)
}
}
47 changes: 47 additions & 0 deletions crates/lovely-core/src/log.rs
@@ -0,0 +1,47 @@
// Exports for convenience.
pub use log::{info, error, warn, debug, trace, LevelFilter};

use log::{Level, Log, Metadata, Record, SetLoggerError};

use crate::hud::MSG_TX;

static LOGGER: LovelyLogger = LovelyLogger {
use_console: false,
use_hud: true,
};

struct LovelyLogger {
use_console: bool,
use_hud: bool,
}

impl Log for LovelyLogger {
fn enabled(&self, metadata: &Metadata) -> bool {
metadata.level() <= Level::Debug
}

fn log(&self, record: &Record) {
if !self.enabled(record.metadata()) {
return;
}

let msg = format!("{} - {}", record.level(), record.args());
if self.use_console {
println!("{msg}");
}

if self.use_hud {
MSG_TX
.get()
.unwrap()
.send(msg)
.expect("Failed to pump hud log message");
}
}

fn flush(&self) {}
}

pub fn init() -> Result<(), SetLoggerError> {
log::set_logger(&LOGGER).map(|_| log::set_max_level(LevelFilter::Info))
}
42 changes: 42 additions & 0 deletions crates/lovely-core/src/manifest.rs
Expand Up @@ -13,21 +13,63 @@ pub struct Manifest {
pub struct PatchManifest {
pub manifest: Manifest,
pub patches: Vec<Patch>,

// A table of variable name = value bindings. These are interpolated
// into injected source code as the *last* step in the patching process.
#[serde(default)]
pub vars: HashMap<String, String>,

// A table of arguments, read and parsed from the environment command line.
// Binds double-hyphenated argument names (--arg) to a value, with additional metadata
// available to produce help messages, set default values, and apply other behavior.
#[serde(default)]
pub args: HashMap<String, PatchArgs>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct PatchArgs {
// An optional help string. This will be printed out in the calling console
// (if available) when the --help argument is supplied.
pub help: Option<String>,

// An optional default value. Not including a default value will cause Lovely
// to panic if this argument is missing or could not be parsed.
// Consider this to be both a "default value" and a "required" field, depending
// on whether or not it's set.
pub default: Option<String>,

// This field allows for a patch author to force lovely to parse incoming arguments
// with the exact name that they are defined by.
// This disables lovely's automatic underscore to hyphen conversion.
#[serde(default)]
pub name_override: bool,

// This field allows for arguments (--arg) to be passed without implicit values,
// treating it essentially as a flag. If it exists in the args, it's true, if not,
// then we set it to false.
#[serde(default)]
pub treat_as_flag: bool,
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "kebab-case")]
pub enum Patch {
// A patch which applies some change to a series of line(s) after a line with a match
// to the provided pattern has been found.
Pattern(PatternPatch),
Copy(CopyPatch),
Module(ModulePatch),
}

#[derive(Serialize, Deserialize, Debug)]
pub struct PatternPatch {
// The pattern that the line will be matched against. Very simple,
// supports only `?` (one occurance of any character) and `*` (any numver of any character).
// Patterns are matched against a left-trimmed version of the line, so whitespace does not
// need to be considered.
pub pattern: String,

// The position to insert the target at. `PatternAt::At` replaces the matched line entirely.
pub position: PatternAt,
pub target: String,
pub payload_files: Option<Vec<String>>,
Expand Down
32 changes: 31 additions & 1 deletion crates/lovely-core/src/sys.rs
@@ -1,4 +1,4 @@
use std::ffi::c_void;
use std::ffi::{c_void, CString};

use libc::FILE;
use libloading::{Library, Symbol};
Expand Down Expand Up @@ -105,3 +105,33 @@ pub static lua_isstring: Lazy<Symbol<unsafe extern "C" fn(*mut LuaState, isize)
LUA_LIB.get(b"lua_isstring").unwrap()
}
});

/// Load the provided buffer as a lua module with the specified name.
/// # Safety
/// Makes a lot of FFI calls, mutates internal C lua state.
pub unsafe fn load_module<F: Fn(*mut LuaState, *const u8, isize, *const u8) -> u32>(state: *mut LuaState, name: &str, buffer: &str, lual_loadbuffer: &F) {
let buf_cstr = CString::new(buffer).unwrap();
let buf_len = buf_cstr.as_bytes().len();

let p_name = format!("@{name}");
let p_name_cstr = CString::new(p_name).unwrap();

// Push the global package.loaded table onto the top of the stack, saving its index.
let stack_top = lua_gettop(state);
lua_getfield(state, LUA_GLOBALSINDEX, b"package\0".as_ptr() as _);
lua_getfield(state, -1, b"loaded\0".as_ptr() as _);

// This is the index of the `package.loaded` table.
let field_index = lua_gettop(state);

// Load the buffer and execute it via lua_pcall, pushing the result to the top of the stack.
lual_loadbuffer(state, buf_cstr.into_raw() as _, buf_len as _, p_name_cstr.into_raw() as _);

lua_pcall(state, 0, -1, 0);

// Insert pcall results onto the package.loaded global table.
let module_cstr = CString::new(name).unwrap();

lua_setfield(state, field_index, module_cstr.into_raw() as _);
lua_settop(state, stack_top);
}
9 changes: 5 additions & 4 deletions crates/lovely-win/Cargo.toml
Expand Up @@ -10,11 +10,12 @@ crate-type = ["cdylib"]
[dependencies]
lovely-core = { version ="*", path = "../lovely-core" }

dirs = "5.0.1"
getargs = "0.5.0"
libc = "0.2.153"
libc = "0.2.141"
once_cell = "1.19.0"
widestring = "1.0.2"
hudhook = "0.6.5"
getargs = "0.5.0"
dirs = "5.0.1"

[dependencies.retour]
version = "0.4.0-alpha.1"
Expand All @@ -36,4 +37,4 @@ features = [
]

[build-dependencies]
forward-dll = "0.1.13"
forward-dll = "0.1.13"
53 changes: 25 additions & 28 deletions crates/lovely-win/src/lib.rs
@@ -1,26 +1,20 @@
use std::panic;
use std::slice;
use std::ffi::{c_void, CStr, CString};
use std::fs;
use std::path::PathBuf;
use std::sync::mpsc;

use getargs::{Arg, Options};
use lovely_core::log::*;
use lovely_core::{hud, PatchTable};
use lovely_core::sys::{self, LuaState};

use lovely_core::PatchTable;
use getargs::{Arg, Options};
use once_cell::sync::OnceCell;
use retour::static_detour;
use widestring::U16CString;
use windows::core::{s, w, PCWSTR};
use windows::Win32::Foundation::{HINSTANCE, HWND};
use windows::Win32::System::Console::{
AllocConsole,
GetConsoleMode,
GetStdHandle,
SetConsoleMode,
CONSOLE_MODE,
ENABLE_VIRTUAL_TERMINAL_PROCESSING,
STD_OUTPUT_HANDLE
};
use windows::Win32::System::LibraryLoader::{GetProcAddress, LoadLibraryW};
use windows::Win32::UI::WindowsAndMessaging::{MessageBoxW, MESSAGEBOX_STYLE};

Expand All @@ -34,17 +28,20 @@ static_detour! {
}

unsafe extern "C" fn lua_loadbuffer_detour(state: *mut LuaState, buf_ptr: *const u8, size: isize, name_ptr: *const u8) -> u32 {
let patch_table = PATCH_TABLE.get().unwrap();

// Install native function overrides *once*.
if HAS_INIT.get().is_none() {
let closure = override_print as *const c_void;
sys::lua_pushcclosure(state, closure, 0);
sys::lua_setfield(state, sys::LUA_GLOBALSINDEX, b"print\0".as_ptr() as _);

patch_table.inject_metadata(state);

HAS_INIT.set(()).unwrap();
}

let name = CStr::from_ptr(name_ptr as _).to_str().unwrap();
let patch_table = PATCH_TABLE.get().unwrap();

if !patch_table.needs_patching(name) {
return LuaLoadbuffer_Detour.call(state, buf_ptr, size, name_ptr);
Expand Down Expand Up @@ -84,6 +81,7 @@ unsafe extern "system" fn DllMain(_: HINSTANCE, reason: u32, _: *const c_void) -

std::panic::set_hook(Box::new(|x| unsafe {
let message = format!("lovely-injector has crashed: \n{x}");
error!("{message}");

let message = U16CString::from_str(message);
MessageBoxW(
Expand All @@ -93,18 +91,16 @@ unsafe extern "system" fn DllMain(_: HINSTANCE, reason: u32, _: *const c_void) -
MESSAGEBOX_STYLE(0),
);
}));

// Setup console redirection, replacing Love's own implementation.
let _ = AllocConsole();

// Enable virtual terminal processing to allow for fancy colored text.
let stdout = GetStdHandle(STD_OUTPUT_HANDLE).unwrap();
lovely_core::log::init().unwrap_or_else(|e| panic!("Failed to initialize logger: {e:?}"));

let mut mode = CONSOLE_MODE(0);
GetConsoleMode(stdout, &mut mode as *mut _).unwrap();
// Create the UI
let (tx, rx) = mpsc::channel::<String>();
std::thread::spawn(move || {
hud::open(rx);
});

let mode = mode.0 | ENABLE_VIRTUAL_TERMINAL_PROCESSING.0;
SetConsoleMode(stdout, CONSOLE_MODE(mode)).unwrap();
hud::MSG_TX.set(tx).unwrap();

let args = std::env::args().skip(1).collect::<Vec<_>>();
let mut opts = Options::new(args.iter().map(String::as_str));
Expand All @@ -122,16 +118,16 @@ unsafe extern "system" fn DllMain(_: HINSTANCE, reason: u32, _: *const c_void) -

// Stop here if we're runnning in vanilla mode. Don't install hooks, don't setup patches, etc.
if vanilla {
println!("[LOVELY] Running in vanilla mode");
info!("Running in vanilla mode");
return 1;
}

if !mod_dir.is_dir() {
println!("[LOVELY] Creating mods directory at {mod_dir:?}");
info!("Creating mods directory at {mod_dir:?}");
fs::create_dir_all(&mod_dir).unwrap();
}

println!("[LOVELY] Using mods directory at {mod_dir:?}");
info!("Using mods directory at {mod_dir:?}");

// Patch files are stored within the root of mod subdirectories within the mods dir.
// - MOD_DIR/lovely.toml
Expand Down Expand Up @@ -164,6 +160,7 @@ unsafe extern "system" fn DllMain(_: HINSTANCE, reason: u32, _: *const c_void) -
/// Native lua API access. It's unsafe, it's unchecked, it will probably eat your firstborn.
pub unsafe extern "C" fn override_print(state: *mut LuaState) -> isize {
let argc = sys::lua_gettop(state);
let mut out = String::new();

for i in 0..argc {
let mut str_len = 0_isize;
Expand All @@ -173,14 +170,14 @@ pub unsafe extern "C" fn override_print(state: *mut LuaState) -> isize {
let arg_str = String::from_utf8(str_buf.to_vec()).unwrap();

if i > 1 {
print!("\t");
out.push('\t');
}

print!("[GAME] {arg_str}");

out.push_str(&format!("[GAME] {arg_str}"));
sys::lua_settop(state, -(1) - 1);
}
println!();

info!("{out}");

0
}