Fix more issues that have been found in code review.

Signed-off-by: Tzachi Dar <tzachi.dar@gmail.com>