Improving socket server

- Small refactoring
- Now handles socket disconnect properly