Add some controller logic