Adding controller code