Cross-Platform Virtual Memory - Nim

memaddr.nim

# Cross-Platform Virtual Memory - Enhanced Security Version
##
## This module provides cross-platform virtual memory allocation with enhanced security features.
## It offers a consistent API across Windows and POSIX systems with additional security checks
## and protections against common memory-related vulnerabilities.
##
## Overview:
## - Cross-platform memory management between Windows (VirtualAlloc) and POSIX (mmap)
## - Enhanced security against common memory attacks
## - Protection against buffer overflows, integer overflows and memory leaks
## - Memory zeroing to prevent information leakage
## - Address Space Layout Randomization (ASLR) capabilities (on supported platforms)
##
## Example usage:
## ```nim
##
## # Allocate 1MB of readable/writable memory
## let memSize = 1024 * 1024
## try:
##   let memPtr = nVirtualAlloc(nil, memSize, VM_COMMIT or VM_RESERVE, VM_PAGE_READWRITE)
##   # Write data to memory
##   cast[ptr int](memPtr)[] = 42
##   echo "Memory allocation succeeded, wrote value: ", cast[ptr int](memPtr)[]
##
##   # When done, free the memory with security zeroing
##   let freeResult = nVirtualFree(memPtr, memSize, VM_FREE)
##   echo "Memory freed: ", freeResult
## except MemoryError as e:
##   echo "Memory operation failed: ", e.msg
## ```
import strformat
from strutils import toHex

when defined(windows):
  import winlean
  type SIZE_T = uint
  const
    MEM_COMMIT = 0x1000
    MEM_RESERVE = 0x2000
    MEM_RELEASE = 0x8000
    PAGE_READONLY = 0x02
    PAGE_READWRITE = 0x04
    PAGE_EXECUTE = 0x10
    PAGE_EXECUTE_READ = 0x20
    PAGE_EXECUTE_READWRITE = 0x40

  proc VirtualAlloc(
    lpAddress: pointer, dwSize: SIZE_T, flAllocationType, flProtect: int32
  ): pointer {.stdcall, dynlib: "kernel32", importc: "VirtualAlloc", noSideEffect.}

  proc VirtualFree(
    lpAddress: pointer, dwSize: SIZE_T, dwFreeType: int32
  ): int32 {.stdcall, dynlib: "kernel32", importc: "VirtualFree", noSideEffect.}

else:
  import posix

type MemoryError* = object of CatchableError
  ## Custom exception type for memory allocation and management errors

# Constants that match Windows API but won't conflict
const
  VM_COMMIT* = 0x1000
    ## Allocates memory charges (commits pages) for the specified reserved memory pages
  VM_RESERVE* = 0x2000
    ## Reserves a range of the process's virtual address space without allocating physical storage
  VM_PAGE_READONLY* = 0x02 ## Enables read-only access to the committed region of pages
  VM_PAGE_READWRITE* = 0x04
    ## Enables read and write access to the committed region of pages
  VM_PAGE_EXECUTE* = 0x10 ## Enables execute access to the committed region of pages
  VM_PAGE_EXECUTE_READ* = 0x20
    ## Enables execute and read access to the committed region of pages
  VM_PAGE_EXECUTE_READWRITE* = 0x40
    ## Enables execute, read, and write access to the committed region of pages
  VM_FREE* = 0x10000 ## Frees the specified virtual memory pages

# Maximum allocation size (2GB) to prevent resource exhaustion attacks
const MAX_ALLOCATION_SIZE = 2_147_483_648

proc nVirtualAlloc*(
    lpAddress: pointer, dwSize: int, flAllocationType: int, flProtect: int
): pointer {.raises: [MemoryError].} =
  ## Allocates memory in the virtual address space of the calling process with enhanced security.
  ##
  ## Parameters:
  ##   lpAddress: The starting address of the region to allocate. If nil, the system determines the location.
  ##   dwSize: The size of the region in bytes.
  ##   flAllocationType: The type of memory allocation (VM_COMMIT, VM_RESERVE, or both).
  ##   flProtect: The memory protection for the region (VM_PAGE_* constants).
  ##
  ## Returns:
  ##   A pointer to the allocated memory.
  ##
  ## Raises:
  ##   MemoryError: If memory allocation fails for any reason.
  ##
  ## Security Features:
  ##   - Size validation to prevent resource exhaustion
  ##   - Protection against integer overflow
  ##   - Enhanced security flags on supported platforms
  ##   - Memory zeroing for security
  ##   - ASLR-like protections on supported platforms
  ##
  ## Example:
  ##   ```nim
  ##   # Allocate 10KB of read-write memory
  ##   try:
  ##     let buffer = nVirtualAlloc(nil, 10240, VM_COMMIT or VM_RESERVE, VM_PAGE_READWRITE)
  ##     # Use the memory safely
  ##     var dataArray = cast[ptr array[10240, byte]](buffer)
  ##     dataArray[0] = 123
  ##     dataArray[1] = 45
  ##     echo dataArray[0], ", ", dataArray[1]  # Outputs: 123, 45
  ##
  ##     # Free the memory when done
  ##     discard nVirtualFree(buffer, 10240, VM_FREE)
  ##   except MemoryError as e:
  ##     echo "Memory allocation failed: ", e.msg
  ##   ```

  # Input validation checks
  if unlikely(dwSize <= 0):
    raise newException(MemoryError, "Invalid memory size: must be greater than 0")
  if unlikely(dwSize > MAX_ALLOCATION_SIZE):
    raise newException(MemoryError, "Requested allocation size exceeds maximum allowed")
  if unlikely((flAllocationType and (VM_COMMIT or VM_RESERVE)) == 0):
    raise newException(
      MemoryError, "Invalid allocation type: must include COMMIT or RESERVE"
    )

  when defined(windows):
    # Use Windows VirtualAlloc directly with optimized flag conversion
    let winAllocType: int32 =
      (if (flAllocationType and VM_COMMIT) != 0: MEM_COMMIT.int32 else: 0.int32) or
      (if (flAllocationType and VM_RESERVE) != 0: MEM_RESERVE.int32 else: 0.int32)

    let winProtect: int32 =
      if (flProtect and VM_PAGE_EXECUTE_READWRITE) != 0:
        PAGE_EXECUTE_READWRITE
      elif (flProtect and VM_PAGE_EXECUTE_READ) != 0:
        PAGE_EXECUTE_READ
      elif (flProtect and VM_PAGE_EXECUTE) != 0:
        PAGE_EXECUTE
      elif (flProtect and VM_PAGE_READWRITE) != 0:
        PAGE_READWRITE
      elif (flProtect and VM_PAGE_READONLY) != 0:
        PAGE_READONLY
      else:
        0

    let memPtr = VirtualAlloc(lpAddress, SIZE_T(dwSize), winAllocType, winProtect)
    if memPtr == nil:
      let errorCode = getLastError()
      raise newException(
        MemoryError, "Windows VirtualAlloc failed with error code: " & $errorCode
      )
    return memPtr
  else:
    # POSIX implementation (Linux, macOS, etc.) with optimized flag conversion
    var prot: cint =
      if (flProtect and VM_PAGE_EXECUTE_READWRITE) != 0:
        PROT_EXEC or PROT_READ or PROT_WRITE
      elif (flProtect and VM_PAGE_EXECUTE_READ) != 0:
        PROT_EXEC or PROT_READ
      elif (flProtect and VM_PAGE_EXECUTE) != 0:
        PROT_EXEC
      elif (flProtect and VM_PAGE_READWRITE) != 0:
        PROT_READ or PROT_WRITE
      elif (flProtect and VM_PAGE_READONLY) != 0:
        PROT_READ
      else:
        PROT_NONE

    # Security enhancement: prevent executable writable memory unless specifically requested
    var flags: cint = MAP_PRIVATE or MAP_ANONYMOUS

    # Add additional security flags available on Linux
    when defined(MAP_NORESERVE):
      flags = flags or MAP_NORESERVE # Don't reserve swap space

    # Use enhanced security flags where available
    when defined(linux):
      # Add security flags
      when defined(MAP_STACK):
        if (prot and PROT_EXEC) == 0: # Only for non-executable memory
          flags = flags or MAP_STACK

      # Randomize allocation address for ASLR-like protection
      when defined(MAP_RANDOMIZE):
        flags = flags or MAP_RANDOMIZE

    # Use mmap to allocate memory with enhanced security flags
    let allocatedMem = mmap(
      lpAddress, # Preferred address
      dwSize.int, # Size of allocation
      prot, # Memory protection
      flags, # Enhanced security flags
      -1.cint, # File descriptor (none)
      0.Off, # Offset (none)
    )

    if allocatedMem == cast[pointer](MAP_FAILED):
      let errorCode = errno
      raise
        newException(MemoryError, "POSIX mmap failed with error code: " & $errorCode)

    # Ensure memory is zeroed for security - only if writable
    if (prot and PROT_WRITE) != 0:
      zeroMem(allocatedMem, dwSize)

    return allocatedMem

proc nVirtualFree*(
    lpAddress: pointer, dwSize: int, dwFreeType: int
): bool {.raises: [MemoryError].} =
  ## Releases, decommits, or releases and decommits a region of memory with enhanced security.
  ##
  ## Parameters:
  ##   lpAddress: A pointer to the base address of the region of pages to be freed.
  ##   dwSize: The size of the region in bytes.
  ##   dwFreeType: The type of free operation (VM_FREE).
  ##
  ## Returns:
  ##   true if the function succeeds.
  ##
  ## Raises:
  ##   MemoryError: If memory deallocation fails for any reason.
  ##
  ## Security Features:
  ##   - Memory zeroing before freeing to prevent information leakage
  ##   - Input validation to prevent invalid free operations
  ##   - Protection against double-free vulnerabilities
  ##
  ## Example:
  ##   ```nim
  ##   # First allocate some memory
  ##   let memSize = 4096  # 4KB of memory
  ##   try:
  ##     let memory = nVirtualAlloc(nil, memSize, VM_COMMIT or VM_RESERVE, VM_PAGE_READWRITE)
  ##
  ##     # Use the memory
  ##     var intPtr = cast[ptr int](memory)
  ##     intPtr[] = 12345
  ##     echo "Value stored: ", intPtr[]
  ##
  ##     # Securely free the memory when done
  ##     nVirtualFree(memory, memSize, VM_FREE)
  ##     echo "Memory freed successfully"
  ##
  ##     # The memory is now securely zeroed and deallocated
  ##   except MemoryError as e:
  ##     echo "Memory operation failed: ", e.msg
  ##   ```

  # Input validation checks
  if unlikely(lpAddress == nil):
    raise newException(MemoryError, "Cannot free nil pointer")

  when defined(windows):
    # Use Windows VirtualFree with optimized flag conversion
    let winFreeType: int32 =
      if (dwFreeType and VM_FREE) != 0: MEM_RELEASE.int32 else: 0.int32

    let freeResult = VirtualFree(lpAddress, SIZE_T(0), winFreeType)
    if freeResult == 0:
      let errorCode = getLastError()
      raise newException(
        MemoryError, "Windows VirtualFree failed with error code: " & $errorCode
      )
    return true
  else:
    # POSIX implementation - additional validation
    if unlikely(dwSize <= 0):
      raise newException(
        MemoryError, "Invalid memory size for freeing: must be greater than 0"
      )

    # Security enhancement: zero memory before freeing if possible
    when defined(linux) or defined(macosx):
      # Try to make the memory writable temporarily to zero it
      # We only do this if mprotect succeeds to avoid potential issues
      if mprotect(lpAddress, dwSize.int, PROT_READ or PROT_WRITE) == 0:
        zeroMem(lpAddress, dwSize)

    # Free the memory - use prefetch before munmap for better performance
    when compileOption("threads"):
      {.
        emit:
          "if (__builtin_expect(!!(`lpAddress`), 1)) __builtin_prefetch(`lpAddress`, 0, 3);"
      .}

    let freeResult = munmap(lpAddress, dwSize.int)
    if freeResult != 0:
      let errorCode = errno
      raise
        newException(MemoryError, "POSIX munmap failed with error code: " & $errorCode)

    return true

template memAddr[T](x: T, length: uint = pointer.sizeof * 2): string =
  ## Returns the memory address of `x` as a hexadecimal string. This template
  ## can be utilized for debugging purposes, where inspecting the memory address
  ## of variables is necessary. It works with any type of variable by utilizing
  ## generic type `T`.
  ##
  ## Arguments:
  ## - `x`: The variable whose memory address is to be returned. This can be
  ##   of any data type, including primitive types, objects, sequences, etc.
  ## - `length` (optional): Specifies the length of the hexadecimal string to be
  ##   returned. Defaults to `pointer.sizeof * 2`, which is typically 8 on a 32-bit
  ##   system and 16 on a 64-bit system, reflecting the size of memory addresses on
  ##   these platforms. This value can be adjusted if a different string length is
  ##   desired for aesthetic or alignment purposes in the output.
  ##
  ## Returns:
  ## A string representing the memory address of `x` in hexadecimal format,
  ## prefixed with "0x". This format is widely recognized and used for representing
  ## memory addresses in debugging and low-level programming contexts.
  ##
  ## Example Usage:
  ## ```
  ## var myVar: int = 42
  ## echo memAddr(myVar)  # Outputs something like "0x7ffdf2c2b6c0"
  ## ```
  ##
  ## *Note:*
  ## Working with raw memory addresses is a low-level operation that should be
  ## approached with caution. This functionality is primarily intended for
  ## debugging or interfacing with low-level system APIs or external C libraries,
  ## and should not generally be used in high-level application logic. Misuse
  ## of memory addresses can lead to security vulnerabilities, memory corruption,
  ## or application crashes.
  "0x" & $(cast[uint](addr x)).toHex(length)

when isMainModule:
  # Example showing how to allocate executable memory for shellcode
  # Translation of C# Code: https://www.fergonez.net/post/shellcode-csharp
  # Example should work in Windows / MacOS / Linux, without chnage
  block shellcodeExample:
    # Example shellcode (x86_64) - this is just a simple calculation routine
    # It Give different result when executed in different environments ?
    const shellcode: array[30, byte] = [
      byte 0xBA, 0xC5, 0x9D, 0x1C, 0x81, 0xEB, 0x0B, 0x0F, 0xBE, 0xC0, 0x33, 0xC2, 0x69,
      0xD0, 0x93, 0x01, 0x00, 0x01, 0x8A, 0x01, 0x48, 0xFF, 0xC1, 0x84, 0xC0, 0x75,
      0xEC, 0x8B, 0xC2, 0xC3,
    ]

    echo "Allocating executable memory for shellcode demonstration..."

    # Allocate memory with execute permissions
    let codeSize = shellcode.len
    try:
      let execMem =
        nVirtualAlloc(nil, codeSize, VM_COMMIT or VM_RESERVE, VM_PAGE_EXECUTE_READWRITE)

      # Copy shellcode to executable memory safely
      copyMem(execMem, unsafeAddr shellcode[0], codeSize)

      echo "Shellcode loaded into executable memory at address: ", execMem.memAddr(16)

      # We prepare the string that we're going to pass to the shellcode function.
      let stringInput = "Fergo\0" # For Safety will use Null Terminated String

      # We declare a type for a function that takes a single cstring
      # (a string in C language) argument and returns a uint32
      # (32-bit unsigned integer).
      type HashFunction = proc(s: cstring): uint32 {.stdcall.}

      # We cast the shellcode address to a HashPrototype function pointer.
      # In simple terms, we're saying, "Treat the code in this memory
      # address as a function that matches our HashPrototype type."
      let hashFunction = cast[HashFunction](execMem)

      # We call the shellcode function, passing it our string.
      # It returns the result of hashing our string, according to the code
      # in the shellcode.
      let hashedResult = hashFunction(stringInput.cstring)

      # We print the result in hexadecimal form.
      echo(fmt"[*] Here is hashed result of {stringInput} in hex: {hashedResult:0X}")

      # Free the executable memory
      try:
        discard nVirtualFree(execMem, codeSize, VM_FREE)
        echo "Executable memory freed successfully"
      except MemoryError as e:
        echo "Failed to free executable memory: ", e.msg
    except MemoryError as e:
      echo "Failed to allocate executable memory: ", e.msg