Fix #362, update code with upstream