drgn/tests/__init__.py

import functools
from typing import Any, NamedTuple, Optional
import unittest

from drgn import (
    Architecture,
    FindObjectFlags,
    Object,
    Platform,
    PlatformFlags,
    Program,
    Type,
    TypeKind,
    class_type,
    enum_type,
    float_type,
    int_type,
    struct_type,
    typedef_type,
    union_type,
)


coord_type = class_type('coord', 12, (
    (int_type('int', 4, True), 'x', 0),
    (int_type('int', 4, True), 'y', 32),
    (int_type('int', 4, True), 'z', 64),
))
point_type = struct_type('point', 8, (
    (int_type('int', 4, True), 'x', 0),
    (int_type('int', 4, True), 'y', 32),
))
line_segment_type = struct_type('line_segment', 16, (
    (point_type, 'a'),
    (point_type, 'b', 64),
))
option_type = union_type('option', 4, (
    (int_type('int', 4, True), 'i'),
    (float_type('float', 4), 'f'),
))
color_type = enum_type('color', int_type('unsigned int', 4, False),
                       (('RED', 0), ('GREEN', 1), ('BLUE', 2)))
pid_type = typedef_type('pid_t', int_type('int', 4, True))


MOCK_32BIT_PLATFORM = Platform(Architecture.UNKNOWN,
                               PlatformFlags.IS_LITTLE_ENDIAN)
MOCK_PLATFORM = Platform(Architecture.UNKNOWN,
                         PlatformFlags.IS_64_BIT |
                         PlatformFlags.IS_LITTLE_ENDIAN)


class MockMemorySegment(NamedTuple):
    buf: bytes
    virt_addr: Optional[int] = None
    phys_addr: Optional[int] = None


def mock_memory_read(data, address, count, offset, physical):
    return data[offset:offset + count]


class MockObject(NamedTuple):
    name: str
    type: Type
    address: Optional[int] = None
    value: Any = None


def mock_program(platform=MOCK_PLATFORM, *, segments=None, types=None,
                 objects=None):
    def mock_find_type(kind, name, filename):
        if filename:
            return None
        for type in types:
            if type.kind == kind:
                try:
                    type_name = type.name
                except AttributeError:
                    try:
                        type_name = type.tag
                    except AttributeError:
                        continue
                if type_name == name:
                    return type
        return None

    def mock_object_find(prog, name, flags, filename):
        if filename:
            return None
        for obj in objects:
            if obj.name == name:
                if obj.value is not None:
                    if flags & FindObjectFlags.CONSTANT:
                        break
                elif obj.type.kind == TypeKind.FUNCTION:
                    if flags & FindObjectFlags.FUNCTION:
                        break
                elif flags & FindObjectFlags.VARIABLE:
                    break
        else:
            return None
        return Object(prog, obj.type, address=obj.address, value=obj.value)

    prog = Program(platform)
    if segments is not None:
        for segment in segments:
            if segment.virt_addr is not None:
                prog.add_memory_segment(
                    segment.virt_addr, len(segment.buf),
                    functools.partial(mock_memory_read, segment.buf))
            if segment.phys_addr is not None:
                prog.add_memory_segment(
                    segment.phys_addr, len(segment.buf),
                    functools.partial(mock_memory_read, segment.buf), True)
    if types is not None:
        prog.add_type_finder(mock_find_type)
    if objects is not None:
        prog.add_object_finder(mock_object_find)
    return prog


class ObjectTestCase(unittest.TestCase):
    def setUp(self):
        super().setUp()
        self.prog = mock_program()
        # For testing, we want to compare the raw objects rather than using the
        # language's equality operator.
        def object_equality_func(a, b, msg=None):
            if a.prog_ is not b.prog_:
                raise self.failureException(msg or 'objects have different program')
            if a.type_ != b.type_:
                raise self.failureException(msg or f'object types differ: {a.type_!r} != {b.type_!r}')
            if a.address_ != b.address_:
                a_address = 'None' if a.address_ is None else hex(a.address_)
                b_address = 'None' if b.address_ is None else hex(b.address_)
                raise self.failureException(msg or f'object addresses differ: {a_address} != {b_address}')
            if a.byteorder_ != b.byteorder_:
                raise self.failureException(msg or f'object byteorders differ: {a.byteorder_} != {b.byteorder_}')
            if a.bit_offset_ != b.bit_offset_:
                raise self.failureException(msg or f'object bit offsets differ: {a.bit_offset_} != {b.bit_offset_}')
            if a.bit_field_size_ != b.bit_field_size_:
                raise self.failureException(msg or f'object bit field sizes differ: {a.bit_field_size_} != {b.bit_field_size_}')
            exc_a = exc_b = False
            try:
                value_a = a.value_()
            except Exception:
                exc_a = True
            try:
                value_b = b.value_()
            except Exception:
                exc_b = True
            if exc_a and not exc_b:
                raise self.failureException(msg or f'exception raised while reading {a!r}')
            if not exc_a and exc_b:
                raise self.failureException(msg or f'exception raised while reading {b!r}')
            if not exc_a and value_a != value_b:
                raise self.failureException(msg or f'object values differ: {value_a!r} != {value_b!r}')
        self.addTypeEqualityFunc(Object, object_equality_func)

    def bool(self, value):
        return Object(self.prog, '_Bool', value=value)

    def int(self, value):
        return Object(self.prog, 'int', value=value)

    def unsigned_int(self, value):
        return Object(self.prog, 'unsigned int', value=value)

    def long(self, value):
        return Object(self.prog, 'long', value=value)

    def double(self, value):
        return Object(self.prog, 'double', value=value)