I can think of a number of ways of fixing this, but yours has the virtue of minimizing word changes to the spec and lines of code to an implementation. This is very late in the process to be making cryptographic changes, but I feel like this one is worth doing.