Refactored ws server code