Added onMessage and refactored code to support it.