Cleaned up some code and added disconnect method