Refactor Socket code