Simplify logic of replacing socket send/receive functions