Add NULL check for preventing crash