Improve hook registration

Change-Id: If2e9578af48acfb548cca8dfa05d114b480b58a8