diff --git a/drivers/staging/rtl8723au/core/rtw_mlme.c b/drivers/staging/rtl8723au/core/rtw_mlme.c
index b562bd6564a9..900fc520d7ce 100644
--- a/drivers/staging/rtl8723au/core/rtw_mlme.c
+++ b/drivers/staging/rtl8723au/core/rtw_mlme.c
@@ -1867,6 +1867,11 @@ int rtw_set_key23a(struct rtw_adapter *adapter,
 	struct mlme_priv *pmlmepriv = &adapter->mlmepriv;
 	int res = _SUCCESS;
 
+	if (keyid >= 4) {
+		res = _FAIL;
+		goto exit;
+	}
+
 	pcmd = (struct cmd_obj *)kzalloc(sizeof(struct cmd_obj), GFP_KERNEL);
 	if (!pcmd) {
 		res = _FAIL;  /* try again */
@@ -1893,10 +1898,10 @@ int rtw_set_key23a(struct rtw_adapter *adapter,
 			  "psecuritypriv->dot11PrivacyAlgrthm =%d\n",
 			  psetkeyparm->algorithm));
 	}
-	psetkeyparm->keyid = (u8)keyid;/* 0~3 */
+	psetkeyparm->keyid = keyid;/* 0~3 */
 	psetkeyparm->set_tx = set_tx;
 	if (is_wep_enc(psetkeyparm->algorithm))
-		pmlmepriv->key_mask |= CHKBIT(psetkeyparm->keyid);
+		pmlmepriv->key_mask |= BIT(psetkeyparm->keyid);
 
 	DBG_8723A("==> rtw_set_key23a algorithm(%x), keyid(%x), key_mask(%x)\n",
 		  psetkeyparm->algorithm, psetkeyparm->keyid,
diff --git a/drivers/staging/rtl8723au/core/rtw_pwrctrl.c b/drivers/staging/rtl8723au/core/rtw_pwrctrl.c
index 8a95ae1d353c..8755d0a56025 100644
--- a/drivers/staging/rtl8723au/core/rtw_pwrctrl.c
+++ b/drivers/staging/rtl8723au/core/rtw_pwrctrl.c
@@ -81,8 +81,8 @@ int ips_leave23a(struct rtw_adapter * padapter)
 		{
 			DBG_8723A("==>%s, channel(%d), processing(%x)\n", __func__, padapter->mlmeextpriv.cur_channel, pwrpriv->bips_processing);
 			set_channel_bwmode23a(padapter, padapter->mlmeextpriv.cur_channel, HAL_PRIME_CHNL_OFFSET_DONT_CARE, HT_CHANNEL_WIDTH_20);
-			for (keyid = 0;keyid<4;keyid++) {
-				if (pmlmepriv->key_mask & CHKBIT(keyid)) {
+			for (keyid = 0; keyid < 4; keyid++) {
+				if (pmlmepriv->key_mask & BIT(keyid)) {
 					if (keyid == psecuritypriv->dot11PrivacyKeyIndex)
 						result = rtw_set_key23a(padapter, psecuritypriv, keyid, 1);
 					else
diff --git a/drivers/staging/rtl8723au/os_dep/ioctl_cfg80211.c b/drivers/staging/rtl8723au/os_dep/ioctl_cfg80211.c
index d79d690e3302..1e3129f12f84 100644
--- a/drivers/staging/rtl8723au/os_dep/ioctl_cfg80211.c
+++ b/drivers/staging/rtl8723au/os_dep/ioctl_cfg80211.c
@@ -470,7 +470,7 @@ exit:
 }
 
 static int set_group_key(struct rtw_adapter *padapter, u8 *key, u8 alg,
-			 int keyid)
+			 u8 keyid)
 {
 	u8 keylen;
 	struct cmd_obj *pcmd;
@@ -480,6 +480,11 @@ static int set_group_key(struct rtw_adapter *padapter, u8 *key, u8 alg,
 
 	DBG_8723A("%s\n", __func__);
 
+	if (keyid >= 4) {
+		res = _FAIL;
+		goto exit;
+	}
+
 	pcmd = kzalloc(sizeof(struct cmd_obj), GFP_KERNEL);
 	if (!pcmd) {
 		res = _FAIL;
@@ -492,9 +497,9 @@ static int set_group_key(struct rtw_adapter *padapter, u8 *key, u8 alg,
 		goto exit;
 	}
 
-	psetkeyparm->keyid = (u8) keyid;
+	psetkeyparm->keyid = keyid;
 	if (is_wep_enc(alg))
-		padapter->mlmepriv.key_mask |= CHKBIT(psetkeyparm->keyid);
+		padapter->mlmepriv.key_mask |= BIT(psetkeyparm->keyid);
 
 	psetkeyparm->algorithm = alg;
 
@@ -529,7 +534,7 @@ exit:
 }
 
 static int set_wep_key(struct rtw_adapter *padapter, u8 *key, u8 keylen,
-		       int keyid)
+		       u8 keyid)
 {
 	u8 alg;
 
@@ -552,7 +557,8 @@ static int rtw_cfg80211_ap_set_encryption(struct net_device *dev,
 					  u32 param_len)
 {
 	int ret = 0;
-	u32 wep_key_idx, wep_key_len;
+	u32 wep_key_len;
+	u8 wep_key_idx;
 	struct sta_info *psta = NULL, *pbcmc_sta = NULL;
 	struct rtw_adapter *padapter = netdev_priv(dev);
 	struct mlme_priv *pmlmepriv = &padapter->mlmepriv;