Fix up connection logic