Fix missed inludes

Signed-off-by: Atanas Kerezov <kerezov.atanas@gmail.com>