diff --git a/drivers/net/wireless/ath/ath5k/attach.c b/drivers/net/wireless/ath/ath5k/attach.c
index aabad4f13e2a47d3f826a5848f164d99f255b793..eb125d69fb78daafd880e57fea1aec46d1d67808 100644
--- a/drivers/net/wireless/ath/ath5k/attach.c
+++ b/drivers/net/wireless/ath/ath5k/attach.c
@@ -314,6 +314,9 @@ int ath5k_hw_attach(struct ath5k_softc *sc)
 	}
 
 	/* Crypto settings */
+	common->keymax = (sc->ah->ah_version == AR5K_AR5210 ?
+			  AR5K_KEYTABLE_SIZE_5210 : AR5K_KEYTABLE_SIZE_5211);
+
 	ah->ah_aes_support = srev >= AR5K_SREV_AR5212_V4 &&
 		(ee->ee_version >= AR5K_EEPROM_VERSION_5_0 &&
 		 !AR5K_EEPROM_AES_DIS(ee->ee_misc5));
diff --git a/drivers/net/wireless/ath/ath5k/base.c b/drivers/net/wireless/ath/ath5k/base.c
index 45789c8990d3e705fe890d4952edd903b5054de4..7933646583a9a0a77805d34a43dc03895069e55e 100644
--- a/drivers/net/wireless/ath/ath5k/base.c
+++ b/drivers/net/wireless/ath/ath5k/base.c
@@ -2570,6 +2570,7 @@ static int
 ath5k_init(struct ath5k_softc *sc)
 {
 	struct ath5k_hw *ah = sc->ah;
+	struct ath_common *common = ath5k_hw_common(ah);
 	int ret, i;
 
 	mutex_lock(&sc->lock);
@@ -2605,8 +2606,8 @@ ath5k_init(struct ath5k_softc *sc)
 	 * Reset the key cache since some parts do not reset the
 	 * contents on initial power up or resume from suspend.
 	 */
-	for (i = 0; i < AR5K_KEYTABLE_SIZE; i++)
-		ath5k_hw_reset_key(ah, i);
+	for (i = 0; i < common->keymax; i++)
+		ath_hw_keyreset(common, (u16)i);
 
 	ath5k_hw_set_ack_bitrate_high(ah, true);
 	ret = 0;
@@ -3287,9 +3288,6 @@ ath5k_set_key(struct ieee80211_hw *hw, enum set_key_cmd cmd,
 	if (modparam_nohwcrypt)
 		return -EOPNOTSUPP;
 
-	if (sc->opmode == NL80211_IFTYPE_AP)
-		return -EOPNOTSUPP;
-
 	switch (key->cipher) {
 	case WLAN_CIPHER_SUITE_WEP40:
 	case WLAN_CIPHER_SUITE_WEP104:
@@ -3298,7 +3296,6 @@ ath5k_set_key(struct ieee80211_hw *hw, enum set_key_cmd cmd,
 	case WLAN_CIPHER_SUITE_CCMP:
 		if (sc->ah->ah_aes_support)
 			break;
-
 		return -EOPNOTSUPP;
 	default:
 		WARN_ON(1);
@@ -3309,27 +3306,25 @@ ath5k_set_key(struct ieee80211_hw *hw, enum set_key_cmd cmd,
 
 	switch (cmd) {
 	case SET_KEY:
-		ret = ath5k_hw_set_key(sc->ah, key->keyidx, key,
-				       sta ? sta->addr : NULL);
-		if (ret) {
-			ATH5K_ERR(sc, "can't set the key\n");
-			goto unlock;
+		ret = ath_key_config(common, vif, sta, key);
+		if (ret >= 0) {
+			key->hw_key_idx = ret;
+			/* push IV and Michael MIC generation to stack */
+			key->flags |= IEEE80211_KEY_FLAG_GENERATE_IV;
+			if (key->cipher == WLAN_CIPHER_SUITE_TKIP)
+				key->flags |= IEEE80211_KEY_FLAG_GENERATE_MMIC;
+			if (key->cipher == WLAN_CIPHER_SUITE_CCMP)
+				key->flags |= IEEE80211_KEY_FLAG_SW_MGMT;
+			ret = 0;
 		}
-		__set_bit(key->keyidx, common->keymap);
-		key->hw_key_idx = key->keyidx;
-		key->flags |= (IEEE80211_KEY_FLAG_GENERATE_IV |
-			       IEEE80211_KEY_FLAG_GENERATE_MMIC);
 		break;
 	case DISABLE_KEY:
-		ath5k_hw_reset_key(sc->ah, key->keyidx);
-		__clear_bit(key->keyidx, common->keymap);
+		ath_key_delete(common, key);
 		break;
 	default:
 		ret = -EINVAL;
-		goto unlock;
 	}
 
-unlock:
 	mmiowb();
 	mutex_unlock(&sc->lock);
 	return ret;