//-------------------------------------------------------------------
//	newsched.c
//
//	This module exploits the Pentium's breakpoint capability to
//	intercept every call to the kernel's 'schedule()' function, 
//	diverting control to an alternate function.  Its purpose is 
//	purely to demonstrate the feasibility of 'dynamic patching' 
//	as a way to alter Linux's task-scheduling behavior, without
//	the need to recompile any kernel sources; thus in this demo
//	the only action taken by our substitute scheduling function
//	is to increment a counter, then transfer control to Linux's
//	original scheduler -- with the breakpoint being temporarily 
//	disabled by setting the Resume-Flag in the EFLAGS register. 
//
//	NOTE: Written and tested using Linux kernel version 2.4.26.
//
//	programmer: ALLAN CRUSE
//	written on: 06 NOV 2004
//	revised on: 12 NOV 2004 -- per suggestion by Michael Elliot
//	revised on: 13 NOV 2004 -- for quitting in cleanup_module()
//-------------------------------------------------------------------

#include <linux/module.h>	// for init_module() 
#include <linux/proc_fs.h>	// for create_proc_read_entry() 

#define DEBUG_TRAP_ID	0x01

static char modname[] = "newsched";
static unsigned short oldidtr[3], newidtr[3];
static unsigned long long *oldidt, *newidt;
static unsigned long kpage, selector_cs;
extern void schedule( void );

MODULE_LICENSE("GPL"); 

static unsigned long old_scheduler;
static unsigned long new_scheduler;
static unsigned long sched_count = 0;
static unsigned long quitting = 0;	// <-- added 11/13/2004

//-- OUR SUBSTITUTE SCHEDULER ROUTINE --//
asmlinkage void new_schedule( void );
asm("	.text				");
asm("	.type	new_schedule, @function	");
asm("new_schedule:			");
// increment count of 'schedule()' calls
asm("	lock incl sched_count		");
// transfer control to Linux's scheduler
// via IRET-instruction with RF-flag set
asm("	pushfl				");
asm("	bts	$16, (%esp) 		");		
asm("	pushl	%cs 			");
asm("	pushl  	old_scheduler 		");
asm("	iret				");	
//--------------------------------------//

static void dbg_intercept( unsigned long *tos )
{
	unsigned long	dbstatus;

	// examine debug status (and reset register DR6) 
	asm(" movl  %dr6, %eax ");
	asm(" movl  %%eax, %0 " : "=m" (dbstatus) );
	asm(" xorl  %eax, %eax ");	// <-- added 11/12/2004
	asm(" movl  %eax, %dr6 ");	// <-- added 11/12/2004
	
	// if DB-bit is set, skip next instruction (unless quitting)
	if (( !quitting )&&( dbstatus & (1<<13) )) tos[ 10 ] += 3;

	// if B0-bit is set, take detour to new scheduler
	if ( dbstatus & 1 ) tos[ 10 ] = new_scheduler;
	
	// stop task-switcher from disabling breakpoint
	if ( !quitting )	// <-- added 11/13/2004
		{
		asm(" movl  %dr7, %eax ");
		asm(" andl  $0xFFF0FFFC, %eax ");
		asm(" orl   $0x00002003, %eax ");
		asm(" movl  %eax, %dr7 ");
		}
}	

//----  OUR DEBUG-FAULT EXCEPTION-HANDLER  -----//
asmlinkage void isr_debug_fault( void );
asm("	.text					");
asm("	.type	isr_debug_fault, @function	");
asm("isr_debug_fault:				");
asm("	pushal					");
asm("	pushl	%ds				");
asm("	pushl	%es				");
//
asm("	movl	%ss, %eax			");
asm("	movl	%eax, %ds			");
asm("	movl	%eax, %es			");
//
asm("	pushl	%esp				");
asm("	call	dbg_intercept			");
asm("	addl	$4, %esp			");
//
asm("	popl	%es				");
asm("	popl	%ds				");
asm("	popal	   				");
asm(" 	iret					");	
//----------------------------------------------//

static int my_get_info( char *buf, char **start, off_t off, int count )
{
	int	len = 0;
	len += sprintf( buf+len, "Number of \'schedule()\' calls: " );
	len += sprintf( buf+len, "%d \n", sched_count );
	return	len;
}

static void load_IDTR( void *regimage )
{
	asm(" lidt %0 " : : "m" (*(unsigned short *)regimage) );
}


static void setup_breakpoint0( void *address )
{
	asm(" movl %0, %%eax " : : "m" (*(unsigned long*)address) );
	asm(" movl %eax, %dr0 ");
	asm(" movl %dr7, %eax ");
	asm(" andl $0xFFF0FFFC, %eax ");
	asm(" orl  $0x00002003, %eax ");
	asm(" movl %eax, %dr7 ");
}	


static void clear_breakpoint0( void *address )
{
	asm(" movl %dr7, %eax ");
	asm(" andl $0xFFF0DFFC, %eax ");
	asm(" movl %eax, %dr7 ");
}	


int init_module( void )
{
	unsigned long long	gate_desc;
	unsigned long		breakpt;
	
	printk( "<1>\nInstalling \'%s\' module\n", modname );
	
	// initialize our function-pointers
	old_scheduler = (unsigned long)schedule;
	new_scheduler = (unsigned long)new_schedule;

	// allocate a page of kernel memory to hold our new IDT
	kpage = get_free_page( GFP_KERNEL );
	if ( !kpage ) return -ENOMEM;

	// copy the Interrupt Descriptor Table to this 'writable' page 
	asm(" sidt oldidtr \n sidt newidtr \n movw %cs, selector_cs ");
	memcpy( newidtr+1, &kpage, sizeof( kpage ) );
	oldidt = (unsigned long long*)(*(unsigned long*)(oldidtr+1)); 
	newidt = (unsigned long long*)(*(unsigned long*)(newidtr+1)); 
	memcpy( newidt, oldidt, 256 * sizeof( unsigned long long ) );

	// modify our copy of the Interrupt Descriptor Table 
	gate_desc = (unsigned long long)isr_debug_fault;
	gate_desc &= 0x00000000FFFFFFFFLL;
	gate_desc |= ( gate_desc << 32 );
	gate_desc &= 0xFFFF00000000FFFFLL;
	gate_desc |= 0x00008E0000000000LL;
	gate_desc |= ( selector_cs << 16 );
	newidt[ DEBUG_TRAP_ID ] = gate_desc;	
	
	// now activate the new Interrupt Descriptor Table
	load_IDTR( newidtr );
	smp_call_function( load_IDTR, newidtr, 1, 1 );	

	// then set the debug breakpoint
	breakpt = (unsigned long)schedule;
	smp_call_function( setup_breakpoint0, &breakpt, 1, 1 );
	setup_breakpoint0( &breakpt );
	
	// create the pseudo-file that lets users see 'schedule()' count
	create_proc_info_entry( modname, 0, NULL, my_get_info );
	return	0;  //SUCCESS
}


void cleanup_module( void )
{
	// clear the debug breakpoint
	quitting = 1;		// added 11/13/2004
	smp_call_function( clear_breakpoint0, NULL, 1, 1 );
	clear_breakpoint0( NULL );
	
	// deactivate our altered Interrupt Descriptor Table
	smp_call_function( load_IDTR, oldidtr, 1, 1 );
	load_IDTR( oldidtr );

	// delete the pseudo-file 
	remove_proc_entry( modname, NULL );

	// release the allocated page of kernel memory
	if ( kpage ) free_page( kpage );
	printk( "<1>Removing \'%s\' module\n", modname );
}