00001
00002
00003
00004
00005
00006
00007
00008
00009
00010
00011
00012
00013
00014
00015
00016
00017
00018
00019
00020
00021
00022
00023
00024
00025
00026
00027
00028
00029
00030
00031
00032
00033
00034
00035
00036 #include "stdafx.h"
00037 #include "DNSSocket.h"
00038
00039 #include "ErrorHandlerMacros.h"
00040 #include "Array_Ptr.h"
00041 #include "OSManager.h"
00042 #include "ThreadPool.h"
00043 #include "GenericCriticalSection.h"
00044
00045 #include "DNSManager.h"
00046 #include "DNSUDPSocket.h"
00047 #include "DNSTCPSocket.h"
00048
00049 #include <algorithm>
00050
00051 #ifdef _MEMORY_DEBUG
00052 #define new DEBUG_NEW
00053 #define malloc DEBUG_MALLOC
00054 static char THIS_FILE[] = __FILE__;
00055 #endif
00056
00057 namespace KomodiaDNS
00058 {
00059
00060 #define CDNSSocket_Class "CDNSSocket"
00061
00062 CDNSSocket::CDNSSocket(CDNSManager* pFather,
00063 const std::string strDNSServer,
00064 BOOL bTCP,
00065 BOOL bAsync,
00066 BOOL bAsyncConnect) : CErrorHandler(),
00067 m_pFather(pFather),
00068 m_bUseTCP(bTCP),
00069 m_bAsync(bAsync),
00070 m_strDNSServer(strDNSServer),
00071 m_pThreadManager(NULL),
00072 m_pTCP(NULL),
00073 m_pUDP(NULL),
00074 m_pCSection(NULL),
00075 m_bOwners(FALSE),
00076 m_bAsyncConnect(bAsyncConnect),
00077 m_aConnectionStatus(csNone),
00078 m_pCSectionDeque(NULL)
00079 {
00080 try
00081 {
00082
00083 SetName(CDNSSocket_Class);
00084
00085
00086 m_iTimeout=INITIAL_TIMEOUT;
00087
00088
00089 m_pCSection=COSManager::CreateCriticalSection();
00090 m_pCSectionDeque=COSManager::CreateCriticalSection();
00091 }
00092 ERROR_HANDLER("CDNSSocket")
00093 }
00094
00095 CDNSSocket::~CDNSSocket()
00096 {
00097 try
00098 {
00099
00100 std::for_each(m_aPendingData.begin(),
00101 m_aPendingData.end(),
00102 DeleteProc);
00103
00104
00105 m_aPendingData.clear();
00106
00107
00108
00109 if (m_pTCP)
00110 m_pTCP->DeleteSocketFromThread();
00111
00112 if (m_pUDP)
00113 m_pUDP->DeleteSocketFromThread();
00114
00115 if (m_bOwners)
00116
00117 delete m_pThreadManager;
00118
00119
00120 delete m_pCSection;
00121 delete m_pCSectionDeque;
00122 }
00123 ERROR_HANDLER("~CDNSSocket")
00124 }
00125
00126 void CDNSSocket::DeleteProc(WaitingRequests& rRequest)
00127 {
00128 try
00129 {
00130 delete [] rRequest.pData;
00131 }
00132 ERROR_HANDLER_STATIC(CDNSSocket_Class,"DeleteProc")
00133 }
00134
00135 BOOL CDNSSocket::Initialize()
00136 {
00137 try
00138 {
00139
00140
00141 if (m_bUseTCP)
00142 {
00143
00144 if (m_pTCP)
00145 delete m_pTCP;
00146
00147
00148 m_pTCP=new CDNSTCPSocket(this);
00149
00150
00151 if (!m_pTCP->Create())
00152 {
00153
00154 ReportError("Initialize","Failed to create socket!");
00155
00156
00157 return FALSE;
00158 }
00159
00160
00161 if (!m_bAsync)
00162 if (!m_pTCP->ReBlock())
00163 {
00164
00165 ReportError("Initialize","Failed to switch socket to blocking mode!");
00166
00167
00168 return FALSE;
00169 }
00170 else
00171 ;
00172 else
00173
00174 m_pTCP->AllowBlockedBuffer(TRUE);
00175 }
00176 else
00177 {
00178
00179 if (m_pUDP)
00180 delete m_pUDP;
00181
00182
00183 m_pUDP=new CDNSUDPSocket(this);
00184
00185
00186 if (!m_pUDP->Create())
00187 {
00188
00189 ReportError("Initialize","Failed to create socket!");
00190
00191
00192 return FALSE;
00193 }
00194
00195
00196 if (!m_bAsync)
00197 if (!m_pUDP->ReBlock())
00198 {
00199
00200 ReportError("Initialize","Failed to switch socket to blocking mode!");
00201
00202
00203 return FALSE;
00204 }
00205 else
00206 ;
00207 else
00208 if (!m_pUDP->Listen())
00209 {
00210
00211 ReportError("Initialize","Failed to listen on socket to blocking mode!");
00212
00213
00214 return FALSE;
00215 }
00216 }
00217
00218
00219 return TRUE;
00220 }
00221 ERROR_HANDLER_RETURN("Initialize",FALSE)
00222 }
00223
00224 void CDNSSocket::SetConnectionStatus(ConnectionStatus aStatus)
00225 {
00226 try
00227 {
00228
00229 CCriticalAutoRelease aRelease(m_pCSection);
00230
00231
00232 m_aConnectionStatus=aStatus;
00233 }
00234 ERROR_HANDLER("SetConnectionStatus")
00235 }
00236
00237 CThreadPool* CDNSSocket::GetThreadManager()const
00238 {
00239 return m_pThreadManager;
00240 }
00241
00242 BOOL CDNSSocket::Send(const char *pBuffer,
00243 unsigned short bufLen,
00244 LPVOID lpLocalIDData)
00245 {
00246 try
00247 {
00248
00249 int iError;
00250 iError=0;
00251
00252
00253 if (m_bUseTCP)
00254 {
00255
00256 BOOL bReconnect;
00257 bReconnect=FALSE;
00258
00259
00260 CCriticalAutoRelease aRelease(m_pCSection);
00261
00262
00263 if (m_aConnectionStatus==csDisconnected)
00264
00265 if (!Initialize())
00266 {
00267
00268 ReportError("Send","Failed to reinitialize!");
00269
00270
00271 return FALSE;
00272 }
00273 else
00274 m_aConnectionStatus=csNone;
00275
00276
00277 if (m_aConnectionStatus==csNone)
00278 {
00279
00280 SetConnectionStatus(csPending);
00281
00282
00283 aRelease.Exit();
00284
00285
00286 BOOL bResult;
00287 bResult=m_pTCP->Connect(0,
00288 m_strDNSServer.c_str(),
00289 53,
00290 !m_bAsyncConnect ||
00291 bReconnect);
00292
00293
00294 if (!bResult)
00295 {
00296
00297 SetConnectionStatus(csNone);
00298
00299
00300 ReportError("Send","Failed to connect!");
00301
00302
00303 iError=m_pTCP->GetSystemLastError();
00304 }
00305 else if (!m_bAsyncConnect ||
00306 !m_bAsync ||
00307 bReconnect)
00308
00309 SetConnectionStatus(csConnected);
00310 }
00311 else
00312 aRelease.Exit();
00313
00314
00315 if (m_aConnectionStatus!=csConnected &&
00316 !iError)
00317
00318 AddData(pBuffer,
00319 bufLen,
00320 lpLocalIDData);
00321
00322 else if (!iError &&
00323 m_aConnectionStatus==csConnected)
00324 {
00325
00326 if (!Flush())
00327 {
00328
00329 ReportError("Send","Failed to flush!");
00330
00331
00332 iError=1;
00333 }
00334 else if (!m_pTCP->Send(pBuffer,bufLen))
00335 {
00336
00337 ReportError("Send","Failed to send!");
00338
00339
00340 iError=m_pTCP->GetSystemLastError();
00341 }
00342 }
00343 }
00344 else if (m_pUDP->Send(0,
00345 m_strDNSServer.c_str(),
00346 53,
00347 pBuffer,
00348 bufLen)==CSpoofBase::GetErrorCode())
00349 {
00350
00351 ReportError("Send","Failed to send UDP packet!");
00352
00353
00354 iError=m_pUDP->GetSystemLastError();
00355 }
00356
00357
00358 if (iError)
00359
00360 return FALSE;
00361 else
00362
00363 return TRUE;
00364 }
00365 ERROR_HANDLER_RETURN("Send",FALSE)
00366 }
00367
00368 void CDNSSocket::SocketConnected(int iErrorCode)
00369 {
00370 try
00371 {
00372 {
00373
00374 CCriticalAutoRelease aRelease(m_pCSection);
00375
00376 if (iErrorCode)
00377
00378 SetConnectionStatus(csError);
00379 else
00380 SetConnectionStatus(csConnected);
00381 }
00382
00383
00384 Flush();
00385 }
00386 ERROR_HANDLER("SocketConnected")
00387 }
00388
00389 void CDNSSocket::AddData(const char* pData,
00390 unsigned short usSize,
00391 LPVOID lpData)
00392 {
00393 try
00394 {
00395
00396 if (!pData ||
00397 !usSize)
00398 {
00399
00400 ReportError("AddData","Recieved null data!");
00401
00402
00403 return;
00404 }
00405
00406
00407 WaitingRequests aData;
00408
00409
00410 aData.pData=new char[usSize];
00411 memcpy(aData.pData,
00412 pData,
00413 usSize);
00414
00415
00416 aData.lpData=lpData;
00417 aData.usSize=usSize;
00418
00419
00420 CCriticalAutoRelease aRelease(m_pCSectionDeque);
00421
00422
00423 m_aPendingData.push_back(aData);
00424 }
00425 ERROR_HANDLER("AddData")
00426 }
00427
00428 CDNSSocket::WaitingRequests CDNSSocket::GetData()const
00429 {
00430
00431 WaitingRequests aEmptyData;
00432 aEmptyData.pData=NULL;
00433 aEmptyData.usSize=0;
00434 aEmptyData.lpData=NULL;
00435
00436 try
00437 {
00438
00439 CCriticalAutoRelease aRelease(m_pCSectionDeque);
00440
00441
00442 if (m_aPendingData.empty())
00443 return aEmptyData;
00444
00445
00446 WaitingRequests aData;
00447 aData=m_aPendingData.front();
00448
00449
00450 m_aPendingData.pop_front();
00451
00452
00453 return aData;
00454 }
00455 ERROR_HANDLER_RETURN("GetData",aEmptyData)
00456 }
00457
00458 BOOL CDNSSocket::Flush()
00459 {
00460 try
00461 {
00462
00463 if (m_aConnectionStatus!=csConnected)
00464 return TRUE;
00465
00466
00467 WaitingRequests aData;
00468 aData=GetData();
00469
00470
00471 while (aData.pData)
00472 {
00473
00474 CArray_ptr<char> pProtection(aData.pData);
00475
00476
00477 if (!m_pTCP->Send(aData.pData,
00478 aData.usSize))
00479 {
00480
00481 ReportError("Flush","Failed to send!");
00482
00483
00484 return FALSE;
00485 }
00486
00487
00488 if (m_aConnectionStatus!=csConnected)
00489 aData.pData=NULL;
00490 else
00491 aData=GetData();
00492 }
00493
00494
00495 return TRUE;
00496 }
00497 ERROR_HANDLER_RETURN("Flush",FALSE)
00498 }
00499
00500 CDNSAnswers* CDNSSocket::Receive()
00501 {
00502 if (m_bAsync)
00503 return NULL;
00504
00505 try
00506 {
00507
00508 char* pBuffer;
00509 pBuffer=new char[2000];
00510
00511
00512 int iReceived=0;
00513
00514 if (m_bUseTCP)
00515 iReceived=m_pTCP->Receive(pBuffer,
00516 2000);
00517 else
00518 iReceived=m_pUDP->Receive(pBuffer,
00519 2000);
00520
00521
00522 if (iReceived==CSpoofBase::GetErrorCode() ||
00523 !iReceived)
00524 {
00525
00526 delete [] pBuffer;
00527
00528
00529 return FALSE;
00530 }
00531
00532
00533 return SocketReceive(pBuffer,
00534 iReceived,
00535 m_bUseTCP);
00536 }
00537 ERROR_HANDLER_RETURN("Receive",NULL)
00538 }
00539
00540 BOOL CDNSSocket::SetConnectionTimeout(int iMS)
00541 {
00542 try
00543 {
00544
00545 m_iTimeout=iMS;
00546
00547
00548 if (m_bUseTCP)
00549 return m_pTCP->SetConnectionTimeout(iMS);
00550 else
00551 return TRUE;
00552 }
00553 ERROR_HANDLER_RETURN("SetTimeout",FALSE)
00554 }
00555
00556 CDNSAnswers* CDNSSocket::ParseAnswer(char* pBuffer,
00557 char* pOriginalBuffer,
00558 int iBufferLengh,
00559 BOOL bFromThread,
00560 BOOL bTCP)
00561 {
00562 try
00563 {
00564
00565 if (m_pThreadManager &&
00566 !bFromThread)
00567 {
00568
00569 if (iBufferLengh<=0)
00570 {
00571
00572 ReportError("ParseAnswer","Invalid buffer length!");
00573
00574
00575 return NULL;
00576 }
00577
00578
00579 ParseThreadData* pData;
00580 pData=new ParseThreadData;
00581
00582
00583 pData->pData=pBuffer;
00584 pData->pOriginalData=pOriginalBuffer;
00585 pData->iDataSize=iBufferLengh;
00586 pData->pSocket=this;
00587 pData->bTCP=bTCP;
00588
00589
00590 m_pThreadManager->SubmitJob(ThreadProc,
00591 (LPVOID)pData);
00592
00593
00594 return NULL;
00595 }
00596
00597
00598 CDNSAnswers* pAnswers;
00599 pAnswers=new CDNSAnswers;
00600
00601
00602 int iResult;
00603 iResult=pAnswers->Parse(pBuffer,
00604 iBufferLengh);
00605
00606 if (!iResult)
00607 {
00608
00609 delete pAnswers;
00610 pAnswers=NULL;
00611
00612
00613 ReportError("ParseAnswer","Failed to parse answers");
00614
00615
00616 delete [] pOriginalBuffer;
00617
00618
00619 return NULL;
00620 }
00621
00622 if (m_bAsync)
00623
00624 m_pFather->OnDNSReceive(pAnswers);
00625
00626 if (iResult<iBufferLengh)
00627 {
00628
00629 if (bTCP)
00630 iResult+=2;
00631
00632
00633 iBufferLengh-=iResult;
00634 pBuffer+=iResult;
00635
00636
00637 return ParseAnswer(pBuffer,
00638 pOriginalBuffer,
00639 iBufferLengh,
00640 TRUE,
00641 bTCP);
00642 }
00643 else
00644 delete [] pOriginalBuffer;
00645
00646
00647 return pAnswers;
00648 }
00649 ERROR_HANDLER_RETURN("ParseAnswer",NULL)
00650 }
00651
00652 void CDNSSocket::SocketError(int iErrorCode)
00653 {
00654 try
00655 {
00656
00657 m_pFather->OnDNSError(iErrorCode,
00658 NULL);
00659 }
00660 ERROR_HANDLER("SocketError")
00661 }
00662
00663 void CDNSSocket::SocketClosed()
00664 {
00665 try
00666 {
00667
00668 SetConnectionStatus(csDisconnected);
00669 }
00670 ERROR_HANDLER("SocketClosed")
00671 }
00672
00673 CDNSAnswers* CDNSSocket::SocketReceive(char* pBuffer,
00674 int iBufferLengh,
00675 BOOL bTCP)
00676 {
00677 try
00678 {
00679
00680 char* pOriginalBuffer;
00681 pOriginalBuffer=pBuffer;
00682
00683
00684 if (bTCP)
00685 {
00686
00687 pBuffer+=2;
00688 iBufferLengh-=2;
00689
00690
00691 if (!Flush())
00692
00693 ReportError("SocketReceive","Failed to flush!");
00694 }
00695
00696
00697 return ParseAnswer(pBuffer,
00698 pOriginalBuffer,
00699 iBufferLengh,
00700 !m_bAsync,
00701 bTCP);
00702 }
00703 ERROR_HANDLER_RETURN("SocketReceive",NULL)
00704 }
00705
00706 void CDNSSocket::SetMultithreaded(int iThreadNumber)
00707 {
00708 try
00709 {
00710
00711 if (m_pThreadManager)
00712 return;
00713
00714
00715 m_pThreadManager=new CThreadPool(iThreadNumber);
00716
00717
00718 m_bOwners=TRUE;
00719 }
00720 ERROR_HANDLER("SetMultithreaded")
00721 }
00722
00723 void CDNSSocket::SetMultithreaded(CThreadPool* pThreadManager)
00724 {
00725 try
00726 {
00727
00728 if (m_pThreadManager)
00729 delete m_pThreadManager;
00730
00731
00732 m_pThreadManager=pThreadManager;
00733
00734
00735 m_bOwners=FALSE;
00736 }
00737 ERROR_HANDLER("SetMultithreaded")
00738 }
00739
00740 void CDNSSocket::ThreadProc(LPVOID lpParam)
00741 {
00742 try
00743 {
00744
00745 ParseThreadData* pData;
00746 pData=(ParseThreadData*)lpParam;
00747
00748
00749 pData->pSocket->ParseAnswer(pData->pData,
00750 pData->pOriginalData,
00751 pData->iDataSize,
00752 TRUE,
00753 pData->bTCP);
00754
00755
00756 delete pData;
00757 }
00758 ERROR_HANDLER_STATIC(CDNSSocket_Class,"ThreadProc")
00759 }
00760
00761
00762 }