Merge Android 24Q2 Release (ab/11526283) to aosp-main-future

Bug: 337098550
Merged-In: I53a278f1317f5307441103dc42dc5be6f20d7075
Change-Id: I3a53ebf7339cdb886d3c1b733b8eee1847f8832a
diff --git a/wifi/aidl/default/aidl_struct_util.cpp b/wifi/aidl/default/aidl_struct_util.cpp
index b62b3a0..d9e023c 100644
--- a/wifi/aidl/default/aidl_struct_util.cpp
+++ b/wifi/aidl/default/aidl_struct_util.cpp
@@ -1085,6 +1085,41 @@
     return true;
 }
 
+// TODO (b/324519882): Remove logs after validating the structure size.
+void logAidlLinkLayerStatsSize(StaLinkLayerStats& aidl_stats) {
+    unsigned long expectedMaxRadios = 5;
+    unsigned long expectedMaxLinks = 5;
+    unsigned long expectedMaxChannelStats = 512;
+    unsigned long expectedMaxPeers = 3;
+    unsigned long expectedMaxRateStats = 1024;
+
+    unsigned long maxChannelStats = 0, maxPeers = 0, maxRateStats = 0;
+    for (size_t i = 0; i < aidl_stats.radios.size(); i++) {
+        maxChannelStats =
+                std::max(maxChannelStats, (unsigned long)aidl_stats.radios[i].channelStats.size());
+    }
+    for (size_t i = 0; i < aidl_stats.iface.links.size(); i++) {
+        maxPeers = std::max(maxPeers, (unsigned long)aidl_stats.iface.links[i].peers.size());
+        for (size_t j = 0; j < aidl_stats.iface.links[i].peers.size(); j++) {
+            maxRateStats =
+                    std::max(maxRateStats,
+                             (unsigned long)aidl_stats.iface.links[i].peers[j].rateStats.size());
+        }
+    }
+
+    if (aidl_stats.radios.size() > expectedMaxRadios ||
+        aidl_stats.iface.links.size() > expectedMaxLinks ||
+        maxChannelStats > expectedMaxChannelStats || maxPeers > expectedMaxPeers ||
+        maxRateStats > expectedMaxRateStats) {
+        LOG(INFO) << "StaLinkLayerStats exceeds expected vector size";
+        LOG(INFO) << "  numRadios: " << aidl_stats.radios.size();
+        LOG(INFO) << "  numLinks: " << aidl_stats.iface.links.size();
+        LOG(INFO) << "  maxChannelStats: " << maxChannelStats;
+        LOG(INFO) << "  maxPeers: " << maxPeers;
+        LOG(INFO) << "  maxRateStats: " << maxRateStats;
+    }
+}
+
 bool convertLegacyPeerInfoStatsToAidl(const legacy_hal::WifiPeerInfo& legacy_peer_info_stats,
                                       StaPeerInfo* aidl_peer_info_stats) {
     if (!aidl_peer_info_stats) {
@@ -2741,9 +2776,8 @@
     if (!convertAidlRttConfigToLegacy(aidl_config, &(legacy_config->rtt_config))) {
         return false;
     }
-    legacy_config->tx_ltf_repetition_count = aidl_config.txLtfRepetitionCount;
-    legacy_config->ntb_min_measurement_time_millis = aidl_config.ntbMinMeasurementTimeMillis;
-    legacy_config->ntb_max_measurement_time_millis = aidl_config.ntbMaxMeasurementTimeMillis;
+    legacy_config->ntb_min_measurement_time = aidl_config.ntbMinMeasurementTime;
+    legacy_config->ntb_max_measurement_time = aidl_config.ntbMaxMeasurementTime;
     return true;
 }
 
@@ -2887,11 +2921,10 @@
     aidl_capabilities->bwSupport = convertLegacyRttBwBitmapToAidl(legacy_capabilities.bw_support);
     aidl_capabilities->mcVersion = legacy_capabilities.mc_version;
     // Initialize 11az parameters to default
-    aidl_capabilities->azPreambleSupport = RttPreamble::INVALID;
-    aidl_capabilities->azBwSupport = RttBw::BW_UNSPECIFIED;
+    aidl_capabilities->azPreambleSupport = (int)RttPreamble::INVALID;
+    aidl_capabilities->azBwSupport = (int)RttBw::BW_UNSPECIFIED;
     aidl_capabilities->ntbInitiatorSupported = false;
     aidl_capabilities->ntbResponderSupported = false;
-    aidl_capabilities->maxTxLtfRepetitionCount = 0;
     return true;
 }
 
@@ -2914,12 +2947,11 @@
             convertLegacyRttBwBitmapToAidl(legacy_capabilities_v3.rtt_capab.bw_support);
     aidl_capabilities->mcVersion = legacy_capabilities_v3.rtt_capab.mc_version;
     aidl_capabilities->azPreambleSupport =
-            convertLegacyRttPreambleBitmapToAidl(legacy_capabilities_v3.az_preamble_support);
+            (int)convertLegacyRttPreambleBitmapToAidl(legacy_capabilities_v3.az_preamble_support);
     aidl_capabilities->azBwSupport =
-            convertLegacyRttBwBitmapToAidl(legacy_capabilities_v3.az_bw_support);
+            (int)convertLegacyRttBwBitmapToAidl(legacy_capabilities_v3.az_bw_support);
     aidl_capabilities->ntbInitiatorSupported = legacy_capabilities_v3.ntb_initiator_supported;
     aidl_capabilities->ntbResponderSupported = legacy_capabilities_v3.ntb_responder_supported;
-    aidl_capabilities->maxTxLtfRepetitionCount = legacy_capabilities_v3.max_tx_ltf_repetition_count;
     return true;
 }
 
@@ -2994,9 +3026,12 @@
         }
         aidl_result.channelFreqMHz = 0;
         aidl_result.packetBw = RttBw::BW_UNSPECIFIED;
-        aidl_result.txLtfRepetitionCount = 0;
-        aidl_result.ntbMinMeasurementTimeMillis = 0;
-        aidl_result.ntbMaxMeasurementTimeMillis = 0;
+        aidl_result.i2rTxLtfRepetitionCount = 0;
+        aidl_result.r2iTxLtfRepetitionCount = 0;
+        aidl_result.ntbMinMeasurementTime = 0;
+        aidl_result.ntbMaxMeasurementTime = 0;
+        aidl_result.numTxSpatialStreams = 0;
+        aidl_result.numRxSpatialStreams = 0;
         aidl_results->push_back(aidl_result);
     }
     return true;
@@ -3017,9 +3052,12 @@
         aidl_result.channelFreqMHz =
                 legacy_result->frequency != UNSPECIFIED ? legacy_result->frequency : 0;
         aidl_result.packetBw = convertLegacyRttBwToAidl(legacy_result->packet_bw);
-        aidl_result.txLtfRepetitionCount = 0;
-        aidl_result.ntbMinMeasurementTimeMillis = 0;
-        aidl_result.ntbMaxMeasurementTimeMillis = 0;
+        aidl_result.i2rTxLtfRepetitionCount = 0;
+        aidl_result.r2iTxLtfRepetitionCount = 0;
+        aidl_result.ntbMinMeasurementTime = 0;
+        aidl_result.ntbMaxMeasurementTime = 0;
+        aidl_result.numTxSpatialStreams = 0;
+        aidl_result.numRxSpatialStreams = 0;
         aidl_results->push_back(aidl_result);
     }
     return true;
@@ -3041,9 +3079,12 @@
                                              ? legacy_result->rtt_result.frequency
                                              : 0;
         aidl_result.packetBw = convertLegacyRttBwToAidl(legacy_result->rtt_result.packet_bw);
-        aidl_result.txLtfRepetitionCount = legacy_result->tx_ltf_repetition_count;
-        aidl_result.ntbMinMeasurementTimeMillis = legacy_result->ntb_min_measurement_time_millis;
-        aidl_result.ntbMaxMeasurementTimeMillis = legacy_result->ntb_max_measurement_time_millis;
+        aidl_result.i2rTxLtfRepetitionCount = legacy_result->i2r_tx_ltf_repetition_count;
+        aidl_result.r2iTxLtfRepetitionCount = legacy_result->r2i_tx_ltf_repetition_count;
+        aidl_result.ntbMinMeasurementTime = legacy_result->ntb_min_measurement_time;
+        aidl_result.ntbMaxMeasurementTime = legacy_result->ntb_max_measurement_time;
+        aidl_result.numTxSpatialStreams = legacy_result->num_tx_sts;
+        aidl_result.numRxSpatialStreams = legacy_result->num_rx_sts;
         aidl_results->push_back(aidl_result);
     }
     return true;
@@ -3587,13 +3628,13 @@
     if (legacy_twt_capabs.min_wake_duration_micros > legacy_twt_capabs.max_wake_duration_micros) {
         return false;
     }
-    aidl_twt_capabs->minWakeDurationMicros = legacy_twt_capabs.min_wake_duration_micros;
-    aidl_twt_capabs->maxWakeDurationMicros = legacy_twt_capabs.max_wake_duration_micros;
+    aidl_twt_capabs->minWakeDurationUs = legacy_twt_capabs.min_wake_duration_micros;
+    aidl_twt_capabs->maxWakeDurationUs = legacy_twt_capabs.max_wake_duration_micros;
     if (legacy_twt_capabs.min_wake_interval_micros > legacy_twt_capabs.max_wake_interval_micros) {
         return false;
     }
-    aidl_twt_capabs->minWakeIntervalMicros = legacy_twt_capabs.min_wake_interval_micros;
-    aidl_twt_capabs->maxWakeIntervalMicros = legacy_twt_capabs.max_wake_interval_micros;
+    aidl_twt_capabs->minWakeIntervalUs = legacy_twt_capabs.min_wake_interval_micros;
+    aidl_twt_capabs->maxWakeIntervalUs = legacy_twt_capabs.max_wake_interval_micros;
     return true;
 }
 
@@ -3603,16 +3644,16 @@
         return false;
     }
     legacy_twt_request->mlo_link_id = aidl_twt_request.mloLinkId;
-    if (aidl_twt_request.minWakeDurationMicros > aidl_twt_request.maxWakeDurationMicros) {
+    if (aidl_twt_request.minWakeDurationUs > aidl_twt_request.maxWakeDurationUs) {
         return false;
     }
-    legacy_twt_request->min_wake_duration_micros = aidl_twt_request.minWakeDurationMicros;
-    legacy_twt_request->max_wake_duration_micros = aidl_twt_request.maxWakeDurationMicros;
-    if (aidl_twt_request.minWakeIntervalMicros > aidl_twt_request.maxWakeIntervalMicros) {
+    legacy_twt_request->min_wake_duration_micros = aidl_twt_request.minWakeDurationUs;
+    legacy_twt_request->max_wake_duration_micros = aidl_twt_request.maxWakeDurationUs;
+    if (aidl_twt_request.minWakeIntervalUs > aidl_twt_request.maxWakeIntervalUs) {
         return false;
     }
-    legacy_twt_request->min_wake_interval_micros = aidl_twt_request.minWakeIntervalMicros;
-    legacy_twt_request->max_wake_interval_micros = aidl_twt_request.maxWakeIntervalMicros;
+    legacy_twt_request->min_wake_interval_micros = aidl_twt_request.minWakeIntervalUs;
+    legacy_twt_request->max_wake_interval_micros = aidl_twt_request.maxWakeIntervalUs;
     return true;
 }
 
@@ -3664,8 +3705,8 @@
 
     aidl_twt_session->sessionId = twt_session.session_id;
     aidl_twt_session->mloLinkId = twt_session.mlo_link_id;
-    aidl_twt_session->wakeDurationMicros = twt_session.wake_duration_micros;
-    aidl_twt_session->wakeIntervalMicros = twt_session.wake_interval_micros;
+    aidl_twt_session->wakeDurationUs = twt_session.wake_duration_micros;
+    aidl_twt_session->wakeIntervalUs = twt_session.wake_interval_micros;
     switch (twt_session.negotiation_type) {
         case WIFI_TWT_NEGO_TYPE_INDIVIDUAL:
             aidl_twt_session->negotiationType = TwtSession::TwtNegotiationType::INDIVIDUAL;
@@ -3696,7 +3737,7 @@
     aidl_twt_stats->avgRxPktCount = twt_stats.avg_pkt_num_rx;
     aidl_twt_stats->avgTxPktSize = twt_stats.avg_tx_pkt_size;
     aidl_twt_stats->avgRxPktSize = twt_stats.avg_rx_pkt_size;
-    aidl_twt_stats->avgEospDurationMicros = twt_stats.avg_eosp_dur_us;
+    aidl_twt_stats->avgEospDurationUs = twt_stats.avg_eosp_dur_us;
     aidl_twt_stats->eospCount = twt_stats.eosp_count;
 
     return true;
diff --git a/wifi/aidl/default/aidl_struct_util.h b/wifi/aidl/default/aidl_struct_util.h
index 7089363..2574f95 100644
--- a/wifi/aidl/default/aidl_struct_util.h
+++ b/wifi/aidl/default/aidl_struct_util.h
@@ -94,6 +94,7 @@
                                          StaLinkLayerStats* aidl_stats);
 bool convertLegacyLinkLayerStatsToAidl(const legacy_hal::LinkLayerStats& legacy_stats,
                                        StaLinkLayerStats* aidl_stats);
+void logAidlLinkLayerStatsSize(StaLinkLayerStats& aidl_stats);
 bool convertLegacyRoamingCapabilitiesToAidl(
         const legacy_hal::wifi_roaming_capabilities& legacy_caps,
         StaRoamingCapabilities* aidl_caps);
diff --git a/wifi/aidl/default/wifi_legacy_hal.cpp b/wifi/aidl/default/wifi_legacy_hal.cpp
index 55d6f59..cf86120 100644
--- a/wifi/aidl/default/wifi_legacy_hal.cpp
+++ b/wifi/aidl/default/wifi_legacy_hal.cpp
@@ -1867,9 +1867,8 @@
     return {capabs, status};
 }
 
-wifi_error WifiLegacyHal::twtSessionSetup(
-        const std::string& ifaceName, uint32_t cmdId, const wifi_twt_request& request,
-        const on_twt_failure& on_twt_failure_user_callback,
+wifi_error WifiLegacyHal::twtRegisterEvents(
+        const std::string& ifaceName, const on_twt_failure& on_twt_failure_user_callback,
         const on_twt_session_create& on_twt_session_create_user_callback,
         const on_twt_session_update& on_twt_session_update_user_callback,
         const on_twt_session_teardown& on_twt_session_teardown_user_callback,
@@ -1921,13 +1920,18 @@
         on_twt_session_resume_user_callback(id, session_id);
     };
 
-    return global_func_table_.wifi_twt_session_setup(
-            cmdId, getIfaceHandle(ifaceName), request,
+    return global_func_table_.wifi_twt_register_events(
+            getIfaceHandle(ifaceName),
             {onAsyncTwtError, onAsyncTwtSessionCreate, onAsyncTwtSessionUpdate,
              onAsyncTwtSessionTeardown, onAsyncTwtSessionStats, onAsyncTwtSessionSuspend,
              onAsyncTwtSessionResume});
 }
 
+wifi_error WifiLegacyHal::twtSessionSetup(const std::string& ifaceName, uint32_t cmdId,
+                                          const wifi_twt_request& request) {
+    return global_func_table_.wifi_twt_session_setup(cmdId, getIfaceHandle(ifaceName), request);
+}
+
 wifi_error WifiLegacyHal::twtSessionUpdate(const std::string& ifaceName, uint32_t cmdId,
                                            uint32_t sessionId, const wifi_twt_request& request) {
     return global_func_table_.wifi_twt_session_update(cmdId, getIfaceHandle(ifaceName), sessionId,
diff --git a/wifi/aidl/default/wifi_legacy_hal.h b/wifi/aidl/default/wifi_legacy_hal.h
index 121d1b5..3fd567b 100644
--- a/wifi/aidl/default/wifi_legacy_hal.h
+++ b/wifi/aidl/default/wifi_legacy_hal.h
@@ -780,15 +780,16 @@
 
     // TWT functions
     std::pair<wifi_twt_capabilities, wifi_error> twtGetCapabilities(const std::string& ifaceName);
+    wifi_error twtRegisterEvents(
+            const std::string& ifaceName, const on_twt_failure& on_twt_failure_user_callback,
+            const on_twt_session_create& on_twt_session_create_user_callback,
+            const on_twt_session_update& on_twt_session_update_user_callback,
+            const on_twt_session_teardown& on_twt_session_teardown_user_callback,
+            const on_twt_session_stats& on_twt_session_stats_user_callback,
+            const on_twt_session_suspend& on_twt_session_suspend_user_callback,
+            const on_twt_session_resume& on_twt_session_resume_user_callback);
     wifi_error twtSessionSetup(const std::string& ifaceName, uint32_t cmdId,
-                               const wifi_twt_request& request,
-                               const on_twt_failure& on_twt_failure_user_callback,
-                               const on_twt_session_create& on_twt_session_create_user_callback,
-                               const on_twt_session_update& on_twt_session_update_user_callback,
-                               const on_twt_session_teardown& on_twt_session_teardown_user_callback,
-                               const on_twt_session_stats& on_twt_session_stats_user_callback,
-                               const on_twt_session_suspend& on_twt_session_suspend_user_callback,
-                               const on_twt_session_resume& on_twt_session_resume_user_callback);
+                               const wifi_twt_request& request);
     wifi_error twtSessionUpdate(const std::string& ifaceName, uint32_t cmdId, uint32_t sessionId,
                                 const wifi_twt_request& request);
     wifi_error twtSessionSuspend(const std::string& ifaceName, uint32_t cmdId, uint32_t sessionId);
@@ -797,7 +798,7 @@
     wifi_error twtSessionGetStats(const std::string& ifaceName, uint32_t cmdId, uint32_t sessionId);
 
     // Note: Following TWT functions are deprecated
-    // Deprecated
+    // Deprecated by twtRegisterEvegnts
     wifi_error twtRegisterHandler(const std::string& iface_name,
                                   const TwtCallbackHandlers& handler);
     // Deprecated by twtGetCapabilities
diff --git a/wifi/aidl/default/wifi_legacy_hal_stubs.cpp b/wifi/aidl/default/wifi_legacy_hal_stubs.cpp
index 3e4afd0..878abf0 100644
--- a/wifi/aidl/default/wifi_legacy_hal_stubs.cpp
+++ b/wifi/aidl/default/wifi_legacy_hal_stubs.cpp
@@ -182,12 +182,14 @@
     populateStubFor(&hal_fn->wifi_get_rtt_capabilities_v3);
     populateStubFor(&hal_fn->wifi_rtt_range_request_v3);
     populateStubFor(&hal_fn->wifi_twt_get_capabilities);
+    populateStubFor(&hal_fn->wifi_twt_register_events);
     populateStubFor(&hal_fn->wifi_twt_session_setup);
     populateStubFor(&hal_fn->wifi_twt_session_update);
     populateStubFor(&hal_fn->wifi_twt_session_suspend);
     populateStubFor(&hal_fn->wifi_twt_session_resume);
     populateStubFor(&hal_fn->wifi_twt_session_teardown);
     populateStubFor(&hal_fn->wifi_twt_session_get_stats);
+    populateStubFor(&hal_fn->wifi_virtual_interface_create_with_vendor_data);
     return true;
 }
 
diff --git a/wifi/aidl/default/wifi_sta_iface.cpp b/wifi/aidl/default/wifi_sta_iface.cpp
index f0509dc..85e373c 100644
--- a/wifi/aidl/default/wifi_sta_iface.cpp
+++ b/wifi/aidl/default/wifi_sta_iface.cpp
@@ -31,7 +31,11 @@
 WifiStaIface::WifiStaIface(const std::string& ifname,
                            const std::weak_ptr<legacy_hal::WifiLegacyHal> legacy_hal,
                            const std::weak_ptr<iface_util::WifiIfaceUtil> iface_util)
-    : ifname_(ifname), legacy_hal_(legacy_hal), iface_util_(iface_util), is_valid_(true) {
+    : ifname_(ifname),
+      legacy_hal_(legacy_hal),
+      iface_util_(iface_util),
+      is_valid_(true),
+      is_twt_registered_(false) {
     // Turn on DFS channel usage for STA iface.
     legacy_hal::wifi_error legacy_status = legacy_hal_.lock()->setDfsFlag(ifname_, true);
     if (legacy_status != legacy_hal::WIFI_SUCCESS) {
@@ -266,11 +270,150 @@
     return {ifname_, ndk::ScopedAStatus::ok()};
 }
 
+ndk::ScopedAStatus WifiStaIface::registerTwtEventCallbackInternal() {
+    std::weak_ptr<WifiStaIface> weak_ptr_this = weak_ptr_this_;
+
+    // onTwtFailure callback
+    const auto& on_twt_failure = [weak_ptr_this](legacy_hal::wifi_request_id id,
+                                                 legacy_hal::wifi_twt_error_code error_code) {
+        const auto shared_ptr_this = weak_ptr_this.lock();
+        IWifiStaIfaceEventCallback::TwtErrorCode aidl_error_code =
+                aidl_struct_util::convertLegacyHalTwtErrorCodeToAidl(error_code);
+        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
+            LOG(ERROR) << "Callback invoked on an invalid object";
+            return;
+        }
+        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
+            if (!callback->onTwtFailure(id, aidl_error_code).isOk()) {
+                LOG(ERROR) << "Failed to invoke onTwtFailure callback";
+            }
+        }
+    };
+    // onTwtSessionCreate callback
+    const auto& on_twt_session_create = [weak_ptr_this](legacy_hal::wifi_request_id id,
+                                                        legacy_hal::wifi_twt_session twt_session) {
+        const auto shared_ptr_this = weak_ptr_this.lock();
+        TwtSession aidl_twt_session;
+        if (!aidl_struct_util::convertLegacyHalTwtSessionToAidl(twt_session, &aidl_twt_session)) {
+            LOG(ERROR) << "convertLegacyHalTwtSessionToAidl failed";
+            return;
+        }
+
+        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
+            LOG(ERROR) << "Callback invoked on an invalid object";
+            return;
+        }
+        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
+            if (!callback->onTwtSessionCreate(id, aidl_twt_session).isOk()) {
+                LOG(ERROR) << "Failed to invoke onTwtSessionCreate callback";
+            }
+        }
+    };
+    // onTwtSessionUpdate callback
+    const auto& on_twt_session_update = [weak_ptr_this](legacy_hal::wifi_request_id id,
+                                                        legacy_hal::wifi_twt_session twt_session) {
+        const auto shared_ptr_this = weak_ptr_this.lock();
+        TwtSession aidl_twt_session;
+        if (!aidl_struct_util::convertLegacyHalTwtSessionToAidl(twt_session, &aidl_twt_session)) {
+            LOG(ERROR) << "convertLegacyHalTwtSessionToAidl failed";
+            return;
+        }
+
+        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
+            LOG(ERROR) << "Callback invoked on an invalid object";
+            return;
+        }
+        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
+            if (!callback->onTwtSessionUpdate(id, aidl_twt_session).isOk()) {
+                LOG(ERROR) << "Failed to invoke onTwtSessionUpdate callback";
+            }
+        }
+    };
+    // onTwtSessionTeardown callback
+    const auto& on_twt_session_teardown =
+            [weak_ptr_this](legacy_hal::wifi_request_id id, int session_id,
+                            legacy_hal::wifi_twt_teardown_reason_code reason_code) {
+                const auto shared_ptr_this = weak_ptr_this.lock();
+                IWifiStaIfaceEventCallback::TwtTeardownReasonCode aidl_reason_code =
+                        aidl_struct_util::convertLegacyHalTwtReasonCodeToAidl(reason_code);
+                if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
+                    LOG(ERROR) << "Callback invoked on an invalid object";
+                    return;
+                }
+                for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
+                    if (!callback->onTwtSessionTeardown(id, session_id, aidl_reason_code).isOk()) {
+                        LOG(ERROR) << "Failed to invoke onTwtSessionTeardown callback";
+                    }
+                }
+            };
+    // onTwtSessionStats callback
+    const auto& on_twt_session_stats = [weak_ptr_this](legacy_hal::wifi_request_id id,
+                                                       int session_id,
+                                                       legacy_hal::wifi_twt_session_stats stats) {
+        const auto shared_ptr_this = weak_ptr_this.lock();
+        TwtSessionStats aidl_session_stats;
+        if (!aidl_struct_util::convertLegacyHalTwtSessionStatsToAidl(stats, &aidl_session_stats)) {
+            LOG(ERROR) << "convertLegacyHalTwtSessionStatsToAidl failed";
+            return;
+        }
+        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
+            LOG(ERROR) << "Callback invoked on an invalid object";
+            return;
+        }
+        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
+            if (!callback->onTwtSessionStats(id, session_id, aidl_session_stats).isOk()) {
+                LOG(ERROR) << "Failed to invoke onTwtSessionStats callback";
+            }
+        }
+    };
+    // onTwtSessionSuspend callback
+    const auto& on_twt_session_suspend = [weak_ptr_this](legacy_hal::wifi_request_id id,
+                                                         int session_id) {
+        const auto shared_ptr_this = weak_ptr_this.lock();
+        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
+            LOG(ERROR) << "Callback invoked on an invalid object";
+            return;
+        }
+        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
+            if (!callback->onTwtSessionSuspend(id, session_id).isOk()) {
+                LOG(ERROR) << "Failed to invoke onTwtSessionSuspend callback";
+            }
+        }
+    };
+    // onTwtSessionResume callback
+    const auto& on_twt_session_resume = [weak_ptr_this](legacy_hal::wifi_request_id id,
+                                                        int session_id) {
+        const auto shared_ptr_this = weak_ptr_this.lock();
+        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
+            LOG(ERROR) << "Callback invoked on an invalid object";
+            return;
+        }
+        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
+            if (!callback->onTwtSessionResume(id, session_id).isOk()) {
+                LOG(ERROR) << "Failed to invoke onTwtSessionResume callback";
+            }
+        }
+    };
+
+    legacy_hal::wifi_error legacy_status = legacy_hal_.lock()->twtRegisterEvents(
+            ifname_, on_twt_failure, on_twt_session_create, on_twt_session_update,
+            on_twt_session_teardown, on_twt_session_stats, on_twt_session_suspend,
+            on_twt_session_resume);
+
+    if (legacy_status == legacy_hal::WIFI_ERROR_NOT_SUPPORTED) {
+        LOG(INFO) << "twtRegisterEvents is not supported" << legacy_status;
+    } else if (legacy_status != legacy_hal::WIFI_SUCCESS) {
+        LOG(ERROR) << "twtRegisterEvents failed - %d" << legacy_status;
+    }
+    return createWifiStatusFromLegacyError(legacy_status);
+}
+
 ndk::ScopedAStatus WifiStaIface::registerEventCallbackInternal(
         const std::shared_ptr<IWifiStaIfaceEventCallback>& callback) {
     if (!event_cb_handler_.addCallback(callback)) {
         return createWifiStatus(WifiStatusCode::ERROR_UNKNOWN);
     }
+    is_twt_registered_ = registerTwtEventCallbackInternal().isOk();
     return ndk::ScopedAStatus::ok();
 }
 
@@ -435,6 +578,7 @@
     } else {
         return {StaLinkLayerStats{}, createWifiStatus(WifiStatusCode::ERROR_UNKNOWN)};
     }
+    aidl_struct_util::logAidlLinkLayerStatsSize(aidl_stats);
     return {aidl_stats, ndk::ScopedAStatus::ok()};
 }
 
@@ -616,138 +760,16 @@
 
 ndk::ScopedAStatus WifiStaIface::twtSessionSetupInternal(int32_t cmdId,
                                                          const TwtRequest& aidlTwtRequest) {
+    if (!is_twt_registered_) {
+        LOG(INFO) << "twtSessionSetup is not supported as twtRegisterEvents failed";
+        return createWifiStatus(WifiStatusCode::ERROR_NOT_SUPPORTED);
+    }
     legacy_hal::wifi_twt_request legacyHalTwtRequest;
     if (!aidl_struct_util::convertAidlTwtRequestToLegacy(aidlTwtRequest, &legacyHalTwtRequest)) {
         return createWifiStatus(WifiStatusCode::ERROR_INVALID_ARGS);
     }
-    std::weak_ptr<WifiStaIface> weak_ptr_this = weak_ptr_this_;
-
-    // onTwtFailure callback
-    const auto& on_twt_failure = [weak_ptr_this](legacy_hal::wifi_request_id id,
-                                                 legacy_hal::wifi_twt_error_code error_code) {
-        const auto shared_ptr_this = weak_ptr_this.lock();
-        IWifiStaIfaceEventCallback::TwtErrorCode aidl_error_code =
-                aidl_struct_util::convertLegacyHalTwtErrorCodeToAidl(error_code);
-        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
-            LOG(ERROR) << "Callback invoked on an invalid object";
-            return;
-        }
-        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
-            if (!callback->onTwtFailure(id, aidl_error_code).isOk()) {
-                LOG(ERROR) << "Failed to invoke onTwtFailure callback";
-            }
-        }
-    };
-    // onTwtSessionCreate callback
-    const auto& on_twt_session_create = [weak_ptr_this](legacy_hal::wifi_request_id id,
-                                                        legacy_hal::wifi_twt_session twt_session) {
-        const auto shared_ptr_this = weak_ptr_this.lock();
-        TwtSession aidl_twt_session;
-        if (!aidl_struct_util::convertLegacyHalTwtSessionToAidl(twt_session, &aidl_twt_session)) {
-            LOG(ERROR) << "convertLegacyHalTwtSessionToAidl failed";
-            return;
-        }
-
-        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
-            LOG(ERROR) << "Callback invoked on an invalid object";
-            return;
-        }
-        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
-            if (!callback->onTwtSessionCreate(id, aidl_twt_session).isOk()) {
-                LOG(ERROR) << "Failed to invoke onTwtSessionCreate callback";
-            }
-        }
-    };
-    // onTwtSessionUpdate callback
-    const auto& on_twt_session_update = [weak_ptr_this](legacy_hal::wifi_request_id id,
-                                                        legacy_hal::wifi_twt_session twt_session) {
-        const auto shared_ptr_this = weak_ptr_this.lock();
-        TwtSession aidl_twt_session;
-        if (!aidl_struct_util::convertLegacyHalTwtSessionToAidl(twt_session, &aidl_twt_session)) {
-            LOG(ERROR) << "convertLegacyHalTwtSessionToAidl failed";
-            return;
-        }
-
-        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
-            LOG(ERROR) << "Callback invoked on an invalid object";
-            return;
-        }
-        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
-            if (!callback->onTwtSessionUpdate(id, aidl_twt_session).isOk()) {
-                LOG(ERROR) << "Failed to invoke onTwtSessionUpdate callback";
-            }
-        }
-    };
-    // onTwtSessionTeardown callback
-    const auto& on_twt_session_teardown =
-            [weak_ptr_this](legacy_hal::wifi_request_id id, int session_id,
-                            legacy_hal::wifi_twt_teardown_reason_code reason_code) {
-                const auto shared_ptr_this = weak_ptr_this.lock();
-                IWifiStaIfaceEventCallback::TwtTeardownReasonCode aidl_reason_code =
-                        aidl_struct_util::convertLegacyHalTwtReasonCodeToAidl(reason_code);
-                if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
-                    LOG(ERROR) << "Callback invoked on an invalid object";
-                    return;
-                }
-                for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
-                    if (!callback->onTwtSessionTeardown(id, session_id, aidl_reason_code).isOk()) {
-                        LOG(ERROR) << "Failed to invoke onTwtSessionTeardown callback";
-                    }
-                }
-            };
-    // onTwtSessionStats callback
-    const auto& on_twt_session_stats = [weak_ptr_this](legacy_hal::wifi_request_id id,
-                                                       int session_id,
-                                                       legacy_hal::wifi_twt_session_stats stats) {
-        const auto shared_ptr_this = weak_ptr_this.lock();
-        TwtSessionStats aidl_session_stats;
-        if (!aidl_struct_util::convertLegacyHalTwtSessionStatsToAidl(stats, &aidl_session_stats)) {
-            LOG(ERROR) << "convertLegacyHalTwtSessionStatsToAidl failed";
-            return;
-        }
-        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
-            LOG(ERROR) << "Callback invoked on an invalid object";
-            return;
-        }
-        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
-            if (!callback->onTwtSessionStats(id, session_id, aidl_session_stats).isOk()) {
-                LOG(ERROR) << "Failed to invoke onTwtSessionStats callback";
-            }
-        }
-    };
-    // onTwtSessionSuspend callback
-    const auto& on_twt_session_suspend = [weak_ptr_this](legacy_hal::wifi_request_id id,
-                                                         int session_id) {
-        const auto shared_ptr_this = weak_ptr_this.lock();
-        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
-            LOG(ERROR) << "Callback invoked on an invalid object";
-            return;
-        }
-        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
-            if (!callback->onTwtSessionSuspend(id, session_id).isOk()) {
-                LOG(ERROR) << "Failed to invoke onTwtSessionSuspend callback";
-            }
-        }
-    };
-    // onTwtSessionResume callback
-    const auto& on_twt_session_resume = [weak_ptr_this](legacy_hal::wifi_request_id id,
-                                                        int session_id) {
-        const auto shared_ptr_this = weak_ptr_this.lock();
-        if (!shared_ptr_this.get() || !shared_ptr_this->isValid()) {
-            LOG(ERROR) << "Callback invoked on an invalid object";
-            return;
-        }
-        for (const auto& callback : shared_ptr_this->getEventCallbacks()) {
-            if (!callback->onTwtSessionResume(id, session_id).isOk()) {
-                LOG(ERROR) << "Failed to invoke onTwtSessionResume callback";
-            }
-        }
-    };
-
-    legacy_hal::wifi_error legacy_status = legacy_hal_.lock()->twtSessionSetup(
-            ifname_, cmdId, legacyHalTwtRequest, on_twt_failure, on_twt_session_create,
-            on_twt_session_update, on_twt_session_teardown, on_twt_session_stats,
-            on_twt_session_suspend, on_twt_session_resume);
+    legacy_hal::wifi_error legacy_status =
+            legacy_hal_.lock()->twtSessionSetup(ifname_, cmdId, legacyHalTwtRequest);
     return createWifiStatusFromLegacyError(legacy_status);
 }
 
diff --git a/wifi/aidl/default/wifi_sta_iface.h b/wifi/aidl/default/wifi_sta_iface.h
index eb8f745..5713928 100644
--- a/wifi/aidl/default/wifi_sta_iface.h
+++ b/wifi/aidl/default/wifi_sta_iface.h
@@ -103,6 +103,7 @@
   private:
     // Corresponding worker functions for the AIDL methods.
     std::pair<std::string, ndk::ScopedAStatus> getNameInternal();
+    ndk::ScopedAStatus registerTwtEventCallbackInternal();
     ndk::ScopedAStatus registerEventCallbackInternal(
             const std::shared_ptr<IWifiStaIfaceEventCallback>& callback);
     std::pair<int32_t, ndk::ScopedAStatus> getFeatureSetInternal();
@@ -157,6 +158,7 @@
     std::weak_ptr<WifiStaIface> weak_ptr_this_;
     bool is_valid_;
     aidl_callback_util::AidlCallbackHandler<IWifiStaIfaceEventCallback> event_cb_handler_;
+    bool is_twt_registered_;
 
     DISALLOW_COPY_AND_ASSIGN(WifiStaIface);
 };