Simplify socket handling and some code cleanup