Fix possible NULL reference bug