event_ctx->unexpected_event = 1;
                        goto out;
                }
-               /* update topology event if hdcp is not desired */
-               status = mod_hdcp_add_display_topology(hdcp);
        } else if (is_in_hdcp1_states(hdcp)) {
                status = mod_hdcp_hdcp1_execution(hdcp, event_ctx, &input->hdcp1);
        } else if (is_in_hdcp1_dp_states(hdcp)) {
                        mod_hdcp_hdcp1_destroy_session(hdcp);
 
                }
-               if (hdcp->auth.trans_input.hdcp1.add_topology == PASS) {
-                       status = mod_hdcp_remove_display_topology(hdcp);
-                       if (status != MOD_HDCP_STATUS_SUCCESS) {
-                               output->callback_needed = 0;
-                               output->watchdog_timer_needed = 0;
-                               goto out;
-                       }
-               }
+
                HDCP_TOP_RESET_AUTH_TRACE(hdcp);
                memset(&hdcp->auth, 0, sizeof(struct mod_hdcp_authentication));
                memset(&hdcp->state, 0, sizeof(struct mod_hdcp_state));
                                goto out;
                        }
                }
-               if (hdcp->auth.trans_input.hdcp2.add_topology == PASS) {
-                       status = mod_hdcp_remove_display_topology(hdcp);
-                       if (status != MOD_HDCP_STATUS_SUCCESS) {
-                               output->callback_needed = 0;
-                               output->watchdog_timer_needed = 0;
-                               goto out;
-                       }
-               }
+
                HDCP_TOP_RESET_AUTH_TRACE(hdcp);
                memset(&hdcp->auth, 0, sizeof(struct mod_hdcp_authentication));
                memset(&hdcp->state, 0, sizeof(struct mod_hdcp_state));
                set_state_id(hdcp, output, HDCP_INITIALIZED);
        } else if (is_in_cp_not_desired_state(hdcp)) {
-               status = mod_hdcp_remove_display_topology(hdcp);
-               if (status != MOD_HDCP_STATUS_SUCCESS) {
-                       output->callback_needed = 0;
-                       output->watchdog_timer_needed = 0;
-                       goto out;
-               }
                HDCP_TOP_RESET_AUTH_TRACE(hdcp);
                memset(&hdcp->auth, 0, sizeof(struct mod_hdcp_authentication));
                memset(&hdcp->state, 0, sizeof(struct mod_hdcp_state));
        if (status != MOD_HDCP_STATUS_SUCCESS)
                goto out;
 
-       /* add display to connection */
-       hdcp->connection.link = *link;
-       *display_container = *display;
-
        /* reset retry counters */
        reset_retry_counts(hdcp);
 
        /* reset error trace */
        memset(&hdcp->connection.trace, 0, sizeof(hdcp->connection.trace));
 
+       /* add display to connection */
+       hdcp->connection.link = *link;
+       *display_container = *display;
+       status = mod_hdcp_add_display_to_topology(hdcp, display->index);
+       if (status != MOD_HDCP_STATUS_SUCCESS)
+               goto out;
+
        /* request authentication */
        if (current_state(hdcp) != HDCP_INITIALIZED)
                set_state_id(hdcp, output, HDCP_INITIALIZED);
        if (status != MOD_HDCP_STATUS_SUCCESS)
                goto out;
 
-       /* remove display */
-       display->state = MOD_HDCP_DISPLAY_INACTIVE;
-
        /* clear retry counters */
        reset_retry_counts(hdcp);
 
        /* reset error trace */
        memset(&hdcp->connection.trace, 0, sizeof(hdcp->connection.trace));
 
+       /* remove display */
+       status = mod_hdcp_remove_display_from_topology(hdcp, index);
+       if (status != MOD_HDCP_STATUS_SUCCESS)
+               goto out;
+       display->state = MOD_HDCP_DISPLAY_INACTIVE;
+
        /* request authentication for remaining displays*/
        if (get_active_display_count(hdcp) > 0)
                callback_in_ms(hdcp->connection.link.adjust.auth_delay * 1000,
 
        in->process.msg3_desc.msg_id = TA_HDCP_HDCP2_MSG_ID__NULL_MESSAGE;
        in->process.msg3_desc.msg_size = 0;
 }
-enum mod_hdcp_status mod_hdcp_remove_display_topology(struct mod_hdcp *hdcp)
-{
-
-       struct psp_context *psp = hdcp->config.psp.handle;
-       struct ta_dtm_shared_memory *dtm_cmd;
-       struct mod_hdcp_display *display = NULL;
-       uint8_t i;
+enum mod_hdcp_status mod_hdcp_remove_display_from_topology(
+               struct mod_hdcp *hdcp, uint8_t index)
+ {
+       struct psp_context *psp = hdcp->config.psp.handle;
+       struct ta_dtm_shared_memory *dtm_cmd;
+       struct mod_hdcp_display *display =
+                       get_active_display_at_index(hdcp, index);
 
        dtm_cmd = (struct ta_dtm_shared_memory *)psp->dtm_context.dtm_shared_buf;
 
-       for (i = 0; i < MAX_NUM_OF_DISPLAYS; i++) {
-               if (is_display_added(&(hdcp->connection.displays[i]))) {
-
-                       memset(dtm_cmd, 0, sizeof(struct ta_dtm_shared_memory));
-
-                       display = &hdcp->connection.displays[i];
+       if (!display || !is_display_added(display))
+               return MOD_HDCP_STATUS_DISPLAY_NOT_FOUND;
 
-                       dtm_cmd->cmd_id = TA_DTM_COMMAND__TOPOLOGY_UPDATE_V2;
-                       dtm_cmd->dtm_in_message.topology_update_v2.display_handle = display->index;
-                       dtm_cmd->dtm_in_message.topology_update_v2.is_active = 0;
-                       dtm_cmd->dtm_status = TA_DTM_STATUS__GENERIC_FAILURE;
+       memset(dtm_cmd, 0, sizeof(struct ta_dtm_shared_memory));
 
-                       psp_dtm_invoke(psp, dtm_cmd->cmd_id);
+       dtm_cmd->cmd_id = TA_DTM_COMMAND__TOPOLOGY_UPDATE_V2;
+       dtm_cmd->dtm_in_message.topology_update_v2.display_handle = display->index;
+       dtm_cmd->dtm_in_message.topology_update_v2.is_active = 0;
+       dtm_cmd->dtm_status = TA_DTM_STATUS__GENERIC_FAILURE;
 
-                       if (dtm_cmd->dtm_status != TA_DTM_STATUS__SUCCESS)
-                               return MOD_HDCP_STATUS_UPDATE_TOPOLOGY_FAILURE;
+       psp_dtm_invoke(psp, dtm_cmd->cmd_id);
 
-                       display->state = MOD_HDCP_DISPLAY_ACTIVE;
-                       HDCP_TOP_REMOVE_DISPLAY_TRACE(hdcp, display->index);
-               }
-       }
+       if (dtm_cmd->dtm_status != TA_DTM_STATUS__SUCCESS)
+               return MOD_HDCP_STATUS_UPDATE_TOPOLOGY_FAILURE;
 
-       return MOD_HDCP_STATUS_SUCCESS;
-}
+       display->state = MOD_HDCP_DISPLAY_ACTIVE;
+       HDCP_TOP_REMOVE_DISPLAY_TRACE(hdcp, display->index);
+ 
+       return MOD_HDCP_STATUS_SUCCESS;
+ }
 
-enum mod_hdcp_status mod_hdcp_add_display_topology(struct mod_hdcp *hdcp)
+enum mod_hdcp_status mod_hdcp_add_display_to_topology(
+               struct mod_hdcp *hdcp, uint8_t index)
 {
        struct psp_context *psp = hdcp->config.psp.handle;
        struct ta_dtm_shared_memory *dtm_cmd;
-       struct mod_hdcp_display *display = NULL;
+       struct mod_hdcp_display *display =
+                       get_active_display_at_index(hdcp, index);
        struct mod_hdcp_link *link = &hdcp->connection.link;
-       uint8_t i;
 
        if (!psp->dtm_context.dtm_initialized) {
                DRM_ERROR("Failed to add display topology, DTM TA is not initialized.");
                return MOD_HDCP_STATUS_FAILURE;
        }
 
+       if (!display || is_display_added(display))
+               return MOD_HDCP_STATUS_UPDATE_TOPOLOGY_FAILURE;
+
        dtm_cmd = (struct ta_dtm_shared_memory *)psp->dtm_context.dtm_shared_buf;
 
-       for (i = 0; i < MAX_NUM_OF_DISPLAYS; i++) {
-               if (hdcp->connection.displays[i].state == MOD_HDCP_DISPLAY_ACTIVE) {
-                       display = &hdcp->connection.displays[i];
-
-                       memset(dtm_cmd, 0, sizeof(struct ta_dtm_shared_memory));
-
-                       dtm_cmd->cmd_id = TA_DTM_COMMAND__TOPOLOGY_UPDATE_V2;
-                       dtm_cmd->dtm_in_message.topology_update_v2.display_handle = display->index;
-                       dtm_cmd->dtm_in_message.topology_update_v2.is_active = 1;
-                       dtm_cmd->dtm_in_message.topology_update_v2.controller = display->controller;
-                       dtm_cmd->dtm_in_message.topology_update_v2.ddc_line = link->ddc_line;
-                       dtm_cmd->dtm_in_message.topology_update_v2.dig_be = link->dig_be;
-                       dtm_cmd->dtm_in_message.topology_update_v2.dig_fe = display->dig_fe;
-                       dtm_cmd->dtm_in_message.topology_update_v2.dp_mst_vcid = display->vc_id;
-                       dtm_cmd->dtm_in_message.topology_update_v2.max_hdcp_supported_version =
-                               TA_DTM_HDCP_VERSION_MAX_SUPPORTED__2_2;
-                       dtm_cmd->dtm_status = TA_DTM_STATUS__GENERIC_FAILURE;
-
-                       psp_dtm_invoke(psp, dtm_cmd->cmd_id);
-
-                       if (dtm_cmd->dtm_status != TA_DTM_STATUS__SUCCESS)
-                               return MOD_HDCP_STATUS_UPDATE_TOPOLOGY_FAILURE;
-
-                       display->state = MOD_HDCP_DISPLAY_ACTIVE_AND_ADDED;
-                       HDCP_TOP_ADD_DISPLAY_TRACE(hdcp, display->index);
-               }
-       }
+       memset(dtm_cmd, 0, sizeof(struct ta_dtm_shared_memory));
 
-       return MOD_HDCP_STATUS_SUCCESS;
+       dtm_cmd->cmd_id = TA_DTM_COMMAND__TOPOLOGY_UPDATE_V2;
+       dtm_cmd->dtm_in_message.topology_update_v2.display_handle = display->index;
+       dtm_cmd->dtm_in_message.topology_update_v2.is_active = 1;
+       dtm_cmd->dtm_in_message.topology_update_v2.controller = display->controller;
+       dtm_cmd->dtm_in_message.topology_update_v2.ddc_line = link->ddc_line;
+       dtm_cmd->dtm_in_message.topology_update_v2.dig_be = link->dig_be;
+       dtm_cmd->dtm_in_message.topology_update_v2.dig_fe = display->dig_fe;
+       if (is_dp_hdcp(hdcp))
+               dtm_cmd->dtm_in_message.topology_update_v2.is_assr = link->dp.assr_supported;
+
+       dtm_cmd->dtm_in_message.topology_update_v2.dp_mst_vcid = display->vc_id;
+       dtm_cmd->dtm_in_message.topology_update_v2.max_hdcp_supported_version =
+                       TA_DTM_HDCP_VERSION_MAX_SUPPORTED__2_2;
+       dtm_cmd->dtm_status = TA_DTM_STATUS__GENERIC_FAILURE;
+
+       psp_dtm_invoke(psp, dtm_cmd->cmd_id);
+
+       if (dtm_cmd->dtm_status != TA_DTM_STATUS__SUCCESS)
+               return MOD_HDCP_STATUS_UPDATE_TOPOLOGY_FAILURE;
+
+       display->state = MOD_HDCP_DISPLAY_ACTIVE_AND_ADDED;
+       HDCP_TOP_ADD_DISPLAY_TRACE(hdcp, display->index);
+ 
+       return MOD_HDCP_STATUS_SUCCESS;
 }
 
 enum mod_hdcp_status mod_hdcp_hdcp1_create_session(struct mod_hdcp *hdcp)
 
        struct psp_context *psp = hdcp->config.psp.handle;
        struct ta_hdcp_shared_memory *hdcp_cmd;
+       uint8_t i = 0;
 
        hdcp_cmd = (struct ta_hdcp_shared_memory *)psp->hdcp_context.hdcp_shared_buf;
        memset(hdcp_cmd, 0, sizeof(struct ta_hdcp_shared_memory));
                return MOD_HDCP_STATUS_HDCP1_DESTROY_SESSION_FAILURE;
 
        HDCP_TOP_HDCP1_DESTROY_SESSION_TRACE(hdcp);
+       for (i = 0; i < MAX_NUM_OF_DISPLAYS; i++)
+               if (is_display_encryption_enabled(
+                               &hdcp->connection.displays[i])) {
+                       hdcp->connection.displays[i].state =
+                                       MOD_HDCP_DISPLAY_ACTIVE_AND_ADDED;
+                       HDCP_HDCP1_DISABLED_TRACE(hdcp,
+                                       hdcp->connection.displays[i].index);
+               }
 
        return MOD_HDCP_STATUS_SUCCESS;
 }
 {
        struct psp_context *psp = hdcp->config.psp.handle;
        struct ta_hdcp_shared_memory *hdcp_cmd;
+       uint8_t i = 0;
 
        hdcp_cmd = (struct ta_hdcp_shared_memory *)psp->hdcp_context.hdcp_shared_buf;
        memset(hdcp_cmd, 0, sizeof(struct ta_hdcp_shared_memory));
                return MOD_HDCP_STATUS_HDCP2_DESTROY_SESSION_FAILURE;
 
        HDCP_TOP_HDCP2_DESTROY_SESSION_TRACE(hdcp);
+       for (i = 0; i < MAX_NUM_OF_DISPLAYS; i++)
+               if (is_display_encryption_enabled(
+                               &hdcp->connection.displays[i])) {
+                       hdcp->connection.displays[i].state =
+                                       MOD_HDCP_DISPLAY_ACTIVE_AND_ADDED;
+                       HDCP_HDCP2_DISABLED_TRACE(hdcp,
+                                       hdcp->connection.displays[i].index);
+               }
 
        return MOD_HDCP_STATUS_SUCCESS;
 }