Skip to content

Commit

Permalink
changing code buffs if the current one runs out of space
Browse files Browse the repository at this point in the history
  • Loading branch information
@matthewfl
matthewfl committed Jul 23, 2016
1 parent a56edc5 commit ae0dcd1
Show file tree
Hide file tree
Showing 5 changed files with 140 additions and 41 deletions.
93 changes: 80 additions & 13 deletions src/code_buffer.cc
View file
Original file line number Diff line number Diff line change
@@ -1,35 +1,101 @@
#include "jit_internal.h"

#include <sys/mman.h>
#include <mutex>


using namespace std;

using namespace redmagic;

extern "C" void red_asm_compile_buff_near();

CodeBuffer::CodeBuffer(size_t size):
size(size),
owns_buffer(true),
can_write_buffer(true),
buffer_consumed(0)
{
namespace redmagic {
std::mutex code_buffer_mutex;

std::vector<CodeBuffer*> owning_code_buffers;

std::vector<CodeBuffer*> unallocated_code_buffers;
}

CodeBuffer* CodeBuffer::CreateBuffer(size_t size) {
std::lock_guard<std::mutex> lock(code_buffer_mutex);

buffer = (uint8_t*)mmap((void*)&red_asm_compile_buff_near,
for(auto it = unallocated_code_buffers.begin(); it != unallocated_code_buffers.end(); it++) {
if((*it)->getFree() > size) {
CodeBuffer *ret = *it;
unallocated_code_buffers.erase(it);
return ret;
}
}

// we could not find an open buffer that has enough space to satasify this request
if(size < 4 * 1024 * 1024) {
// the min size that we will allocate
size = 4 * 1024 * 1024;
} else if((size & (4*1024 - 1)) != 0) {
// request in page amounts
size += 4 * 1024;
size &= (4*1024 - 1);
}

uint8_t *buffer = (uint8_t*)mmap((void*)&red_asm_compile_buff_near,
size, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE | MAP_ANON, -1, 0);

if(buffer == MAP_FAILED) {
perror("failed to mmap buffer");
}
memset(buffer, 0, size);

CodeBuffer* ret = new CodeBuffer();
owning_code_buffers.push_back(ret);

ret->size = size;
ret->owns_buffer = true;
ret->can_write_buffer = true;
ret->buffer_consumed = 0;
ret->buffer = buffer;

return ret;
}

CodeBuffer::~CodeBuffer() {
if(owns_buffer) {
if(munmap(buffer, size) < 0) {
perror("failed to unmap buffer");
}

void CodeBuffer::Release(CodeBuffer *x) {
std::lock_guard<std::mutex> lock(code_buffer_mutex);

assert(x->owns_buffer);

// this buffer should not already be in the list of currently unused
for(auto it : unallocated_code_buffers) {
assert(x != it);
}

unallocated_code_buffers.push_back(x);
}

// CodeBuffer::CodeBuffer(size_t size):
// size(size),
// owns_buffer(true),
// can_write_buffer(true),
// buffer_consumed(0)
// {

// buffer = (uint8_t*)mmap((void*)&red_asm_compile_buff_near,
// size, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE | MAP_ANON, -1, 0);

// if(buffer == MAP_FAILED) {
// perror("failed to mmap buffer");
// }
// memset(buffer, 0, size);
// }

CodeBuffer::~CodeBuffer() {
// if(owns_buffer) {
// if(munmap(buffer, size) < 0) {
// perror("failed to unmap buffer");
// }
// }
// can't un allocate an owning buffer atm
assert(!owns_buffer);
}

CodeBuffer::CodeBuffer(mem_loc_t start, size_t size, bool override_can_write):
Expand All @@ -54,6 +120,7 @@ CodeBuffer::CodeBuffer(CodeBuffer &&x) {
trampolines_size = x.trampolines_size;
size = x.size;
buffer_consumed = x.buffer_consumed;
assert(!x.owns_buffer); /////// TODO: move an owning buffer?
owns_buffer = x.owns_buffer;
x.owns_buffer = false;
can_write_buffer = x.can_write_buffer;
Expand Down
9 changes: 8 additions & 1 deletion src/jit_internal.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,14 @@ namespace redmagic {

class CodeBuffer final {
public:
CodeBuffer(size_t size);
//
public:
// creates a Code Buffer that "owns" a region of memory that is of at least size
static CodeBuffer* CreateBuffer(size_t size);
// when done with a code buffer release it back to a memory pool
static void Release(CodeBuffer *x);

//CodeBuffer(size_t size);
CodeBuffer(mem_loc_t start, size_t size, bool override_can_write=false);
CodeBuffer();

Expand Down
3 changes: 2 additions & 1 deletion src/manager.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,8 @@ void* Manager::begin_trace(void *id, void *ret_addr) {
//assert(tracer == nullptr);
assert(info->tracer == nullptr);

auto buff = make_shared<CodeBuffer>(4 * 1024 * 1024);
//auto buff = make_shared<CodeBuffer>(4 * 1024 * 1024);
auto buff = CodeBuffer::CreateBuffer(1024 * 1024);
new_head->tracer = l = new Tracer(buff);
l->tracing_from = (mem_loc_t)trace_pc;
l->owning_thread = get_thread_id();
Expand Down
72 changes: 48 additions & 24 deletions src/tracer.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ static int udis_input_hook (ud_t *ud) {
return (int)*((uint8_t*)l);
}

Tracer::Tracer(shared_ptr<CodeBuffer> buffer) {
Tracer::Tracer(CodeBuffer* buffer) {
this->buffer = buffer;

ud_init(&disassm);
Expand Down Expand Up @@ -73,6 +73,8 @@ Tracer::~Tracer() {
// free((void*)stack);

// free((void*)stack);
if(buffer != nullptr)
CodeBuffer::Release(buffer);
}


Expand Down Expand Up @@ -221,7 +223,7 @@ void* Tracer::Start(void *start_addr) {
//red_asm_start_tracing(NULL, (void*)&red_begin_tracing, this, stack - sizeof(stack));

using namespace asmjit;
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
//compiler.mov(x86::rdx, x86::rsp);
// stash the values of the register that we are about to override
compiler.mov(x86::ptr(x86::rsp, static_cast<int32_t>(-TRACE_STACK_OFFSET - sizeof(struct user_regs_struct))), x86::rdx);
Expand Down Expand Up @@ -304,10 +306,32 @@ void Tracer::Run(struct user_regs_struct *other_stack) {
last_location = udis_loc;
local_jump_min_addr = last_local_jump = 0;
assert(current_location == last_location);
assert(protected_malloc);
//assert(generation_lock.owns_lock());

// TODO: manage cases where the buffer runs out of space
assert(buffer->getFree() > 100 * 1024);
// if we somehow have less then 1kb free then we might have overwritten something
// which is why this is asserted as an error
assert(buffer->getFree() > 1024);
if(buffer->getFree() <= 10 * 1024) {
// there is less than 10 kb of space on this buffer, so we are going to make a new one
protected_malloc = false;
auto new_buffer = CodeBuffer::CreateBuffer(1024 * 1024);
{
SimpleCompiler compiler(new_buffer);
compiler.mov(asmjit::x86::r15, asmjit::imm_u(0xdeadcafe));
compiler.mov(asmjit::x86::r15, asmjit::imm_u(generated_location));
}
auto new_gen_l = new_buffer->getRawBuffer() + new_buffer->getOffset();
{
SimpleCompiler compiler(buffer);
compiler.jmp(asmjit::imm_ptr(new_gen_l));
auto written = compiler.finalize();
}
CodeBuffer::Release(buffer);
buffer = new_buffer;
generated_location = new_gen_l;
protected_malloc = true;
}

processes_instructions:
while(ud_disassemble(&disassm)) {
Expand Down Expand Up @@ -421,7 +445,7 @@ void* Tracer::EndTraceFallthrough() {
assert(icount - last_call_instruction < 2);
red_printf("tracer end fallthrough\n");
buffer->setOffset(last_call_generated_op);
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
compiler.mov(asmjit::x86::rdi, asmjit::imm_u(last_call_ret_addr));
compiler.jmp(asmjit::imm_ptr(&red_asm_end_trace));
auto w = compiler.finalize();
Expand All @@ -441,7 +465,7 @@ void* Tracer::EndTraceLoop() {
assert(loop_location);

buffer->setOffset(last_call_generated_op);
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
compiler.jmp(asmjit::imm_ptr(loop_location));
finish_patch();
tracing_from = 0;
Expand All @@ -451,7 +475,7 @@ void* Tracer::EndTraceLoop() {
void* Tracer::TempDisableTrace() {
assert(icount - last_call_instruction < 2);
buffer->setOffset(last_call_generated_op);
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
auto label = compiler.newLabel();
//compiler.mov(asmjit::x86::rdi, asmjit::imm_u(0xfafafafafafafafa));
compiler.lea(asmjit::x86::rdi, asmjit::x86::ptr(label));
Expand All @@ -476,7 +500,7 @@ extern "C" void red_asm_jump_rsi();
void Tracer::TempEnableTrace(void *resume_pc) {
// check that the temp enable instruction is coming in at a expected spot, otherwise fork a new trace
set_pc((mem_loc_t)resume_pc);
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
// the "normal" return address will be set to ris when this returns from the temp disabled region
//compiler.mov(asmjit::x86::rax, asmjit::x86::ptr(asmjit::x86::rsp, -8));
compiler.TestRegister((mem_loc_t)&red_asm_jump_rsi, RSI, (register_t)resume_pc);
Expand All @@ -490,7 +514,7 @@ void Tracer::JumpToNestedLoop(void *nested_trace_id) {
// there should have been a backwards branch instruction that we are going to replace
assert(icount - last_call_instruction < 2);
buffer->setOffset(last_call_generated_op);
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
compiler.mov(asmjit::x86::rdi, asmjit::imm_u(0xfafafafafafafafa));
compiler.mov(asmjit::x86::rsi, asmjit::imm_ptr(nested_trace_id));
compiler.mov(asmjit::x86::rdx, asmjit::imm_u(last_call_ret_addr));
Expand All @@ -508,7 +532,7 @@ void Tracer::JumpToNestedLoop(void *nested_trace_id) {
void* Tracer::ReplaceIsTracedCall() {
assert(icount - last_call_instruction < 2);
buffer->setOffset(last_call_generated_op);
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
compiler.mov(asmjit::x86::rax, asmjit::imm(1));
auto written = compiler.finalize();
write_interrupt_block();
Expand Down Expand Up @@ -584,7 +608,7 @@ void Tracer::write_interrupt_block() {
{
// TODO: DON'T USE SimpleCompiler here, it has a high overhead for starting and cleanup
// will be much faster to just hardcode a jmp with 32 bid address
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
compiler.jmp(asmjit::imm_u(interrupt_block_location));
}
buffer->setOffset(offset);
Expand Down Expand Up @@ -962,7 +986,7 @@ void Tracer::evaluate_instruction() {
}
}

SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);

auto cb_b = compiler.ConditionalJump(alternate_instructions, AlignedInstructions::get_asm_mnem(emit));

Expand Down Expand Up @@ -1044,7 +1068,7 @@ void Tracer::evaluate_instruction() {
// auto written = buffer->writeToEnd(test_register_instructions[ri].instruction);
// written.replace_stump<uint64_t>(0xfafafafafafafafa, rv);

SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
auto r_cb = compiler.TestRegister(ud_insn_off(&disassm), ri, rv);
auto written = compiler.finalize();

Expand All @@ -1065,7 +1089,7 @@ void Tracer::evaluate_instruction() {
}
jump_dest = t;
} else {
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
AlignedInstructions ali(&disassm);
auto v = get_opr_value(opr);
compiler.add_used_registers(ali.registers_used());
Expand Down Expand Up @@ -1105,7 +1129,7 @@ void Tracer::evaluate_instruction() {
buffer->setOffset(last_call_generated_op);
pop_stack();
set_pc(last_call_ret_addr);
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
compiler.call(asmjit::imm_ptr(jump_dest));
auto written = compiler.finalize();
write_interrupt_block();
Expand All @@ -1114,7 +1138,7 @@ void Tracer::evaluate_instruction() {
// we are jumping to another method and isn't the first instruction, which means that this is behaving like a tail call optimization
register_t return_pc = peek_stack();
set_pc(return_pc);
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
// pop the previous return address off the stack
compiler.add(asmjit::x86::rsp, asmjit::imm(8));
compiler.call(asmjit::imm_ptr(jump_dest));
Expand Down Expand Up @@ -1183,7 +1207,7 @@ void Tracer::evaluate_instruction() {
} else {

// TODO: check that the register is pointing at the same location in memory
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
AlignedInstructions ali(&disassm);
auto v = get_opr_value(opr1);
compiler.add_used_registers(ali.registers_used());
Expand Down Expand Up @@ -1228,7 +1252,7 @@ void Tracer::evaluate_instruction() {
// written.replace_stump<uint64_t>(0xfafafafafafafafa, call_pc);
mem_loc_t cont_addr;
{
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
compiler.call(asmjit::imm_ptr(call_pc));
auto cb = compiler.finalize();

Expand All @@ -1244,7 +1268,7 @@ void Tracer::evaluate_instruction() {
// inline this method, so push the return address and continue
// auto written = buffer->writeToEnd(cb_asm_push_stack);
// written.replace_stump<uint64_t>(0xfafafafafafafafa, ret_addr);
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
// compiler.mov(asmjit::x86::r15, asmjit::imm(ret_addr));
// compiler.push(asmjit::imm(ret_addr));
compiler.Push64bitValue(ret_addr);
Expand Down Expand Up @@ -1308,7 +1332,7 @@ void Tracer::replace_rip_instruction() {

// auto written = buffer->writeToEnd(set_register_instructions[dest].instruction);
// written.replace_stump<uint64_t>(0xfafafafafafafafa, val.address);
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
compiler.SetRegister(dest, val.address);
return;
}
Expand All @@ -1330,7 +1354,7 @@ void Tracer::replace_rip_instruction() {
// assert(opr1->type == UD_OP_REG);
int dest = ud_register_to_sys(opr1->base);

SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
auto r = compiler.get_register(dest);
compiler.mov(r, asmjit::imm_u(val.address));
// use the aligned instruction so that the correct size is read from memory
Expand All @@ -1348,7 +1372,7 @@ void Tracer::replace_rip_instruction() {
assert(opr1->index == UD_NONE);
opr_value val = get_opr_value(opr1);
assert(val.is_ptr);
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
compiler.PushMemoryLocationValue(val.address);
auto written = compiler.finalize();
return;
Expand Down Expand Up @@ -1380,7 +1404,7 @@ void Tracer::replace_rip_instruction() {
assert((used_registers & (1 << RSP)) == 0);
used_registers &= ~(1 << RIP);

SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
compiler.add_used_registers(used_registers);
auto scr = compiler.get_scratch_register();
compiler.mov(scr, udis_loc); // load the current rip into the scratch register
Expand Down Expand Up @@ -1428,7 +1452,7 @@ void Tracer::abort() {

void Tracer::run_debugger() {
red_printf("\n----ABORT debugger---\n");
SimpleCompiler compiler(buffer.get());
SimpleCompiler compiler(buffer);
compiler.int3();
//compiler.hlt();
compiler.finalize();
Expand Down
Loading

0 comments on commit ae0dcd1

Please sign in to comment.