/*
 ***************************************************************************
 * Ralink Tech Inc.
 * 4F, No. 2 Technology 5th Rd.
 * Science-based Industrial Park
 * Hsin-chu, Taiwan, R.O.C.
 *
 * (c) Copyright 2002-2006, Ralink Technology, Inc.
 *
 * This program is free software; you can redistribute it and/or modify  * 
 * it under the terms of the GNU General Public License as published by  * 
 * the Free Software Foundation; either version 2 of the License, or     * 
 * (at your option) any later version.                                   * 
 *                                                                       * 
 * This program is distributed in the hope that it will be useful,       * 
 * but WITHOUT ANY WARRANTY; without even the implied warranty of        * 
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         * 
 * GNU General Public License for more details.                          * 
 *                                                                       * 
 * You should have received a copy of the GNU General Public License     * 
 * along with this program; if not, write to the                         * 
 * Free Software Foundation, Inc.,                                       * 
 * 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             * 
 *                                                                       * 
 ************************************************************************

	Module Name:
	assoc.c

	Abstract:

	Revision History:
	Who			When			What
	--------	----------		----------------------------------------------
	John		2004-9-3		porting from RT2500
*/
#include "rt_config.h"
#include <net/iw_handler.h>

UCHAR	CipherWpaTemplate[] = {
		0xdd, 					// WPA IE
		0x16,					// Length
		0x00, 0x50, 0xf2, 0x01,	// oui
		0x01, 0x00,				// Version
		0x00, 0x50, 0xf2, 0x02,	// Multicast
		0x01, 0x00,				// Number of unicast
		0x00, 0x50, 0xf2, 0x02,	// unicast
		0x01, 0x00,				// number of authentication method
		0x00, 0x50, 0xf2, 0x01	// authentication
		};
UCHAR	CipherWpa2Template[] = {
		0x30,					// RSN IE
		0x14,					// Length	
		0x01, 0x00,				// Version
		0x00, 0x0f, 0xac, 0x02,	// group cipher, TKIP
		0x01, 0x00,				// number of pairwise
		0x00, 0x0f, 0xac, 0x02,	// unicast
		0x01, 0x00,				// number of authentication method
		0x00, 0x0f, 0xac, 0x02,	// authentication
		0x00, 0x00,				// RSN capability
		};
UCHAR   CipherWpa2TemplateLen = (sizeof(CipherWpa2Template) / sizeof(UCHAR));	
		
/*  
    ==========================================================================
    Description: 
        association state machine init, including state transition and timer init
    Parameters: 
        S - pointer to the association state machine
    Note:
        The state machine looks like the following 
        
                               ASSOC_IDLE               ASSOC_WAIT_RSP             REASSOC_WAIT_RSP             DISASSOC_WAIT_RSP
    MT2_MLME_ASSOC_REQ       mlme_assoc_req_action    invalid_state_when_assoc   invalid_state_when_assoc       invalid_state_when_assoc
    MT2_MLME_REASSOC_REQ     mlme_reassoc_req_action  invalid_state_when_reassoc invalid_state_when_reassoc     invalid_state_when_reassoc
    MT2_MLME_DISASSOC_REQ    mlme_disassoc_req_action mlme_disassoc_req_action   mlme_disassoc_req_action       mlme_disassoc_req_action
    MT2_PEER_DISASSOC_REQ    peer_disassoc_action     peer_disassoc_action       peer_disassoc_action           peer_disassoc_action
    MT2_PEER_ASSOC_REQ       drop                     drop                       drop                           drop
    MT2_PEER_ASSOC_RSP       drop                     peer_assoc_rsp_action      drop                           drop
    MT2_PEER_REASSOC_REQ     drop                     drop                       drop                           drop
    MT2_PEER_REASSOC_RSP     drop                     drop                       peer_reassoc_rsp_action        drop
    MT2_CLS3ERR              cls3err_action           cls3err_action             cls3err_action                 cls3err_action
    MT2_ASSOC_TIMEOUT        timer_nop                assoc_timeout_action       timer_nop                      timer_nop
    MT2_REASSOC_TIMEOUT      timer_nop                timer_nop                  reassoc_timeout_action         timer_nop
    MT2_DISASSOC_TIMEOUT     timer_nop                timer_nop                  timer_nop                      disassoc_timeout_action
    ==========================================================================
 */
VOID AssocStateMachineInit(
    IN	PRTMP_ADAPTER	pAd, 
    IN  STATE_MACHINE *S, 
    OUT STATE_MACHINE_FUNC Trans[]) 
{
    StateMachineInit(S, (STATE_MACHINE_FUNC*)Trans, MAX_ASSOC_STATE, MAX_ASSOC_MSG, (STATE_MACHINE_FUNC)Drop, ASSOC_IDLE, ASSOC_MACHINE_BASE);

    // first column
    StateMachineSetAction(S, ASSOC_IDLE, MT2_MLME_ASSOC_REQ, (STATE_MACHINE_FUNC)MlmeAssocReqAction);
    StateMachineSetAction(S, ASSOC_IDLE, MT2_MLME_REASSOC_REQ, (STATE_MACHINE_FUNC)MlmeReassocReqAction);
    StateMachineSetAction(S, ASSOC_IDLE, MT2_MLME_DISASSOC_REQ, (STATE_MACHINE_FUNC)MlmeDisassocReqAction);
    StateMachineSetAction(S, ASSOC_IDLE, MT2_PEER_DISASSOC_REQ, (STATE_MACHINE_FUNC)PeerDisassocAction);
   
    // second column
    StateMachineSetAction(S, ASSOC_WAIT_RSP, MT2_MLME_ASSOC_REQ, (STATE_MACHINE_FUNC)InvalidStateWhenAssoc);
    StateMachineSetAction(S, ASSOC_WAIT_RSP, MT2_MLME_REASSOC_REQ, (STATE_MACHINE_FUNC)InvalidStateWhenReassoc);
    StateMachineSetAction(S, ASSOC_WAIT_RSP, MT2_MLME_DISASSOC_REQ, (STATE_MACHINE_FUNC)InvalidStateWhenDisassociate);
    StateMachineSetAction(S, ASSOC_WAIT_RSP, MT2_PEER_DISASSOC_REQ, (STATE_MACHINE_FUNC)PeerDisassocAction);
    StateMachineSetAction(S, ASSOC_WAIT_RSP, MT2_PEER_ASSOC_RSP, (STATE_MACHINE_FUNC)PeerAssocRspAction);
    StateMachineSetAction(S, ASSOC_WAIT_RSP, MT2_ASSOC_TIMEOUT, (STATE_MACHINE_FUNC)AssocTimeoutAction);

    // third column
    StateMachineSetAction(S, REASSOC_WAIT_RSP, MT2_MLME_ASSOC_REQ, (STATE_MACHINE_FUNC)InvalidStateWhenAssoc);
    StateMachineSetAction(S, REASSOC_WAIT_RSP, MT2_MLME_REASSOC_REQ, (STATE_MACHINE_FUNC)InvalidStateWhenReassoc);
    StateMachineSetAction(S, REASSOC_WAIT_RSP, MT2_MLME_DISASSOC_REQ, (STATE_MACHINE_FUNC)InvalidStateWhenDisassociate);
    StateMachineSetAction(S, REASSOC_WAIT_RSP, MT2_PEER_DISASSOC_REQ, (STATE_MACHINE_FUNC)PeerDisassocAction);
    StateMachineSetAction(S, REASSOC_WAIT_RSP, MT2_PEER_REASSOC_RSP, (STATE_MACHINE_FUNC)PeerReassocRspAction);
    StateMachineSetAction(S, REASSOC_WAIT_RSP, MT2_REASSOC_TIMEOUT, (STATE_MACHINE_FUNC)ReassocTimeoutAction);

    // fourth column
    StateMachineSetAction(S, DISASSOC_WAIT_RSP, MT2_MLME_ASSOC_REQ, (STATE_MACHINE_FUNC)InvalidStateWhenAssoc);
    StateMachineSetAction(S, DISASSOC_WAIT_RSP, MT2_MLME_REASSOC_REQ, (STATE_MACHINE_FUNC)InvalidStateWhenReassoc);
    StateMachineSetAction(S, DISASSOC_WAIT_RSP, MT2_MLME_DISASSOC_REQ, (STATE_MACHINE_FUNC)InvalidStateWhenDisassociate);
    StateMachineSetAction(S, DISASSOC_WAIT_RSP, MT2_PEER_DISASSOC_REQ, (STATE_MACHINE_FUNC)PeerDisassocAction);
    StateMachineSetAction(S, DISASSOC_WAIT_RSP, MT2_DISASSOC_TIMEOUT, (STATE_MACHINE_FUNC)DisassocTimeoutAction);

    // timer init
    RTMPInitTimer(pAd, &pAd->MlmeAux.AssocTimer, &AssocTimeout);
    RTMPInitTimer(pAd, &pAd->MlmeAux.ReassocTimer, &ReassocTimeout);
    RTMPInitTimer(pAd, &pAd->MlmeAux.DisassocTimer, &DisassocTimeout);
    
}

/*
    ==========================================================================
    Description:
        Association timeout procedure. After association timeout, this function 
        will be called and it will put a message into the MLME queue
    Parameters:
        Standard timer parameters
    ==========================================================================
 */
VOID AssocTimeout(
    IN	unsigned long data) 
{
    RTMP_ADAPTER *pAd = (RTMP_ADAPTER *)data;
    DBGPRINT(RT_DEBUG_TRACE,"ASSOC - enqueue MT2_ASSOC_TIMEOUT \n");
    MlmeEnqueue(pAd, ASSOC_STATE_MACHINE, MT2_ASSOC_TIMEOUT, 0, NULL);
    RTUSBMlmeUp(pAd);
}

/*
    ==========================================================================
    Description:
        Reassociation timeout procedure. After reassociation timeout, this 
        function will be called and put a message into the MLME queue
    Parameters:
        Standard timer parameters
    ==========================================================================
 */
VOID ReassocTimeout(
    IN	unsigned long data) 
{
    RTMP_ADAPTER *pAd = (RTMP_ADAPTER *)data;
    DBGPRINT(RT_DEBUG_TRACE,"ASSOC - enqueue MT2_REASSOC_TIMEOUT \n");
    MlmeEnqueue(pAd, ASSOC_STATE_MACHINE, MT2_REASSOC_TIMEOUT, 0, NULL);
    RTUSBMlmeUp(pAd);
}

/*
    ==========================================================================
    Description:
        Disassociation timeout procedure. After disassociation timeout, this 
        function will be called and put a message into the MLME queue
    Parameters:
        Standard timer parameters
    ==========================================================================
 */
VOID DisassocTimeout(
    IN	unsigned long data) 
{
    RTMP_ADAPTER *pAd = (RTMP_ADAPTER *)data;
    DBGPRINT(RT_DEBUG_TRACE,"ASSOC - enqueue MT2_DISASSOC_TIMEOUT \n");
    MlmeEnqueue(pAd, ASSOC_STATE_MACHINE, MT2_DISASSOC_TIMEOUT, 0, NULL);
    RTUSBMlmeUp(pAd);
}

/*
	==========================================================================
	Description:
		mlme assoc req handling procedure
	Parameters:
		Adapter - Adapter pointer
		Elem - MLME Queue Element
	Pre:
		the station has been authenticated and the following information is stored in the config
			-# SSID
			-# supported rates and their length
			-# listen interval (Adapter->PortCfg.default_listen_count)
			-# Transmit power  (Adapter->PortCfg.tx_power)
	Post  :
		-# An association request frame is generated and sent to the air
		-# Association timer starts
		-# Association state -> ASSOC_WAIT_RSP
		
	==========================================================================
 */
VOID MlmeAssocReqAction(
	IN PRTMP_ADAPTER pAd, 
	IN MLME_QUEUE_ELEM *Elem) 
{
	UCHAR			 ApAddr[6];
	HEADER_802_11	 AssocHdr;
	UCHAR			 WmeIe[9] = {IE_VENDOR_SPECIFIC, 0x07, 0x00, 0x50, 0xf2, 0x02, 0x00, 0x01, STA_QOS_CAPABILITY};
	UCHAR			 CipherTmp[64];
	UCHAR			 CipherTmpLen;	
	USHORT			 ListenIntv;
	ULONG			 Timeout;
	USHORT			 CapabilityInfo;
	PUCHAR			 pOutBuffer = NULL;
	ULONG			 FrameLen = 0;
	ULONG			 tmp;
	UCHAR			 VarIesOffset;
	USHORT			 Status;
	ULONG			 idx;
	BOOLEAN			 FoundPMK = FALSE;
	USHORT			 NStatus;
	
	// Block all authentication request durning WPA block period
	if (pAd->PortCfg.bBlockAssoc == TRUE)
	{
		DBGPRINT(RT_DEBUG_TRACE, "ASSOC - Block Assoc request durning WPA block period!\n");
		pAd->Mlme.AssocMachine.CurrState = ASSOC_IDLE;
		Status = MLME_STATE_MACHINE_REJECT;
		MlmeEnqueue(pAd, MLME_CNTL_STATE_MACHINE, MT2_ASSOC_CONF, 2, &Status);
	}	
	// check sanity first
	else if (MlmeAssocReqSanity(pAd, Elem->Msg, Elem->MsgLen, ApAddr, &CapabilityInfo, &Timeout, &ListenIntv))
	{	
		RTMPCancelTimer(&pAd->MlmeAux.AssocTimer);
		COPY_MAC_ADDR(pAd->MlmeAux.Bssid, ApAddr);

		// allocate and send out AssocRsp frame
		NStatus = MlmeAllocateMemory(pAd, (PVOID)&pOutBuffer);	//Get an unused nonpaged memory
		if (NStatus != NDIS_STATUS_SUCCESS)
		{
			DBGPRINT(RT_DEBUG_TRACE,"ASSOC - MlmeAssocReqAction() allocate memory failed \n");
			pAd->Mlme.AssocMachine.CurrState = ASSOC_IDLE;
			Status = MLME_FAIL_NO_RESOURCE;
			MlmeEnqueue(pAd, MLME_CNTL_STATE_MACHINE, MT2_ASSOC_CONF, 2, &Status);
			return;
		}
		
		// Add by James 03/06/27
		pAd->PortCfg.AssocInfo.Length = sizeof(NDIS_802_11_ASSOCIATION_INFORMATION); 
		// Association don't need to report MAC address
		pAd->PortCfg.AssocInfo.AvailableRequestFixedIEs =
			NDIS_802_11_AI_REQFI_CAPABILITIES | NDIS_802_11_AI_REQFI_LISTENINTERVAL;
		pAd->PortCfg.AssocInfo.RequestFixedIEs.Capabilities = CapabilityInfo;
		pAd->PortCfg.AssocInfo.RequestFixedIEs.ListenInterval = ListenIntv;		
		// Only reassociate need this
		//COPY_MAC_ADDR(pAd->PortCfg.AssocInfo.RequestFixedIEs.CurrentAPAddress, ApAddr);
		pAd->PortCfg.AssocInfo.OffsetRequestIEs = sizeof(NDIS_802_11_ASSOCIATION_INFORMATION);
		
		// First add SSID
		VarIesOffset = 0;
		NdisMoveMemory(pAd->PortCfg.ReqVarIEs + VarIesOffset, &SsidIe, 1);
		VarIesOffset += 1;
		NdisMoveMemory(pAd->PortCfg.ReqVarIEs + VarIesOffset, &pAd->MlmeAux.SsidLen, 1);
		VarIesOffset += 1;
		NdisMoveMemory(pAd->PortCfg.ReqVarIEs + VarIesOffset, pAd->MlmeAux.Ssid, pAd->MlmeAux.SsidLen);
		VarIesOffset += pAd->MlmeAux.SsidLen;

		// Second add Supported rates
		NdisMoveMemory(pAd->PortCfg.ReqVarIEs + VarIesOffset, &SupRateIe, 1);
		VarIesOffset += 1;
		NdisMoveMemory(pAd->PortCfg.ReqVarIEs + VarIesOffset, &pAd->MlmeAux.SupRateLen, 1);
		VarIesOffset += 1;
		NdisMoveMemory(pAd->PortCfg.ReqVarIEs + VarIesOffset, pAd->MlmeAux.SupRate, pAd->MlmeAux.SupRateLen);
		VarIesOffset += pAd->MlmeAux.SupRateLen;
		// End Add by James

		DBGPRINT(RT_DEBUG_TRACE, "ASSOC - Send ASSOC request...\n");
		MgtMacHeaderInit(pAd, &AssocHdr, SUBTYPE_ASSOC_REQ, 0, ApAddr, ApAddr);
		
		// Build basic frame first
		MakeOutgoingFrame(pOutBuffer,				&FrameLen,
						  sizeof(HEADER_802_11),	&AssocHdr,
						  2,						&CapabilityInfo,
						  2,						&ListenIntv,
						  1,						&SsidIe,
						  1,						&pAd->MlmeAux.SsidLen, 
						  pAd->MlmeAux.SsidLen, 	pAd->MlmeAux.Ssid,
						  1,						&SupRateIe,
						  1,						&pAd->MlmeAux.SupRateLen,
						  pAd->MlmeAux.SupRateLen,	pAd->MlmeAux.SupRate,
						  END_OF_ARGS);

		if (pAd->MlmeAux.ExtRateLen != 0)
		{
			MakeOutgoingFrame(pOutBuffer + FrameLen,	&tmp,
							  1,						&ExtRateIe,
							  1,						&pAd->MlmeAux.ExtRateLen,
							  pAd->MlmeAux.ExtRateLen,	pAd->MlmeAux.ExtRate,							
							  END_OF_ARGS);
			FrameLen += tmp;
		}

		if (pAd->MlmeAux.APEdcaParm.bValid)
		{
			WmeIe[8] |= (pAd->MlmeAux.APEdcaParm.EdcaUpdateCount & 0x0f);
			MakeOutgoingFrame(pOutBuffer + FrameLen,	&tmp,
							  9,						&WmeIe[0],
							  END_OF_ARGS);
			FrameLen += tmp;
		}


		// For WPA / WPA-PSK
		if ((pAd->PortCfg.AuthMode == Ndis802_11AuthModeWPA) || 
			(pAd->PortCfg.AuthMode == Ndis802_11AuthModeWPAPSK))			
		{
			// Copy WPA template to buffer
			CipherTmpLen = sizeof(CipherWpaTemplate);
			NdisMoveMemory(CipherTmp, CipherWpaTemplate, CipherTmpLen);
			// Modify Group cipher
			CipherTmp[11] = ((pAd->PortCfg.GroupCipher == Ndis802_11Encryption2Enabled) ? 0x2 : 0x4);
			// Modify Pairwise cipher
			CipherTmp[17] = ((pAd->PortCfg.PairCipher == Ndis802_11Encryption2Enabled) ? 0x2 : 0x4);
			// Modify AKM
			CipherTmp[23] = ((pAd->PortCfg.AuthMode == Ndis802_11AuthModeWPA) ? 0x1 : 0x2);
			// Make outgoing frame
			MakeOutgoingFrame(pOutBuffer + FrameLen,	&tmp,
							  CipherTmpLen,				&CipherTmp[0],
							  END_OF_ARGS);
			FrameLen += tmp;
			
			// Append Variable IE
			NdisMoveMemory(pAd->PortCfg.ReqVarIEs + VarIesOffset, CipherTmp, CipherTmpLen);
			VarIesOffset += CipherTmpLen;

			// Set Variable IEs Length
			pAd->PortCfg.ReqVarIELen = VarIesOffset;
			pAd->PortCfg.AssocInfo.RequestIELength = VarIesOffset;
			// OffsetResponseIEs follow ReqVarIE
			pAd->PortCfg.AssocInfo.OffsetResponseIEs = sizeof(NDIS_802_11_ASSOCIATION_INFORMATION) + pAd->PortCfg.ReqVarIELen;
		}
		// For WPA2 / WPA2-PSK
		else if ((pAd->PortCfg.AuthMode == Ndis802_11AuthModeWPA2) || 
				 (pAd->PortCfg.AuthMode == Ndis802_11AuthModeWPA2PSK))
		{
			// Copy WPA2 template to buffer
			CipherTmpLen = sizeof(CipherWpa2Template);
			NdisMoveMemory(CipherTmp, CipherWpa2Template, CipherTmpLen);
			// Modify Group cipher
			CipherTmp[7] = ((pAd->PortCfg.GroupCipher == Ndis802_11Encryption2Enabled) ? 0x2 : 0x4);
			// Modify Pairwise cipher
			CipherTmp[13] = ((pAd->PortCfg.PairCipher == Ndis802_11Encryption2Enabled) ? 0x2 : 0x4);
			// Modify AKM
			CipherTmp[19] = ((pAd->PortCfg.AuthMode == Ndis802_11AuthModeWPA2) ? 0x1 : 0x2);
			// Check for WPA PMK cache list
			if (pAd->PortCfg.AuthMode == Ndis802_11AuthModeWPA2)
			{
				// Search chched PMKID, append it if existed
				for (idx = 0; idx < PMKID_NO; idx++)
				{
					if (NdisEqualMemory(ApAddr, &pAd->PortCfg.SavedPMK[idx].BSSID, 6))
					{
						FoundPMK = TRUE;
						break;
					}
				}

				if (FoundPMK)
				{
					// Update length within RSN IE
					CipherTmp[1] += 18;
					// Set PMK number
					*(PUSHORT) &CipherTmp[CipherTmpLen] = 1;
					NdisMoveMemory(&CipherTmp[CipherTmpLen + 2], &pAd->PortCfg.SavedPMK[idx].PMKID, 16);
					CipherTmpLen += 18;
				}
			}

			// Make outgoing frame
			MakeOutgoingFrame(pOutBuffer + FrameLen,	&tmp,
							  CipherTmpLen,				&CipherTmp[0],
							  END_OF_ARGS);
			FrameLen += tmp;

			// Append Variable IE
			NdisMoveMemory(pAd->PortCfg.ReqVarIEs + VarIesOffset, CipherTmp, CipherTmpLen);
			VarIesOffset += CipherTmpLen;

			// Set Variable IEs Length
			pAd->PortCfg.ReqVarIELen = VarIesOffset;
			pAd->PortCfg.AssocInfo.RequestIELength = VarIesOffset;
			// OffsetResponseIEs follow ReqVarIE
			pAd->PortCfg.AssocInfo.OffsetResponseIEs = sizeof(NDIS_802_11_ASSOCIATION_INFORMATION) + pAd->PortCfg.ReqVarIELen;			
		}
		else
		{
			// Do nothing