Add validation of recieved message and refactor code