%include "util.mac"
%include "vxdn.inc"
%include "icedump.inc"
%include "wiat.inc"


global BreakR3SysDynamicDeviceInit
global BreakR3SysDynamicDeviceExit
global BreakR3ThreadInit
global BreakR3ThreadNotExecuteable
global BreakR3VmInit
global BreakR3VmNotExecuteable
global BreakIn
global Parse_BreakR3
global Service_BreakR3


extern SetCBX
extern sdata
extern Parser.error
extern Parser.errorMsg
;extern Error_V86
;extern Error_PM16
;extern Error_PMR0
extern ParseExpression
extern GetSetWiniceInt
extern oINT3.V86
extern oINT3.PMR3


bits 32


segment _LDATA
init_data_begin ICEDUMP_BREAKR3
init_data_item	AllocateThreadDataSlot,	FreeThreadDataSlot
init_data_item	HookWiniceInts,		UnhookWiniceInts
init_data_end


segment _LTEXT
;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
BreakR3SysDynamicDeviceInit:
	init_construct ICEDUMP_BREAKR3
	jc	.error

	clc
	retn

.error:
	call	BreakR3SysDynamicDeviceExit

	stc
	retn


;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
BreakR3SysDynamicDeviceExit:
	init_destruct ICEDUMP_BREAKR3
	clc
	retn


;-------------------------------------------------------------------------------
; EDI: R0TCB
;
; stc on error
;-------------------------------------------------------------------------------
BreakR3ThreadInit:
	push	ecx

	mov	ecx,[TDS]
	and	dword [ecx+edi],byte 0

	pop	ecx
	clc
	retn


;-------------------------------------------------------------------------------
; EDI: R0TCB
;
; stc on error
;-------------------------------------------------------------------------------
BreakR3ThreadNotExecuteable:
	push	ecx
	push	edx

	mov	ecx,[TDS]
	xor	eax,eax
	xchg	eax,[ecx+edi]
	or	eax,eax
	jz	@F

	VMMCall	_HeapFree, eax, byte 0

@@
	pop	edx
	pop	ecx
	clc
	retn


;-------------------------------------------------------------------------------
; ebx: VMCB
;
; stc on error
;-------------------------------------------------------------------------------
BreakR3VmInit:
	VMMCall	Get_Initial_Thread_Handle
	call	BreakR3ThreadInit
	retn


;-------------------------------------------------------------------------------
; ebx: VMCB
;
; stc on error
;-------------------------------------------------------------------------------
BreakR3VmNotExecuteable:
	VMMCall	Get_Initial_Thread_Handle
	call	BreakR3ThreadNotExecuteable
	retn


;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
AllocateThreadDataSlot:
	push	ebx
	push	ecx
	push	edx
	push	edi

	VMMCall	_AllocateThreadDataSlot
	or	eax,eax
	jz	.error

	mov	[TDS],eax
	VMMCall	Get_Sys_VM_Handle
	VMMCall	Get_Initial_Thread_Handle
	mov	ebx,edi
	mov	ecx,[TDS]

@@
	and	dword [ecx+edi],byte 0
	VMMCall	Get_Next_Thread_Handle
	cmp	ebx,edi
	jnz	@B

	pop	edi
	pop	edx
	pop	ecx
	pop	ebx
	clc
	retn

.error:
	pop	edi
	pop	edx
	pop	ecx
	pop	ebx
	stc
	retn


segment _LDATA
	align 4
TDS:	dd 0


segment _LTEXT
;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
FreeThreadDataSlot:
	push	ebx
	push	ecx
	push	edx
	push	edi

	VMMCall	Get_Sys_VM_Handle
	VMMCall	Get_Initial_Thread_Handle
	mov	ebx,edi

.loop:
	call	BreakR3ThreadNotExecuteable

	VMMCall	Get_Next_Thread_Handle
	cmp	ebx,edi
	jnz	.loop

	VMMCall	_FreeThreadDataSlot, dword [TDS]

	pop	edi
	pop	edx
	pop	ecx
	pop	ebx
	clc
	retn


;-------------------------------------------------------------------------------
; hook winice's int03 handler (tWiniceIntHandlers)
;
; stc on error
;-------------------------------------------------------------------------------
HookWiniceInts:
	pushfd
	cli

	push	esi

	mov	al,3
	mov	esi,HookedWiniceInt03
	call	GetSetWiniceInt
	mov	[HookedWiniceInt03.old],eax
	jc	.0

	pop	esi
	popfd
	clc
	retn

.0:
	pop	esi
	popfd
	stc
	retn


;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
UnhookWiniceInts:
	pushfd
	cli

	push	esi

	mov	al,3
	mov	esi,[HookedWiniceInt03.old]
	call	GetSetWiniceInt
	and	dword [HookedWiniceInt03.old],byte 0

	pop	esi
	popfd
	clc
	retn


;-------------------------------------------------------------------------------
; check if int3 was in response to a breakin attempt, restore context then
;-------------------------------------------------------------------------------
HookedWiniceInt03:
%define EFlags 0x28
%define CS     0x24
%define EIP    0x20

	push	ds
	push	es

	pushfd
	cld

	push	eax
	push	ecx
	push	esi
	push	edi

	mov	ax,0x30
	mov	ds,eax
	mov	es,eax

	VMMCall	Get_Cur_Thread_Handle
	add	edi,[TDS]
	mov	esi,[edi]		; esi: WCRS
	or	esi,esi
	jz	@F

	cmp	dword [esi+WCRS.valid],byte 0
	jz	@F

	mov	eax,[esp+EIP]
	test	[esp+EFlags+2],byte 2
	jz	.PM

	movzx	eax,ax
	movzx	edi,word [esp+CS]
	shl	edi,4
	add	eax,edi

.PM:
	dec	eax			; int3 is trap, not fault
	cmp	[esi+WCRS.oINT3],eax
	jnz	@F

	and	dword [esi+WCRS.valid],byte 0

	mov	ecx,(WCRS_size-WCRS.dClient_EAX)/4
	add	esi,byte WCRS.dClient_EAX
	mov	edi,[dClient_EAX]
	rep	movsd

	call	[pSaveClientRegistersForDisplay]

	pop	edi
	pop	esi
	pop	ecx
	pop	eax
	popfd
	pop	es
	lea	esp,[esp+4]		; do not pop ds

	call	[pRestoreAllClientRegisters]
	lea	esp,[esp+4]		; get rid of old handler

	call	[ss:pWiniceMain]
	iretd

@@
	pop	edi
	pop	esi
	pop	ecx
	pop	eax
	popfd
	pop	es
	pop	ds
	jmp	[ss:.old]

%undef EFlags
%undef CS
%undef EIP


segment _LDATA
	align 4
.old:	dd 0


segment _LTEXT
;-------------------------------------------------------------------------------
; EDI: R0TCB
; EBP: Client Registers
;
; stc on error
;-------------------------------------------------------------------------------
BreakIn:
	pushad
	mov	ecx,[TDS]
	mov	eax,[ecx+edi]			; eax: WCRS
	or	eax,eax
	jnz	@F

	VMMCall	_HeapAllocate, dword WCRS_size, byte HEAPZEROINIT
	mov	ecx,[TDS]
	mov	[ecx+edi],eax
	or	eax,eax
	jnz	@F

	popad
	stc
	retn

@@
	mov	dword [eax+WCRS.valid],1

	push	dword [ebp+CRS.EAX]
	pop	dword [eax+WCRS.dClient_EAX]

	push	dword [ebp+CRS.EBX]
	pop	dword [eax+WCRS.dClient_EBX]

	push	dword [ebp+CRS.ECX]
	pop	dword [eax+WCRS.dClient_ECX]

	push	dword [ebp+CRS.EDX]
	pop	dword [eax+WCRS.dClient_EDX]

	push	dword [ebp+CRS.ESI]
	pop	dword [eax+WCRS.dClient_ESI]

	push	dword [ebp+CRS.EDI]
	pop	dword [eax+WCRS.dClient_EDI]

	push	dword [ebp+CRS.EBP]
	pop	dword [eax+WCRS.dClient_EBP]

	push	dword [ebp+CRS.ESP]
	pop	dword [eax+WCRS.dClient_ESP]

	push	dword [ebp+CRS.EIP]
	pop	dword [eax+WCRS.dClient_EIP]

	push	dword [ebp+CRS.EFlags]
	pop	dword [eax+WCRS.dClient_EFLAGS]

	push	dword [ebp+CRS.CS]
	pop	dword [eax+WCRS.dClient_CS]

	push	dword [ebp+CRS.DS]
	pop	dword [eax+WCRS.dClient_DS]

	push	dword [ebp+CRS.SS]
	pop	dword [eax+WCRS.dClient_SS]

	push	dword [ebp+CRS.ES]
	pop	dword [eax+WCRS.dClient_ES]

	push	dword [ebp+CRS.FS]
	pop	dword [eax+WCRS.dClient_FS]

	push	dword [ebp+CRS.GS]
	pop	dword [eax+WCRS.dClient_GS]

	mov	edx,cr2
	mov	[eax+WCRS.dClient_CR2],edx

	mov	edx,cr3
	mov	[eax+WCRS.dClient_CR3],edx

	test	byte [ebp+CRS.EFlags+2],2	; is client in V86 mode
	jnz	.V86

	mov	ecx,[wSelector_WINICE_Code]	; set client CS:EIP
	movzx	ecx,word [ecx]
	mov	edi,[oINT3.PMR3]		; to an INT3
	mov	[eax+WCRS.oINT3],edi		; store address of our INT3
	jmp	short .setCSEIP

.V86:
	mov	edi,[oINT3.V86]
	mov	[eax+WCRS.oINT3],edi		; store address of our INT3
	xor	cl,cl
	shld	ecx,edi,28			; ECX: client CS
	and	edi,byte 0x0F			; EDI: client IP

.setCSEIP:
	mov	[ebp+CRS.CS],ecx
	mov	[ebp+CRS.EIP],edi

	popad
	clc
	retn


;-------------------------------------------------------------------------------
; BREAKR3 <R0TCB>
;-------------------------------------------------------------------------------
Parse_BreakR3:
	mov	edi,Error_BadR0TCB
	call	ParseExpression			; parse <R0TCB>
	jc	near Parser.errorMsg

	mov	[.R0TCB],eax

; save client EAX/CS/EIP
	mov	ebp,[dClient_EAX]
	push	dword [ebp]
	pop	dword [.EAX]

	mov	ebp,[dClient_CS]
	push	dword [ebp]
	pop	dword [.CS]

	mov	ebp,[dClient_EIP]
	push	dword [ebp]
	pop	dword [.EIP]

; set up registers for service
	push	byte SERVICE_BREAKR3
	mov	ebp,[dClient_EAX]
	pop	dword [ebp]

	call	SetCBX

	mov	ebp,[fExecuteMoreCommands]	; set internal Winice flag to 0
	mov	byte [ebp],0

	popad
	retn


segment _LDATA
	align 4
.EAX:		dd 0
.CS:		dd 0
.EIP:		dd 0
.R0TCB:		dd 0
Error_BadR0TCB:	db 'no R0TCB specified',0


segment _LTEXT
Service_BreakR3:
	push	dword [Parse_BreakR3.EAX]
	pop	dword [ebp+CRS.EAX]

	push	dword [Parse_BreakR3.CS]
	pop	dword [ebp+CRS.CS]

	push	dword [Parse_BreakR3.EIP]
	pop	dword [ebp+CRS.EIP]

	mov	edi,[Parse_BreakR3.R0TCB]
	VMMCall	Validate_Thread_Handle
	jnc	@F

	Trace_Out "ICEDUMP: invalid thread handle, R0TCB: #edi"

	jmp	short .ret

@@
	mov	ebp,[edi+TCB_ClientPtr]
	call	BreakIn

.ret:
	popfd
	popad
	retn
