/*++

Copyright (c) 1995-1998 Microsoft Corporation

Module Name: 

    suspend.c

Abstract:
    
    This module implements CpuSuspendThread, CpuGetContext and CpuSetContext.

Author:

    16-Dec-1999  SamerA

Revision History:

--*/

#define _WOW64CPUAPI_

#ifdef _X86_
#include "ia6432.h"
#else
#define _NTDDK_
#include <nt.h>
#include <ntrtl.h>
#include <nturtl.h>
#include <ntos.h>
#include "wow64.h"
#include "wow64cpu.h"
#include "ia64cpu.h"
#endif

#include "cpup.h"

#include <stdio.h>
#include <stdarg.h>


ASSERTNAME;

ULONG_PTR ia32ShowContext = 0;

VOID
CpupDebugPrint(
    IN ULONG_PTR Flags,
    IN PCHAR Format,
    ...)
{
    va_list ArgList;
    int BytesWritten;
    CHAR Buffer[ 512 ];

    if ((ia32ShowContext & Flags) || (Flags == ERRORLOG))
    {
        va_start(ArgList, Format);
        BytesWritten = _vsnprintf(Buffer,
                                  sizeof(Buffer) - 1,
                                  Format,
                                  ArgList);
        if (BytesWritten > 0)
        {
            DbgPrint(Buffer);
        }
        va_end(ArgList);
    }
    
    return;
}


VOID
CpupPrintContext(
    IN PCHAR str,
    IN PCPUCONTEXT cpu
    )
/*++

Routine Description:

    Print out the ia32 context based on the passed in cpu context

Arguments:

    str - String to print out as a header
    cpu - Pointer to the per-thread wow64 ia32 context.

Return Value:

    none

--*/
{
    DbgPrint(str);
    DbgPrint("Context addr(0x%p): EIP=0x%08x\n", &(cpu->Context), cpu->Context.Eip);
    DbgPrint("Context EAX=0x%08x, EBX=0x%08x, ECX=0x%08x, EDX=0x%08x\n",
                        cpu->Context.Eax,
                        cpu->Context.Ebx,
                        cpu->Context.Ecx,
                        cpu->Context.Edx);
    DbgPrint("Context ESP=0x%08x, EBP=0x%08x, ESI=0x%08x, EDI=0x%08x\n",
                        cpu->Context.Esp,
                        cpu->Context.Ebp,
                        cpu->Context.Esi,
                        cpu->Context.Edi);
    try {
        //
        // The stack may not yet be fully formed, so don't
        // let a missing stack cause the process to abort
        //
        DbgPrint("Context stack=0x%08x 0x%08x 0x%08x 0x%08x\n",
                        *((PULONG) cpu->Context.Esp),
                        *(((PULONG) cpu->Context.Esp) + 1),
                        *(((PULONG) cpu->Context.Esp) + 2),
                        *(((PULONG) cpu->Context.Esp) + 3));
    }
    except ((GetExceptionCode() == STATUS_ACCESS_VIOLATION)?1:0) {
        //
        // Got an access violation, so don't print any of the stack
        //
        DbgPrint("Context stack: Can't get stack contents\n");
    }

    DbgPrint("Context EFLAGS=0x%08x\n", cpu->Context.EFlags);
}


NTSTATUS CpupReadBuffer(
    IN HANDLE ProcessHandle,
    IN PVOID Source,
    OUT PVOID Destination,
    IN ULONG Size)
/*++

Routine Description:

    This routine setup the arguments for the remoted  SuspendThread call.
    
Arguments:

    ProcessHandle  - Target process handle to read data from
    Source         - Target base address to read data from
    Destination    - Address of buffer to receive data read from the specified address space
    Size           - Size of data to read

Return Value:

    NTSTATUS.

--*/
{
    NTSTATUS NtStatus = STATUS_SUCCESS;

    if (ProcessHandle == NtCurrentProcess ()) {

        try {

            RtlCopyMemory (Destination,
                           Source,
                           Size);

        } except (EXCEPTION_EXECUTE_HANDLER) {
            NtStatus = GetExceptionCode ();
        }

    } else {
        
        NtStatus = NtReadVirtualMemory (ProcessHandle,
                                        Source,
                                        Destination,
                                        Size,
                                        NULL);
    }

    return NtStatus;
}

NTSTATUS
CpupWriteBuffer(
    IN HANDLE ProcessHandle,
    IN PVOID Target,
    IN PVOID Source,
    IN ULONG Size)
/*++

Routine Description:

    Writes data to memory taken into consideration if the write is cross-process
    or not
    
Arguments:

    ProcessHandle  - Target process handle to write data into
    Target         - Target base address to write data at
    Source         - Address of contents to write in the specified address space
    Size           - Size of data to write
    
Return Value:

    NTSTATUS.

--*/
{
    NTSTATUS NtStatus = STATUS_SUCCESS;

    if (ProcessHandle == NtCurrentProcess ()) {

        try {

            RtlCopyMemory (Target,
                           Source,
                           Size);

        } except (EXCEPTION_EXECUTE_HANDLER) {
            NtStatus = GetExceptionCode ();
        }

    } else {
    
        NtStatus = NtWriteVirtualMemory (ProcessHandle,
                                         Target,
                                         Source,
                                         Size,
                                         NULL);
    }

    return NtStatus;
}

NTSTATUS
GetContextRecord(
    IN PCPUCONTEXT cpu,
    IN OUT PCONTEXT32 Context
    )
/*++

Routine Description:

    Retrevies the context record of the specified CPU

Arguments:

    cpu     - CPU to retreive the context record for.
    Context - IN/OUT pointer to CONTEXT32 to fill in.  Context->ContextFlags
              should be used to determine how much of the context to copy.

Return Value:

    None.

--*/

{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    ULONG ContextFlags;
    

    try 
    {
        ContextFlags = Context->ContextFlags;
        if (ContextFlags & CONTEXT_IA64)
        {
            LOGPRINT((ERRORLOG, "CpuGetContext: Request for ia64 context (0x%x) being FAILED\n", ContextFlags));
            ASSERT((ContextFlags & CONTEXT_IA64) == 0);
        }

        if ((ContextFlags & CONTEXT32_CONTROL) == CONTEXT32_CONTROL) 
        {
            //
            // i386 control registers are:
            // ebp, eip, cs, eflag, esp and ss
            //
            Context->Ebp = cpu->Context.Ebp;
            Context->Eip = cpu->Context.Eip;
            Context->SegCs = KGDT_R3_CODE|3;   // Force reality
            Context->EFlags = SANITIZE_X86EFLAGS(cpu->Context.EFlags);
            Context->Esp = cpu->Context.Esp;
            Context->SegSs = KGDT_R3_DATA|3;   // Force reality
        }

        if ((ContextFlags & CONTEXT32_INTEGER)  == CONTEXT32_INTEGER)
        {
            //
            // i386 integer registers are:
            // edi, esi, ebx, edx, ecx, eax
            //
            Context->Edi = cpu->Context.Edi;
            Context->Esi = cpu->Context.Esi;
            Context->Ebx = cpu->Context.Ebx;
            Context->Edx = cpu->Context.Edx;
            Context->Ecx = cpu->Context.Ecx;
            Context->Eax = cpu->Context.Eax;
        }

        if ((ContextFlags & CONTEXT32_SEGMENTS) == CONTEXT32_SEGMENTS) 
        {
            //
            // i386 segment registers are:
            // ds, es, fs, gs
            // And since they are a constant, force them to be the right values
            //
            Context->SegDs = KGDT_R3_DATA|3;
            Context->SegEs = KGDT_R3_DATA|3;
            Context->SegFs = KGDT_R3_TEB|3;
            Context->SegGs = 0;
        }

        if ((ContextFlags & CONTEXT32_EXTENDED_REGISTERS) == CONTEXT32_EXTENDED_REGISTERS) 
        {
            //
            // Point to the destination area
            //
            PFXSAVE_FORMAT_WX86 xmmi = (PFXSAVE_FORMAT_WX86) &(Context->ExtendedRegisters[0]);

            LOGPRINT((TRACELOG, "CpuGetContext: Request to get Katmai registers(0x%x)\n", ContextFlags));
            
            RtlCopyMemory(xmmi, &(cpu->Context.ExtendedRegisters[0]),
                          MAXIMUM_SUPPORTED_EXTENSION);
            //
            // For performance reasons, the PCPU context has the
            // fp registers un-rotated. So we need to rotate them 
            // to get them into the standard FXSAVE format
            //
            Wow64RotateFpTop(xmmi->StatusWord, (PFLOAT128) &(xmmi->RegisterArea[0]));
        }

        if ((ContextFlags & CONTEXT32_FLOATING_POINT) == CONTEXT32_FLOATING_POINT) 
        {
            //
            // For the ISA transition routine, these floats are
            // in the ExtendedRegister area. So grab the values requested
            // from that area
            //

            //
            // Point to the source area
            //
            PFXSAVE_FORMAT_WX86 xmmi = (PFXSAVE_FORMAT_WX86) &(cpu->Context.ExtendedRegisters[0]);

            //
            // Need space to rotate the registers
            //
            FLOAT128 tmpFloat[NUMBER_OF_387REGS];

            LOGPRINT((TRACELOG, "CpuGetContext: Request to get float registers(0x%x)\n", ContextFlags));

            //
            // Start by grabbing the status/control portion
            //
            Context->FloatSave.ControlWord = xmmi->ControlWord;
            Context->FloatSave.StatusWord = xmmi->StatusWord;
            Context->FloatSave.TagWord = xmmi->TagWord;
            Context->FloatSave.ErrorOffset = xmmi->ErrorOffset;
            Context->FloatSave.ErrorSelector = xmmi->ErrorSelector;
            Context->FloatSave.DataOffset = xmmi->DataOffset;
            Context->FloatSave.DataSelector = xmmi->DataSelector;

            //
            // Don't touch the original PCPU context. Make a copy.
            //
            RtlCopyMemory(tmpFloat, xmmi->RegisterArea,
                   NUMBER_OF_387REGS * sizeof(FLOAT128));

            //
            // For performance reasons, the PCPU context leaves the
            // fp registers un-rotated. So we need to rotate them now
            // to make it follow the proper FSAVE fotmat
            //
            Wow64RotateFpTop(xmmi->StatusWord, tmpFloat);

            //
            // Now get the packed 10-byte fp data registers
            //
            Wow64CopyFpFromIa64Byte16(tmpFloat,
                                      &(Context->FloatSave.RegisterArea[0]),
                                      NUMBER_OF_387REGS);
        }

        if ((ContextFlags & CONTEXT32_DEBUG_REGISTERS) == CONTEXT32_DEBUG_REGISTERS) 
        {
            LOGPRINT((TRACELOG, "CpuGetContext: Request to get debug registers(0x%x)\n", ContextFlags));
            Context->Dr0 = cpu->Context.Dr0;
            Context->Dr1 = cpu->Context.Dr1;
            Context->Dr2 = cpu->Context.Dr2;
            Context->Dr3 = cpu->Context.Dr3;
            Context->Dr6 = cpu->Context.Dr6;
            Context->Dr7 = cpu->Context.Dr7;
        }
    }
    except(EXCEPTION_EXECUTE_HANDLER)
    {
        NtStatus = GetExceptionCode();
    }

    if (ia32ShowContext & LOG_CONTEXT_GETSET) 
    {
        CpupPrintContext("Getting ia32 context: ", cpu);
    }

    return NtStatus;
}

NTSTATUS
CpupGetContext(
    IN OUT PCONTEXT32 Context
    )
/*++

Routine Description:

    This routine extracts the context record for the currently executing thread. 

Arguments:

    Context  - Context record to fill

Return Value:

    NTSTATUS.

--*/
{
    DECLARE_CPU;

    return GetContextRecord(cpu, Context);
}


NTSTATUS
CpupGetContextThread(
    IN HANDLE ThreadHandle,
    IN HANDLE ProcessHandle,
    IN PTEB Teb,
    IN OUT PCONTEXT32 Context)
/*++

Routine Description:

    This routine extract the context record of any thread. This is a generic routine.
    When entered, if the target thread isn't the current thread, then it should be 
    guaranteed that the target thread is suspended at a proper CPU state.

Arguments:

    ThreadHandle   - Target thread handle to retreive the context for
    ProcessHandle  - Open handle to the process that the thread runs in
    Teb            - Pointer to the target's thread TEB
    Context        - Context record to fill                 

Return Value:

    NTSTATUS.

--*/
{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    CONTEXT ContextEM;
    PCPUCONTEXT CpuRemoteContext;
    CPUCONTEXT CpuContext;


    ContextEM.ContextFlags = CONTEXT_FULL | CONTEXT_DEBUG;
    NtStatus = NtGetContextThread(ThreadHandle,
                                  &ContextEM);

    if (!NT_SUCCESS(NtStatus))
    {
        LOGPRINT((ERRORLOG, "CpupGetContextThread: NtGetContextThread (%lx) failed - %lx\n", 
                  ThreadHandle, NtStatus));
        return NtStatus;
    }

    if (ContextEM.StIPSR & (1i64 << PSR_IS))
    {
        Wow64CtxFromIa64(Context->ContextFlags, &ContextEM, Context);
                    
        LOGPRINT((TRACELOG, "Getting context while thread is executing 32-bit instructions - %lx\n", NtStatus));
    }
    else
    {
        LOGPRINT((TRACELOG, "Getting context while thread is executing 64-bit instructions\n"));
        NtStatus = CpupReadBuffer(ProcessHandle,
                                  ((PCHAR)Teb + FIELD_OFFSET(TEB, TlsSlots[WOW64_TLS_CPURESERVED])),
                                  &CpuRemoteContext,
                                  sizeof(CpuRemoteContext));

        if (NT_SUCCESS(NtStatus))
        {
            NtStatus = CpupReadBuffer(ProcessHandle,
                                      CpuRemoteContext,
                                      &CpuContext,
                                      sizeof(CpuContext));

            if (NT_SUCCESS(NtStatus))
            {
                NtStatus = GetContextRecord(&CpuContext, Context);
            }
            else
            {
                LOGPRINT((ERRORLOG, "CpupGetContextThread: Couldn't read CPU context %lx - %lx\n", 
                          CpuRemoteContext, NtStatus));

            }
        }
        else
        {
            LOGPRINT((ERRORLOG, "CpupGetContextThread: Couldn't read CPU context address - %lx\n", 
                      NtStatus));
        }
    }

    return NtStatus;
}



NTSTATUS
SetContextRecord(
    IN OUT PCPUCONTEXT cpu,
    IN PCONTEXT32 Context
    )
/*++

Routine Description:

    Update the CPU's register set for the specified CPU.

Arguments:

    cpu     - CPU to update its registers
    Context - IN pointer to CONTEXT32 to use.  Context->ContextFlags
              should be used to determine how much of the context to update.

Return Value:

    None.

--*/
{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    ULONG ContextFlags;

    
    try 
    {
        ContextFlags = Context->ContextFlags;
        if (ContextFlags & CONTEXT_IA64) 
        {
            LOGPRINT((ERRORLOG, "CpuSetContext: Request with ia64 context (0x%x) FAILED\n", ContextFlags));
            ASSERT((ContextFlags & CONTEXT_IA64) == 0);
        }

        if ((ContextFlags & CONTEXT32_CONTROL) == CONTEXT32_CONTROL) 
        {
            //
            // i386 control registers are:
            // ebp, eip, cs, eflag, esp and ss
            //
            cpu->Context.Ebp = Context->Ebp;
            cpu->Context.Eip = Context->Eip;
            cpu->Context.SegCs = KGDT_R3_CODE|3;   // Force Reality
            cpu->Context.EFlags = SANITIZE_X86EFLAGS(Context->EFlags);
            cpu->Context.Esp = Context->Esp;
            cpu->Context.SegSs = KGDT_R3_DATA|3;   // Force Reality
        }

        if ((ContextFlags & CONTEXT32_INTEGER)  == CONTEXT32_INTEGER)
        {
            //
            // i386 integer registers are:
            // edi, esi, ebx, edx, ecx, eax
            //
            cpu->Context.Edi = Context->Edi;
            cpu->Context.Esi = Context->Esi;
            cpu->Context.Ebx = Context->Ebx;
            cpu->Context.Edx = Context->Edx;
            cpu->Context.Ecx = Context->Ecx;
            cpu->Context.Eax = Context->Eax;
        }

        if ((ContextFlags & CONTEXT32_SEGMENTS) == CONTEXT32_SEGMENTS) 
        {
            //
            // i386 segment registers are:
            // ds, es, fs, gs
            // And since they are a constant, force them to be the right values
            //
            cpu->Context.SegDs = KGDT_R3_DATA|3;
            cpu->Context.SegEs = KGDT_R3_DATA|3;
            cpu->Context.SegFs = KGDT_R3_TEB|3;
            cpu->Context.SegGs = 0;
        }

        //
        // To follow the way ia32 does get/set context, you need to make sure
        // that the older FP context is saved second. That way if both
        // old and new context is passed in, the old takes precedence
        // This happens, for example, when handling an FP exception... The
        // exception handler says both context is available, and older programs
        // only clean up the older FP area...
        //
        if ((ContextFlags & CONTEXT32_EXTENDED_REGISTERS) == CONTEXT32_EXTENDED_REGISTERS) 
        {
            //
            // Point to the destination
            //
            PFXSAVE_FORMAT_WX86 xmmi = (PFXSAVE_FORMAT_WX86) &(cpu->Context.ExtendedRegisters[0]);

            LOGPRINT((TRACELOG, "CpuSetContext: Request to set Katmai registers(0x%x)\n", ContextFlags));

            RtlCopyMemory(xmmi, &(Context->ExtendedRegisters[0]),
                          MAXIMUM_SUPPORTED_EXTENSION);
            //
            // For performance reasons, the PCPU context leaves the
            // fp registers un-rotated. So we need to rotate them back
            // now into the optimized format used for isa transisions
            //
            {
                ULONGLONG RotateFSR = (NUMBER_OF_387REGS -
                                       ((xmmi->StatusWord >> 11) & 0x7)) << 11;
                Wow64RotateFpTop(RotateFSR, (PFLOAT128) &(xmmi->RegisterArea[0])
);
            }
        }

        if ((ContextFlags & CONTEXT32_FLOATING_POINT) == CONTEXT32_FLOATING_POINT) 
        {
            //
            // For the ISA transition routine, these floats need to be
            // in the ExtendedRegister area. So put the values requested
            // into that area
            //
            PFXSAVE_FORMAT_WX86 xmmi = (PFXSAVE_FORMAT_WX86) &(cpu->Context.ExtendedRegisters[0]);

            LOGPRINT((TRACELOG, "CpuSetContext: Request to set float registers(0x%x)\n", ContextFlags));

            //
            // Start by grabbing the status/control portion
            //
            xmmi->ControlWord = (USHORT) (Context->FloatSave.ControlWord & 0xFFFF);
            xmmi->StatusWord = (USHORT) (Context->FloatSave.StatusWord & 0xFFFF);
            xmmi->TagWord = (USHORT) (Context->FloatSave.TagWord & 0xFFFF);
            xmmi->ErrorOffset = Context->FloatSave.ErrorOffset;
            xmmi->ErrorSelector = Context->FloatSave.ErrorSelector;
            xmmi->DataOffset = Context->FloatSave.DataOffset;
            xmmi->DataSelector = Context->FloatSave.DataSelector;

            //
            // Now get the packed 10-byte fp data registers and convert
            // them into the 16-byte format used by FXSAVE (and the
            // ISA transition routine)
            //
            Wow64CopyFpToIa64Byte16(&(Context->FloatSave.RegisterArea[0]),
                                    &(xmmi->RegisterArea[0]),
                                    NUMBER_OF_387REGS);

            //
            // For performance reasons, the PCPU context leaves the
            // fp registers un-rotated. So we need to rotate them back
            // now into the optimized format used for isa transisions
            //
            {
                ULONGLONG RotateFSR = (NUMBER_OF_387REGS -
                                       ((xmmi->StatusWord >> 11) & 0x7)) << 11;
                Wow64RotateFpTop(RotateFSR, (PFLOAT128) &(xmmi->RegisterArea[0]));
            }
        }

        if ((ContextFlags & CONTEXT32_DEBUG_REGISTERS) == CONTEXT32_DEBUG_REGISTERS)
        {
            LOGPRINT((TRACELOG, "CpuSetContext: Request to set debug registers(0x%x)\n", ContextFlags));
            cpu->Context.Dr0 = Context->Dr0;
            cpu->Context.Dr1 = Context->Dr1;
            cpu->Context.Dr2 = Context->Dr2;
            cpu->Context.Dr3 = Context->Dr3;
            cpu->Context.Dr6 = Context->Dr6;
            cpu->Context.Dr7 = Context->Dr7;
        }

        //
        // Whatever they passed in before, it's an X86 context now...
        //
        cpu->Context.ContextFlags = ContextFlags;
    }
    except(EXCEPTION_EXECUTE_HANDLER)
    {
        NtStatus = GetExceptionCode();
    }


    if (ia32ShowContext & LOG_CONTEXT_GETSET) 
    {
        CpupPrintContext("Setting ia32 context: ", cpu);
    }

    return NtStatus;
}

NTSTATUS
CpupSetContext(
    IN PCONTEXT32 Context
    )
/*++

Routine Description:

    This routine sets the context record for the currently executing thread. 

Arguments:

    Context  - Context record to fill 

Return Value:

    NTSTATUS.

--*/
{
    DECLARE_CPU;

    return SetContextRecord(cpu, Context);
}



NTSTATUS
CpupSetContextThread(
    IN HANDLE ThreadHandle,
    IN HANDLE ProcessHandle,
    IN PTEB Teb,
    IN OUT PCONTEXT32 Context)
/*++

Routine Description:

    This routine sets the context record of any thread. This is a generic routine.
    When entered, if the target thread isn't the currently executing thread, then it should be 
    guaranteed that the target thread is suspended at a proper CPU state.

Arguments:

    ThreadHandle   - Target thread handle to retreive the context for
    ProcessHandle  - Open handle to the process that the thread runs in
    Teb            - Pointer to the target's thread TEB
    Context        - Context record to set

Return Value:

    NTSTATUS.

--*/
{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    CONTEXT ContextEM;
    PCPUCONTEXT CpuRemoteContext;
    CPUCONTEXT CpuContext;

    ContextEM.ContextFlags = CONTEXT_FULL | CONTEXT_DEBUG;
    NtStatus = NtGetContextThread(ThreadHandle,
                                  &ContextEM);

    if (!NT_SUCCESS(NtStatus))
    {
        LOGPRINT((ERRORLOG, "CpupGetContextThread: NtGetContextThread (%lx) failed - %lx\n", 
                  ThreadHandle, NtStatus));

        return NtStatus;
    }

    if (ContextEM.StIPSR & (1i64 << PSR_IS))
    {
        Wow64CtxToIa64(Context->ContextFlags, Context, &ContextEM);
        NtStatus = NtSetContextThread(ThreadHandle, &ContextEM);
        LOGPRINT((TRACELOG, "Setting context while thread is executing 32-bit instructions - %lx\n", NtStatus));
    }
    else
    {
        LOGPRINT((TRACELOG, "Setting context while thread is executing 64-bit instructions\n"));
        NtStatus = CpupReadBuffer(ProcessHandle,
                                  ((PCHAR)Teb + FIELD_OFFSET(TEB, TlsSlots[WOW64_TLS_CPURESERVED])),
                                  &CpuRemoteContext,
                                  sizeof(CpuRemoteContext));

        if (NT_SUCCESS(NtStatus))
        {
            NtStatus = CpupReadBuffer(ProcessHandle,
                                      CpuRemoteContext,
                                      &CpuContext,
                                      sizeof(CpuContext));

            if (NT_SUCCESS(NtStatus))
            {    
                NtStatus = SetContextRecord(&CpuContext, Context);

                if (NT_SUCCESS(NtStatus))
                {
                    NtStatus = CpupWriteBuffer(ProcessHandle,
                                               CpuRemoteContext,
                                               &CpuContext,
                                               sizeof(CpuContext));
                }
                else
                {
                    LOGPRINT((ERRORLOG, "CpupSetContextThread: Couldn't read CPU context %lx - %lx\n", 
                              CpuRemoteContext, NtStatus));
                }
            }
        }
        else
        {
            LOGPRINT((ERRORLOG, "CpupSetContextThread: Couldn't read CPU context address - %lx\n", 
                      NtStatus));

        }
    }

    return NtStatus;
}