; SafeFunction.s
;  This program patches the SetFunction() function. This tries to solve problems
;  when two programs try to patch the same function (the system may crash if
;  the programs remove their patch in the wrong order).
; This version is based on the program named SaferPatches which was placed in
; the public domain by the author:
;	Martin Adrian
;	Rullharvsgatan 3A
;	S-431 47 Mölndal
;	SWEDEN
; This version was converted to assembly language for best efficiency, and tries
; to handle better memory failures: SetFunction() is not allowed to fail, but I
; have to make some memory allocations. The problem is solved by throwing an
; "emergency block" into the free memory list on memory shortage. Moreover, all
; allocations are done with pooled allocation functions, in order to reduce
; memory fragmentation. Author:
;	Frédéric Delacroix
;	5 rue d'Artres
;	59269 Quérénaing
;	FRANCE
; This program and the source is fully public domain.

;  version 1.0: Assembly version of Martin Adrian's program.
;  version 1.1: semaphore added to public list for easy rendez-vous.
;		No longer startable from Workbench (the best place is
;		right after SetPatch in the startup-sequence).
;  version 1.2: Enclosed critical parts in Disable()/Enable() (in case
;		an interrupt would try to execute a function whose patch
;		is not yet correctly installed). Now clears the CPU caches.
;  version 1.3: Added a low-memory handler to throw an emergency block of
;		free memory into the system pool on a memory allocation
;		failure during SetFunction(). This reduces the threat of
;		an unexpected NULL return value from SetFunction()...
;  version 1.4: All memory allocations are now done via pools, in order to
;		reduce memory fragmentation.
;  version 1.5: Some minor guide updates, wrote a french documentation.

VERSION	MACRO
	dc.b	'1.5'
	ENDM

	section	Startup,CODE

EMERGENCY_BLOCKSIZE	EQU	32
MEMPUDDLE_SIZE	EQU	32

Start	move.l	4.w,a6
	lea	Variables,a5
	move.l	#20,d7	; return code
	lea	DOS.Name(pc),a1
	moveq	#39,d0	; kickstart >= 3.0 required
	jsr	_LVOOpenLibrary(a6)
	move.l	d0,a4
	move.l	a4,d0
	beq	NoDOS
	move.l	#Bienvenue.MSG,d1
	exg.l	a4,a6
	jsr	_LVOPutStr(a6)
	exg.l	a4,a6
	jsr	_LVOForbid(a6)
	lea	Semaphore.Name,a1
	jsr	_LVOFindSemaphore(a6)
	tst.l	d0
	bne	AlreadyInstalled
	move.l	#MEMF_PUBLIC!MEMF_CLEAR,d0
	move.l	#MEMPUDDLE_SIZE,d1
	move.l	d1,d2
	jsr	_LVOCreatePool(a6)
	move.l	a0,var_MemoryPool(a5)
	beq	NoPool
	suba.l	a0,a0
	jsr	MakeHead
	move.l	d0,var_MainHeader(a5)
	beq.s	NoMemory
	move.l	#_LVOSetFunction,a0
	move.l	#NewSetFunction,d0
	move.l	a6,a1
	jsr	_LVOSetFunction(a6)
	move.l	d0,var_OldSetFunc(a5)
	beq.s	NoSetFunction
	lea	LowMemHandler.Interrupt,a1
	jsr	_LVOAddMemHandler(a6)
	move.l	#EMERGENCY_BLOCKSIZE,d0
	move.l	#MEMF_PUBLIC,d1
	jsr	_LVOAllocMem(a6)
	move.l	d0,var_EmergencyBlock(a5)
	beq.s	.NoFlag
	bset	#SFB_BLOCKAVAILABLE,var_Flags(a5)
.NoFlag	lea	var_SetFuncSemaphore(a5),a1
	move.l	#Semaphore.Name,LN_NAME(a1)
	jsr	_LVOAddSemaphore(a6)
	jsr	_LVOPermit(a6)
	move.l	#Installed.MSG,d1
	exg.l	a4,a6
	jsr	_LVOPutStr(a6)
	exg.l	a6,a4
; cut the seglist
	lea	Start-4(pc),a0
	clr.l	(a0)	; tchac
	moveq	#0,d7
CloseDOS
	move.l	a4,a1
	jsr	_LVOCloseLibrary(a6)
NoDOS	move.l	d7,d0
	rts

NoSetFunction
	move.l	var_MainHeader(a5),a1
	move.l	#libh_SIZEOF,d0
	jsr	FreeMemory
NoMemory
	move.l	var_MemoryPool(a5),a0
	jsr	_LVODeletePool(a6)
NoPool	move.l	#Failed.MSG,d1
PrintAndGo
	exg.l	a4,a6
	jsr	_LVOPutStr(a6)
	exg.l	a4,a6
	bra.s	CloseDOS

AlreadyInstalled
	move.l	#AlreadyInstalled.MSG,d1
	moveq	#5,d7
	bra.s	PrintAndGo

DOS.Name	dc.b	'dos.library',0
Bienvenue.MSG	dc.b	'SafeFunction '
	VERSION
	dc.b	' by F.Delacroix, public domain.',10,0
Installed.MSG	dc.b	'  Patch successfully installed',10,0
Failed.MSG	dc.b	'  Failed to install patch ! (not enough memory ?!)',10,0
AlreadyInstalled.MSG	dc.b	'  Patch already active (you cannot remove it).',10,0

	section	ResidentPatch,CODE

	STRUCTURE	LibHeader,0
	 APTR	libh_LibBase	; points to patched library
	 APTR	libh_Patches	; points to list of patches
	 APTR	libh_NextLib	; points to next library header
	LABEL	libh_SIZEOF

	STRUCTURE	LibPatch,0
	 UWORD	libp_JMP	; JMP instruction
	 APTR	libp_OldFunc	; old (before patch) function
	 APTR	libp_NewFunc	; new (after patch) function
	 WORD	libp_Offset	; offset in jump table
	 APTR	libp_NextPatch	; points to next patch in library
	LABEL	libp_SIZEOF

MAGIC_OFFSET	EQU	$2357	; prime numbers :-]


MakeHead	; (D0)LibHeader=MakeHead(Library)(A0)
	movem.l	d1/a0-a2,-(sp)
	move.l	a0,a2
	move.l	#libh_SIZEOF,d0
	bsr	AllocMemory
	move.l	d0,a0
	move.l	a0,d0
	beq.s	.Fail
	move.l	a2,libh_LibBase(a0)
.Fail	movem.l	(sp)+,d1/a0-a2
	rts

MakePatch	; (D0)Patches=MakePatch(Offset,Next,NewFunc)(D0,A0,A1)
	movem.l	d1-d2/a0-a3,-(sp)
	move.l	d0,d2
	move.l	a0,a2
	move.l	a1,a3
	move.l	#libp_SIZEOF,d0
	bsr	AllocMemory
	move.l	d0,a0
	move.l	a0,d0
	beq.s	.Fail
	move.w	#$4EF9,libp_JMP(a0)	; opcode for JMP
	move.l	a3,libp_NewFunc(a0)
	move.w	d2,libp_Offset(a0)
	move.l	a2,libp_NextPatch(a0)
.Fail	movem.l	(sp)+,d1-d2/a0-a3
	rts

SF_OFFSET	EQUR	d7
SF_LIBRARY	EQUR	d6
SF_NEWFUNC	EQUR	d5
SF_OLDFUNC	EQUR	d4

OldSetFunction	; (D0)OldFunc=OldSetFunction(NewFunc)(A0)
	move.w	SF_OFFSET,d0
	ext.l	d0
	exg.l	d0,a0
	move.l	SF_LIBRARY,a1
	move.l	var_OldSetFunc(a5),-(sp)
	rts

NewSetFunction	; (D0)OldFunc=NewSetFunction(Offset,Library,NewFunc)(A0,A1,D0)
	movem.l	a0-a3/a5/d1/d4-d7,-(sp)
	lea	Variables,a5
	moveq	#0,SF_OLDFUNC
	move.w	a0,SF_OFFSET
	move.l	a1,SF_LIBRARY
	move.l	d0,SF_NEWFUNC
	lea	var_SetFuncSemaphore(a5),a0
	jsr	_LVOObtainSemaphore(a6)
	move.l	var_MainHeader(a5),a0

.NextLib
	move.l	libh_NextLib(a0),d0
	beq	.NoMoreLib
	move.l	d0,a1
	cmp.l	libh_LibBase(a1),SF_LIBRARY
	beq.s	.LibFound
	move.l	a1,a0
	bra	.NextLib
.LibFound
	move.l	a0,var_HeadPtr(a5)
	move.l	a1,var_HeadPtr1(a5)

	move.l	libh_Patches(a1),var_Patch(a5)
.NextPatch
	move.l	var_Patch(a5),a0
	move.l	libp_NextPatch(a0),a1
	move.l	a1,var_Patch1(a5)
	beq.s	.MakeANewPatch
	cmp.w	libp_Offset(a1),SF_OFFSET
	ble.s	.SearchPatchOn

; make a new patch (none installed here)
.MakeANewPatch
	move.w	SF_OFFSET,d0
	move.l	var_Patch1(a5),a0
	move.l	SF_NEWFUNC,a1
	bsr	MakePatch
	move.l	d0,var_Patch1(a5)
	beq	.Done
	move.l	SF_NEWFUNC,a0
	jsr	_LVODisable(a6)
	bsr	OldSetFunction
	move.l	var_Patch1(a5),a0
	move.l	d0,libp_OldFunc(a0)
	jsr	_LVOEnable(a6)
	tst.l	d0
	beq.s	.OldSetFuncFailed	; (should not happen, but let's be careful)
	move.l	var_Patch(a5),a0
	move.l	var_Patch1(a5),libp_NextPatch(a0)
	move.l	var_Patch1(a5),SF_OLDFUNC
	bra	.Done
.OldSetFuncFailed
	move.l	var_Patch1(a5),a1
	move.l	#libp_SIZEOF,d0
	bsr	FreeMemory
	bra	.Done

.SearchPatchOn
	move.l	var_Patch1(a5),a0
	cmp.w	libp_Offset(a0),SF_OFFSET
	bne	.NotThisPatch
	cmp.l	a0,SF_NEWFUNC
	beq.s	.RemoveLast
	cmp.l	libp_OldFunc(a0),SF_NEWFUNC
	bne.s	.NotRemoveLast
; user wants to remove the last patch
.RemoveLast
	move.l	libp_OldFunc(a0),a0
	bsr	OldSetFunction	; no need for Disable() here: let SetFunction() do it
	move.l	d0,SF_OLDFUNC
	move.l	var_Patch1(a5),a0
	cmp.l	libp_NewFunc(a0),SF_OLDFUNC
	bne.s	.BigTrouble
	move.l	var_Patch1(a5),a1
	move.l	var_Patch(a5),a0
	move.l	libp_NextPatch(a1),libp_NextPatch(a0)
	move.l	#libp_SIZEOF,d0
	bsr	FreeMemory
	bra.s	.LastRemoved
.BigTrouble	; if execution arrives here, than someone has changed
	move.l	SF_OLDFUNC,a0	; the vector from the outside...
	move.l	a0,SF_OLDFUNC	; we are in BIG TROUBLE, we just restore
	beq	.Done		; the original vector and pray.
	bsr	OldSetFunction
	moveq	#0,SF_OLDFUNC
	bra	.Done
.LastRemoved
	move.l	var_Patch(a5),a0
	cmp.w	#MAGIC_OFFSET,libp_Offset(a0)
	bne	.Done
	tst.l	libp_NextPatch(a0)
	bne	.Done
; remove unused library header
	move.l	var_HeadPtr(a5),a0
	move.l	var_HeadPtr1(a5),a1
	move.l	libh_NextLib(a1),libh_NextLib(a0)
	move.l	#libh_SIZEOF,d0
	bsr	FreeMemory
	move.l	var_Patch(a5),a1
	move.l	#libp_SIZEOF,d0
	bsr	FreeMemory
	bra	.Done

.NotRemoveLast
	move.l	var_Patch1(a5),a0
.RemoveLoop
	move.l	libp_NextPatch(a0),var_Patch2(a5)
	beq.s	.AddPatch
	move.l	var_Patch2(a5),a0
	cmp.w	libp_Offset(a0),SF_OFFSET
	bne.s	.AddPatch
	cmp.l	a0,SF_NEWFUNC
	beq.s	.RemoveThisOne
	cmp.l	libp_OldFunc(a0),SF_NEWFUNC
	beq.s	.RemoveThisOne
	move.l	a0,var_Patch1(a5)
	bra.s	.RemoveLoop

.RemoveThisOne
	move.l	var_Patch1(a5),a0
	move.l	var_Patch2(a5),a1
	move.l	libp_NextPatch(a1),libp_NextPatch(a0)
; no need for Disable() for a single instruction
	move.l	libp_OldFunc(a1),libp_OldFunc(a0)
	move.l	libp_NewFunc(a1),SF_OLDFUNC
	move.l	#libp_SIZEOF,d0
	bsr	FreeMemory
	bra	.Done

.AddPatch
; let's make a new patch on this function
	move.w	SF_OFFSET,d0
	move.l	var_Patch(a5),a0
	move.l	libp_NextPatch(a0),a0
	move.l	SF_NEWFUNC,a1
	bsr	MakePatch
	move.l	d0,var_Patch1(a5)
	beq	.Done
	move.l	SF_NEWFUNC,a0
	jsr	_LVODisable(a6)
	bsr	OldSetFunction
	move.l	var_Patch1(a5),a0
	move.l	d0,libp_OldFunc(a0)
	jsr	_LVOEnable(a6)
	tst.l	d0
	beq.s	.OldSetFuncFailed2
	move.l	var_Patch(a5),a0
	move.l	var_Patch1(a5),libp_NextPatch(a0)
	move.l	var_Patch1(a5),SF_OLDFUNC
	bra	.Done
.OldSetFuncFailed2
	move.l	var_Patch1(a5),a1
	move.l	#libp_SIZEOF,d0
	bsr	FreeMemory
	bra	.Done

.NotThisPatch
	move.l	var_Patch1(a5),var_Patch(a5)
	bra	.NextPatch

.NoMoreLib
; create a new library header
	move.l	a0,var_HeadPtr(a5)
	move.l	SF_LIBRARY,a0
	bsr	MakeHead
	move.l	d0,var_HeadPtr1(a5)
	beq.s	.Done
	move.l	var_HeadPtr(a5),a0
	move.l	d0,libh_NextLib(a0)
	move.w	SF_OFFSET,d0
	suba.l	a0,a0
	move.l	SF_NEWFUNC,a1
	bsr	MakePatch
	move.l	d0,var_Patch(a5)
	beq.s	.CantMakePatch
	move.w	#MAGIC_OFFSET,d0
	move.l	var_Patch(a5),a0
	sub.l	a1,a1
	bsr	MakePatch
	move.l	var_HeadPtr1(a5),a0
	move.l	d0,libh_Patches(a0)
	beq.s	.NoMagic
	move.l	SF_NEWFUNC,a0
	jsr	_LVODisable(a6)
	bsr	OldSetFunction
	move.l	var_Patch(a5),a0
	move.l	d0,libp_OldFunc(a0)
	jsr	_LVOEnable(a6)
	tst.l	d0
	beq.s	.OldSetFuncFailed3
	move.l	var_HeadPtr(a5),a0
	move.l	var_HeadPtr1(a5),libh_NextLib(a0)
	move.l	var_Patch(a5),SF_OLDFUNC
	bra.s	.Done
.OldSetFuncFailed3
	move.l	var_HeadPtr1(a5),a1
	move.l	libh_Patches(a1),a1
	move.l	#libp_SIZEOF,d0
	bsr.s	FreeMemory
.NoMagic
	move.l	var_Patch(a5),a1
	move.l	#libp_SIZEOF,d0
	bsr.s	FreeMemory
.CantMakePatch
	move.l	var_HeadPtr1(a5),a1
	move.l	#libh_SIZEOF,d0
	bsr.s	FreeMemory
.Done	lea	var_SetFuncSemaphore(a5),a0
	jsr	_LVOReleaseSemaphore(a6)
	jsr	_LVOCacheClearU(a6)
	move.l	SF_OLDFUNC,d0
	movem.l	(sp)+,a0-a3/a5/d1/d4-d7
	rts

; flag definitions for var_Flags
	BITDEF	SF,ALLOCATING,0
	BITDEF	SF,BLOCKAVAILABLE,1

AllocMemory	; (D0)Memory=AllocMemory(Size)(D0)
	move.l	var_MemoryPool(a5),a0
	jsr	_LVOForbid(a6)
	bset	#SFB_ALLOCATING,var_Flags(a5)
	jsr	_LVOAllocPooled(a6)
	bclr	#SFB_ALLOCATING,var_Flags(a5)
	jsr	_LVOPermit(a6)
	btst	#SFB_BLOCKAVAILABLE,var_Flags(a5)
	bne.s	.NoNeed
	move.l	d0,-(sp)
	beq.s	.Nope
	move.l	#EMERGENCY_BLOCKSIZE,d0
	move.l	#MEMF_PUBLIC,d1
	jsr	_LVOAllocMem(a6)
	move.l	d0,var_EmergencyBlock(a5)
	beq.s	.Nope
	bset	#SFB_BLOCKAVAILABLE,var_Flags(a5)
.Nope	move.l	(sp)+,d0
.NoNeed	rts

FreeMemory	; FreeMemory(Memory,Size)(A1,D0)
	move.l	var_MemoryPool(a5),a0
	jmp	_LVOFreePooled(a6)

LowMemHandler.Entry
	move.b	var_Flags(a1),d0
	and.b	#SFF_ALLOCATING!SFF_BLOCKAVAILABLE,d0
	beq.s	.DoNothing
	cmp.w	#EMERGENCY_BLOCKSIZE,memh_RequestSize(a0)
	bgt.s	.DoNothing
	bclr	#SFB_BLOCKAVAILABLE,var_Flags(a1)
	move.l	var_EmergencyBlock(a1),a1
	move.l	#EMERGENCY_BLOCKSIZE,d0
	jsr	_LVOFreeMem(a6)
	move.l	#MEM_ALL_DONE,d0
	rts
.DoNothing
	move.l	#MEM_DID_NOTHING,d0
	rts

LowMemHandler.Interrupt
	dc.l	0,0
	dc.b	NT_INTERRUPT,-120
	dc.l	LowMemHandler.Name
	dc.l	Variables	; IS_DATA field
	dc.l	LowMemHandler.Entry

Semaphore.Name	dc.b	'SafeFunction.Semaphore',0
LowMemHandler.Name	dc.b	'SafeFunction.EmergencyHandler',0

	section	Variables,BSS

	rsreset	; these variables are protected by the semaphore
var_HeadPtr	rs.l	1
var_HeadPtr1	rs.l	1
var_Patch	rs.l	1
var_Patch1	rs.l	1
var_Patch2	rs.l	1
var_MainHeader	rs.l	1	; first LibHeader
var_OldSetFunc	rs.l	1	; points to real SetFunction()
var_SetFuncSemaphore	rs.b	SS_SIZE
var_EmergencyBlock	rs.l	1
var_MemoryPool	rs.l	1
var_Flags	rs.b	1
var_Flags1	rs.b	1	; word-alignment, for now...
var_SIZEOF	rs.w	0

Variables	ds.b	var_SIZEOF