<commit_before>package usercorn

import (
	"encoding/binary"
	"errors"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"

	"github.com/lunixbochs/usercorn/go/models"
)

type Unicorn struct {
	uc.Unicorn
	arch   *models.Arch
	os     *models.OS
	bits   int
	Bsz    int
	order  binary.ByteOrder
	memory []*models.Mmap
}

func NewUnicorn(arch *models.Arch, os *models.OS, order binary.ByteOrder) (*Unicorn, error) {
	Uc, err := uc.NewUnicorn(arch.UC_ARCH, arch.UC_MODE)
	if err != nil {
		return nil, err
	}
	return &Unicorn{
		Unicorn: Uc,
		arch:    arch,
		os:      os,
		bits:    arch.Bits,
		Bsz:     arch.Bits / 8,
		order:   order,
	}, nil
}

func (u *Unicorn) Arch() *models.Arch {
	return u.arch
}

func (u *Unicorn) OS() string {
	return u.os.Name
}

func (u *Unicorn) Bits() uint {
	return uint(u.bits)
}

func (u *Unicorn) ByteOrder() binary.ByteOrder {
	return u.order
}

func (u *Unicorn) mapping(addr, size uint64) *models.Mmap {
	for _, m := range u.memory {
		if addr < m.Addr && addr+size > m.Addr {
			return m
		}
		if addr >= m.Addr && addr < m.Addr+m.Size {
			return m
		}
	}
	return nil
}

func (u *Unicorn) Disas(addr, size uint64) (string, error) {
	mem, err := u.MemRead(addr, size)
	if err != nil {
		return "", err
	}
	return models.Disas(mem, addr, u.arch, u.Bsz)
}

func (u *Unicorn) MemMapProt(addr, size uint64, prot int) error {
	mmap, err := u.MemReserve(addr, size, true)
	if err != nil {
		return err
	}
	mmap.Prot = prot
	return u.Unicorn.MemMapProt(mmap.Addr, mmap.Size, prot)
}

func (u *Unicorn) MemMap(addr, size uint64) error {
	return u.MemMapProt(addr, size, uc.PROT_ALL)
}

func (u *Unicorn) MemProtect(addr, size uint64, prot int) error {
	if mmap := u.mapping(addr, size); mmap != nil {
		mmap.Prot = prot
	}
	return u.Unicorn.MemProtect(addr, size, prot)
}

func (u *Unicorn) MemUnmap(addr, size uint64) error {
	err := u.Unicorn.MemUnmap(addr, size)
	if err != nil {
		return err
	}
	for {
		mmap := u.mapping(addr, size)
		if mmap == nil {
			break
		}
		left := mmap.Addr
		right := left + mmap.Size
		// if unmap overlaps an edge, shrink that edge
		if addr <= left && addr+size > left {
			left = addr + size
		}
		if addr < right && addr+size >= right {
			right = addr
		}
		inMiddle := left < addr && right > addr+size
		// if our mapping now has size <= 0, delete it
		// also delete if the unmap was fully in the middle (as we'll split the mapping into each side)
		if left >= right || inMiddle {
			// delete by copying to a new array
			tmp := make([]*models.Mmap, 0, len(u.memory)-1)
			for _, v := range u.memory {
				if v != mmap {
					tmp = append(tmp, v)
				}
			}
			u.memory = tmp
		} else {
			// otherwise resize our existing mapping
			mmap.Addr = left
			mmap.Size = right - left
		}
		// if unmap range is fully in the middle, split and create a new mapping for each side
		if inMiddle {
			left := &models.Mmap{
				Addr: mmap.Addr, Size: addr - mmap.Addr,
				Prot: mmap.Prot, File: mmap.File, Desc: mmap.Desc,
			}
			right := &models.Mmap{
				Addr: addr + size, Size: (mmap.Addr + mmap.Size) - (addr + size),
				Prot: mmap.Prot, File: mmap.File, Desc: mmap.Desc,
			}
			u.memory = append(u.memory, left)
			u.memory = append(u.memory, right)
		}

	}
	return err
}

func (u *Unicorn) Mappings() []*models.Mmap {
	return u.memory
}

func (u *Unicorn) MemReserve(addr, size uint64, force bool) (*models.Mmap, error) {
	if addr == 0 {
		addr = BASE
	}
	addr, size = align(addr, size, true)
	if force {
		u.MemUnmap(addr, size)
		mmap := &models.Mmap{Addr: addr, Size: size, Prot: uc.PROT_ALL}<commit_msg>change up memory unmapping/protection<commit_after>package usercorn

import (
	"encoding/binary"
	"errors"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"

	"github.com/lunixbochs/usercorn/go/models"
)

type Unicorn struct {
	uc.Unicorn
	arch   *models.Arch
	os     *models.OS
	bits   int
	Bsz    int
	order  binary.ByteOrder
	memory []*models.Mmap
}

func NewUnicorn(arch *models.Arch, os *models.OS, order binary.ByteOrder) (*Unicorn, error) {
	Uc, err := uc.NewUnicorn(arch.UC_ARCH, arch.UC_MODE)
	if err != nil {
		return nil, err
	}
	return &Unicorn{
		Unicorn: Uc,
		arch:    arch,
		os:      os,
		bits:    arch.Bits,
		Bsz:     arch.Bits / 8,
		order:   order,
	}, nil
}

func (u *Unicorn) Arch() *models.Arch {
	return u.arch
}

func (u *Unicorn) OS() string {
	return u.os.Name
}

func (u *Unicorn) Bits() uint {
	return uint(u.bits)
}

func (u *Unicorn) ByteOrder() binary.ByteOrder {
	return u.order
}

func (u *Unicorn) mapping(addr, size uint64) *models.Mmap {
	for _, m := range u.memory {
		if addr < m.Addr && addr+size > m.Addr {
			return m
		}
		if addr >= m.Addr && addr < m.Addr+m.Size {
			return m
		}
	}
	return nil
}

func (u *Unicorn) Disas(addr, size uint64) (string, error) {
	mem, err := u.MemRead(addr, size)
	if err != nil {
		return "", err
	}
	return models.Disas(mem, addr, u.arch, u.Bsz)
}

func (u *Unicorn) MemMapProt(addr, size uint64, prot int) error {
	mmap, err := u.MemReserve(addr, size, true)
	if err != nil {
		return err
	}
	mmap.Prot = prot
	return u.Unicorn.MemMapProt(mmap.Addr, mmap.Size, prot)
}

func (u *Unicorn) MemMap(addr, size uint64) error {
	return u.MemMapProt(addr, size, uc.PROT_ALL)
}

func (u *Unicorn) MemProtect(addr, size uint64, prot int) error {
	// TODO: mapping a subregion should split the mapping?
	addr, size = align(addr, size, true)
	if mmap := u.mapping(addr, size); mmap != nil {
		mmap.Prot = prot
	}
	return u.Unicorn.MemProtect(addr, size, prot)
}

func (u *Unicorn) MemUnmap(addr, size uint64) error {
	// TODO: alignment check?
	for {
		mmap := u.mapping(addr, size)
		if mmap == nil {
			break
		}
		left := mmap.Addr
		right := left + mmap.Size
		// if unmap overlaps an edge, shrink that edge
		if addr <= left && addr+size > left {
			left = addr + size
		}
		if addr < right && addr+size >= right {
			right = addr
		}
		inMiddle := left < addr && right > addr+size
		// unmap in Unicorn
		if left >= right {
			u.Unicorn.MemUnmap(mmap.Addr, mmap.Size)
		} else {
			if left > mmap.Addr {
				u.Unicorn.MemUnmap(mmap.Addr, left-mmap.Addr)
			}
			if right < mmap.Addr+mmap.Size {
				u.Unicorn.MemUnmap(right, (mmap.Addr+mmap.Size)-right)
			}
		}
		// if our mapping now has size <= 0, delete it
		// also delete if the unmap was fully in the middle (as we'll split the mapping into each side)
		if left >= right || inMiddle {
			// delete by copying to a new array
			tmp := make([]*models.Mmap, 0, len(u.memory)-1)
			for _, v := range u.memory {
				if v != mmap {
					tmp = append(tmp, v)
				}
			}
			u.memory = tmp
		} else {
			// otherwise resize our existing mapping
			mmap.Addr = left
			mmap.Size = right - left
		}
		// if unmap range is fully in the middle, split and create a new mapping for each side
		if inMiddle {
			left := &models.Mmap{
				Addr: mmap.Addr, Size: addr - mmap.Addr,
				Prot: mmap.Prot, File: mmap.File, Desc: mmap.Desc,
			}
			right := &models.Mmap{
				Addr: addr + size, Size: (mmap.Addr + mmap.Size) - (addr + size),
				Prot: mmap.Prot, File: mmap.File, Desc: mmap.Desc,
			}
			u.memory = append(u.memory, left)
			u.memory = append(u.memory, right)
		}

	}
	return nil
}

func (u *Unicorn) Mappings() []*models.Mmap {
	return u.memory
}

func (u *Unicorn) MemReserve(addr, size uint64, force bool) (*models.Mmap, error) {
	if addr == 0 {
		addr = BASE
	}
	addr, size = align(addr, size, true)
	if force {
		u.MemUnmap(addr, size)
		mmap := &models.Mmap{Addr: addr, Size: size, Prot: uc.PROT_ALL}
{'fraction_non_alphanumeric': 0.14204168326032127, 'fraction_numerical': 0.005044470994291783, 'mean_word_length': 2.9548556430446196, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1491083', 'n_tokens_mistral': 3247, 'n_tokens_neox': 2916, 'n_words': 845}