Add logic for the WS server
Fixed small bug