Happy 62nd Independence Day

Saare Jahan se Aacha,
Hindustan Hamara..........

Jai Hind.

Execute/Run code after main() function returns

Recently one of my friends asked me whether it is possible to run/execute code after main() returns; provided the condition that we cannot use atexit() to register a function to run after main() returns.

Here’s my take at it:

#include <iostream>
using namespace std;

void lastfn(void)
{
 cout<<"This is executed after main()"<<endl;
 exit(0); //No more messing with the stack.
}

/*
Lets make main() 'naked' so that there is no prolog or epilog 
appended to the function. So when we enter into the function,
return address is on the top of the stack.
*/
__declspec (naked) void  __cdecl main(int count, char* vector[])
{
 long retaddr;

 __asm{
  pop retaddr  //pop the actual return address
  push lastfn  //push the new address where we 
     //want control after main
 }
 
 cout<<"Main routine"<<endl;

 __asm{
  retn
 }
}

A little explanation:
In this program, I have made the main() function naked. This will instruct the compiler NOT to emit the prolog and the epilog code. And as we know that it is the prolog code that make the stack frame for a new function call by pushing the frame pointer 'ebp' and making the new frame pointer equal to the current stack pointer 'esp'.

A general prolog looks something like this:

  push        ebp                ; Save ebp
  mov         ebp, esp           ; Set stack frame pointer
  sub         esp, localbytes    ; Allocate space for locals
  push                ; Save registers

The epilog code does the reverse of whatever has been done in the prolog code. A typical epilog looks something like below

pop     ; Restore registers
mov  esp, ebp      ; Restore stack pointer
pop  ebp           ; Restore ebp
ret                ; Return from function

Also as we know that when a function is called the stack looks something like this

High Address

|                   |
|-------------------|
|  Input arguments  |
|-------------------|
|  Return Address   |
|-------------------|
|   Saved ebp       |
|-------------------|
|  Local variables  | top of stack
|-------------------|
|                   |

Low Address

But without the prolog code the stack at the function call time looks something like this

High Address

|                   |
|-------------------|
|  Input arguments  |
|-------------------|
|  Return Address   | top of stack
|-------------------|
|                   |

Low Address

This is because there is no code emitted by the compiler that does the task of saving the frame pointer and making room for the local variables. So as we can see, the return address is at the top of the stack, which we can easily remove from the stack using a 'pop' and 'push' our address there.

PS: If someone reading this has a different solution or idea, please leave a comment.