You're allowed to do that. You're not allowed to hit the ball twice in an attempt to make runs. That's my understanding anyway. Batters alway swat a ball that is rolling towards their stumps.
Yeah I think the rule specifically mentions you're allowed to hit it the second time to guard your wicket.