Fix key mask

Fixes a minor bug with the return of key masks.