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


global CallBackSysDynamicDeviceInit
global CallBackSysDynamicDeviceExit
global AllocCallBacks
global SetCB
global SetCBX
global CheckWCRS
global ReinforceINT3
global entryV86CB
global entryPMCB
global entryPMR0
global wcrs
global oINT3.V86
global oINT3.PMR3
global IDT.original


extern sdata
extern StaticEntryV86CB
extern StaticEntryPMCB
extern StaticSimIretCB
extern Service_CDPlayer
extern Service_Dump
extern Service_Load
extern Service_MP3Player
extern Service_Pedump
extern Service_SuspendResumeKill
extern Service_Trace
extern Service_TraceX
extern Service_BreakR3
extern Service_FDump
extern GetSetWiniceInt


bits 32


segment _LDATA
init_data_begin ICEDUMP_CALLBACK
init_data_item	HookWiniceInts,		UnhookWiniceInts
init_data_end


segment _LTEXT
;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
CallBackSysDynamicDeviceInit:
	push	ecx
	push	edi

; init int3 callback addresses
	mov	eax,[wSelector_WINICE_Code]
	sub	[oINT3.PMR3],eax

	mov	edi,0xF0000		; start of ROM BIOS
	mov	ecx,0x10000		; length of area we search through
	mov	al,0xCC			; opcode of INT3
	repnz	scasb			; may cause page fault
	jz	@F

	Trace_Out "ICEDUMP: CallBack failed to find an INT3 for V86 mode"

	pop	edi
	pop	ecx
	stc
	retn

@@
	dec	edi			; adjust EDI
	mov	[oINT3.V86],edi

	pop	edi
	pop	ecx

	init_construct ICEDUMP_CALLBACK
	jc	.error

	clc
	retn

.error:
	call	CallBackSysDynamicDeviceExit

	stc
	retn


;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
CallBackSysDynamicDeviceExit:
	init_destruct ICEDUMP_CALLBACK
	clc
	retn


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

	push	esi

	mov	al,3
	mov	esi,HookedWiniceInt3
	call	GetSetWiniceInt
	mov	[HookedWiniceInt3.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,[HookedWiniceInt3.old]
	call	GetSetWiniceInt
	and	dword [HookedWiniceInt3.old],byte 0

	pop	esi
	popfd
	clc
	retn


;-------------------------------------------------------------------------------
; check if int3 was to return from a callback service, restore context then
;-------------------------------------------------------------------------------
HookedWiniceInt3:
%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

	cmp	dword [wcrs+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	[wcrs+WCRS.oINT3],eax
	jnz	@F

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

; restore client's int3 IDT descriptor
	pushfd
	cli

	push	eax
	sidt	[esp-2]
	pop	eax
	lea	eax,[eax+8*3]		; eax: int3 IDT descriptor offset

	push	dword [IDT.client]
	pop	dword [eax]
	push	dword [IDT.client+4]
	pop	dword [eax+4]

	popfd

	mov	ecx,(WCRS_size-WCRS.dClient_EAX)/4
	mov	esi,wcrs+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
wcrs:	istruc WCRS
	iend


segment _LTEXT
;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
AllocCallBacks:
	push	edx
	push	esi

	xor	edx,edx			; refdata, not used
	cmp	dword [sdata+PMCB],byte 0
	jnz	.gotPMCB		; already allocated

	mov	esi,StaticEntryPMCB
	VMMCall	Allocate_PM_Call_Back
	jnb	@F

	xor	eax,eax			; oops...

@@
	mov	[sdata+PMCB+CallBack.offset],ax
	shr	eax,16
	mov	[sdata+PMCB+CallBack.selector],ax

.gotPMCB:

	xor	edx,edx			; refdata, not used
	cmp	dword [sdata+SimIretCB],byte 0
	jnz	.gotSimIretCB; already allocated

	mov	esi,StaticSimIretCB
	VMMCall	Allocate_PM_Call_Back
	jnb	@F

	xor	eax,eax			; oops...

@@
	mov	[sdata+SimIretCB+CallBack.offset],ax
	shr	eax,16
	mov	[sdata+SimIretCB+CallBack.selector],ax

.gotSimIretCB:
	cmp	dword [sdata+V86CB],byte 0
	jnz	.gotV86CB		; already allocated

	mov	esi,StaticEntryV86CB
	VMMCall	Allocate_V86_Call_Back
	jnb	@F

	xor	eax,eax			; oops...

@@
	mov	[sdata+V86CB+CallBack.offset],ax
	shr	eax,16
	mov	[sdata+V86CB+CallBack.selector],ax

.gotV86CB:
	pop	esi
	pop	edx

	cmp	dword [sdata+PMCB],byte 0
	jz	.error

	cmp	dword [sdata+V86CB],byte 0
	jz	.error

	clc
	retn

.error:
	stc
	retn


;------------------------------------------------------------------------------
; set up client CS:EIP to one of our callbacks
;------------------------------------------------------------------------------
SetCBX:
	push	ebx
	push	ecx
	push	edx

; get V86 CB address
	movzx	ebx,word [sdata+V86CB+CallBack.selector]
	movzx	ecx,word [sdata+V86CB+CallBack.offset]

	mov	edx,[dClient_EFLAGS]	; is client in V86 mode?
	test	byte [edx+2],2
	jnz	.setCSEIP

; get PM CB address
	movzx	ebx,word [sdata+PMCB+CallBack.selector]
	movzx	ecx,word [sdata+PMCB+CallBack.offset]

	mov	edx,[dClient_CS]	; is client in ring-0?
	test	byte [edx],3
	jnz	.setCSEIP

	mov	ecx,entryPMR0
	jmp	short .setEIP

.setCSEIP:
	mov	edx,[dClient_CS]	; set client CS
	mov	[edx],ebx

.setEIP:
	mov	edx,[dClient_EIP]	; set client (E)IP
	mov	[edx],ecx

	pop	edx
	pop	ecx
	pop	ebx
	retn


;------------------------------------------------------------------------------
; set up client CS:EIP to one of our callbacks
;------------------------------------------------------------------------------
SetCB:
	call	CheckWCRS
	call	SetCSEIP
	call	ReinforceINT3
	retn


;------------------------------------------------------------------------------
;
;------------------------------------------------------------------------------
CheckWCRS:
	push	eax
	push	esi

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

	mov	esi,msgForceCB
	call	[pPrintToCommandWindow]

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

	pushfd
	cli

	push	eax
	sidt	[esp-2]
	pop	eax
	lea	eax,[eax+8*3]			; eax: int3 IDT descriptor offset

	push	dword [IDT.client]
	pop	dword [eax]
	push	dword [IDT.client+4]
	pop	dword [eax+4]

	popfd

@@
	pop	esi
	pop	eax
	retn


;------------------------------------------------------------------------------
;
;------------------------------------------------------------------------------
SetCSEIP:
	push	eax
	push	ebx
	push	ecx

; get V86 CB address
	mov	eax,[oINT3.V86]
	movzx	ebx,word [sdata+V86CB+CallBack.selector]
	movzx	ecx,word [sdata+V86CB+CallBack.offset]

	mov	edx,[dClient_EFLAGS]	; is client in V86 mode?
	test	byte [edx+2],2
	jnz	.setCSEIP

; get PM CB address
	mov	eax,[oINT3.PMR3]
	movzx	ebx,word [sdata+PMCB+CallBack.selector]
	movzx	ecx,word [sdata+PMCB+CallBack.offset]

	mov	edx,[dClient_CS]	; is client in ring-0?
	test	byte [edx],3
	jnz	.setCSEIP

	mov	ecx,entryPMR0
	mov	eax,[oINT3.PMR0]
	jmp	short .setEIP

.setCSEIP:
	mov	edx,[dClient_CS]	; set client CS
	mov	[edx],ebx

.setEIP:
	mov	edx,[dClient_EIP]	; set client (E)IP
	mov	[edx],ecx

	mov	[wcrs+WCRS.oINT3],eax	; set address of final INT3
	mov	dword [wcrs+WCRS.valid],1

	pop	ecx
	pop	ebx
	pop	eax
	retn


;------------------------------------------------------------------------------
;
;------------------------------------------------------------------------------
ReinforceINT3:
	push	eax

; reinforce original int3 IDT handler (hopefully that of winice ;-)
	pushfd
	cli

	push	eax
	sidt	[esp-2]
	pop	eax
	lea	eax,[eax+8*3]			; eax: int3 IDT descriptor offset

; save client's int3 IDT descriptor
	push	dword [eax]
	pop	dword [IDT.client]
	push	dword [eax+4]
	pop	dword [IDT.client+4]

; reinforce original int3 IDT descriptor
	push	dword [IDT.original]
	pop	dword [eax]
	push	dword [IDT.original+4]
	pop	dword [eax+4]

	popfd

	pop	eax
	retn


PMINT3:
	int3


segment _LDATA
oINT3:
.V86:	dd 0
.PMR3:	dd PMINT3+4
.PMR0:	dd PMINT3
IDT:
.original:	dd 0,0
.client:	dd 0,0
msgForceCB:	db 'warning: a service was already in progress, prepare for a crash...',0


segment _LTEXT
entryV86CB:
	pushad
	pushfd
	cld
	mov	edi,[oINT3.V86]
	xor	cl,cl
	shld	ecx,edi,28		; ECX: client CS
	mov	ebx,edi			; EBX: client IP
	and	ebx,byte 0x0F
	jmp	short setCSEIP

entryPMCB:
	pushad
	pushfd
	cld
	mov	ecx,[wSelector_WINICE_Code]	; set client CS:EIP
	movzx	ecx,word [ecx]
	mov	ebx,[oINT3.PMR3]		; to an INT3

setCSEIP:
	mov     [ebp+CRS.CS],ecx
	mov	[ebp+CRS.EIP],ebx
	jmp	short service


;-------------------------------------------------------------------------------
entryPMR0:
	pushfd
	push	cs
	push	dword [oINT3.PMR0]	; we'll return to an INT3
	push	dword service.retR0	; 'fake' error code, must NOT be modified
	pushad				; simulate Client Register Structure
	mov	ebp,esp			; on the stack
	pushfd
	cld

service:
;	cmp	dword [ebp+CRS.EAX],byte SERVICE_VERSION
;	jnz	@F
;	mov	dword [ebp+CRS.EAX],ICEDUMP_VERSION
;	popfd
;	popad
;	retn
;@@
	mov	eax,[ebp+CRS.EAX]
	cmp	eax,byte SERVICE_MAX
	jae	.ret

	jmp	[ServiceTable+4*eax]

.ret:
	popfd
	popad
	retn

.retR0:
	iretd


segment _LDATA
	align 4
ServiceTable:
	dd	service.ret
	dd	Service_Dump
	dd	Service_Load
	dd	Service_SuspendResumeKill
	dd	Service_SuspendResumeKill
	dd	Service_SuspendResumeKill
	dd	Service_SuspendResumeKill
	dd	Service_CDPlayer
	dd	Service_MP3Player
	dd	Service_Pedump
	dd	Service_Trace
	dd	Service_TraceX
	dd	Service_BreakR3
	dd	Service_FDump
