MySQLClient.cpp

00001 
00002 /*
00003 Copyright (c) 2010 Donatien Garnier (donatiengar [at] gmail [dot] com)
00004  
00005 Permission is hereby granted, free of charge, to any person obtaining a copy
00006 of this software and associated documentation files (the "Software"), to deal
00007 in the Software without restriction, including without limitation the rights
00008 to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
00009 copies of the Software, and to permit persons to whom the Software is
00010 furnished to do so, subject to the following conditions:
00011  
00012 The above copyright notice and this permission notice shall be included in
00013 all copies or substantial portions of the Software.
00014  
00015 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
00016 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
00017 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
00018 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
00019 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
00020 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
00021 THE SOFTWARE.
00022 */
00023 
00024 #include "MySQLClient.h"
00025 #include "sha1.h" //For 4.1+ passwords
00026 #include "mycrypt.h" //For 4.0- passwords
00027 
00028 #define __DEBUG
00029 #include "dbg/dbg.h"
00030 
00031 #define MYSQL_TIMEOUT_MS 15000
00032 #define MYSQL_PORT 3306
00033 
00034 #define BUF_SIZE 256
00035 
00036 #define CLIENT_LONG_PASSWORD     1
00037 #define CLIENT_CONNECT_WITH_DB   8
00038 #define CLIENT_PROTOCOL_41       512
00039 #define CLIENT_INTERACTIVE         1024
00040 #define CLIENT_SECURE_CONNECTION 32768
00041 
00042 #define MIN(a,b) ((a)<(b)?(a):(b))
00043 #define ABS(a) (((a)>0)?(a):0)
00044 
00045 //MySQL commands
00046 #define COM_QUIT  0x01 //Exit
00047 #define COM_QUERY 0x03 //Execute an SQL query
00048 
00049 //#define htons( x ) ( (( x << 8 ) & 0xFF00) | (( x >> 8 ) & 0x00FF) )
00050 #define ntohs( x ) (htons(x))
00051 
00052 /*#define htonl( x ) ( (( x << 24 ) & 0xFF000000)  \
00053                    | (( x <<  8 ) & 0x00FF0000)  \
00054                    | (( x >>  8 ) & 0x0000FF00)  \
00055                    | (( x >> 24 ) & 0x000000FF)  )*/
00056 #define htonl( x ) (x)
00057 #define ntohl( x ) (htonl(x))
00058 
00059 MySQLClient::MySQLClient() : NetService(false) /*Not owned by the pool*/, m_pCbItem(NULL), m_pCbMeth(NULL), m_pCb(NULL),
00060 m_pTCPSocket(NULL), m_watchdog(), m_timeout(MYSQL_TIMEOUT_MS*1000), m_pDnsReq(NULL), m_closed(true), 
00061 m_host(), m_user(), m_password(), m_db(), m_state(MYSQL_CLOSED)
00062 {
00063   m_buf = new byte[BUF_SIZE];
00064   m_pPos = m_buf;
00065   m_len = 0;
00066   m_size = BUF_SIZE;
00067 }
00068 
00069 MySQLClient::~MySQLClient()
00070 {
00071   close();
00072   delete[] m_buf;
00073 }
00074 
00075 //High Level setup functions
00076 MySQLResult MySQLClient::open(Host& host, const string& user, const string& password, const string& db, void (*pMethod)(MySQLResult)) //Non blocking
00077 {
00078   setOnResult(pMethod);
00079   setup(host, user, password, db);
00080   return MYSQL_PROCESSING;
00081 }
00082 
00083 #if 0 //Ref only
00084 template<class T> 
00085 MySQLResult MySQLClient::open(Host& host, const string& user, const string& password, const string& db, T* pItem, void (T::*pMethod)(MySQLResult)) //Non blocking
00086 {
00087   setOnResult(pItem, pMethod);
00088   setup(host, user, password, db);
00089   return MYSQL_PROCESSING;
00090 }
00091 #endif
00092 
00093 MySQLResult MySQLClient::sql(string& sqlCommand)
00094 {
00095   if(m_state!=MYSQL_COMMANDS)
00096     return MYSQL_SETUP;
00097   sendCommand(COM_QUERY, (byte*)sqlCommand.data(), sqlCommand.length());
00098   return MYSQL_PROCESSING;
00099 }
00100 
00101 MySQLResult MySQLClient::exit()
00102 {
00103   sendCommand(COM_QUIT, NULL, 0);
00104   close();
00105   return MYSQL_OK;
00106 }
00107   
00108 void MySQLClient::setOnResult( void (*pMethod)(MySQLResult) )
00109 {
00110   m_pCb = pMethod;
00111   m_pCbItem = NULL;
00112   m_pCbMeth = NULL;
00113 }
00114 
00115 #if 0 //Ref only
00116 template<class T> 
00117 void MySQLClient::setOnResult( T* pItem, void (T::*pMethod)(MySQLResult) )
00118 {
00119   m_pCb = NULL;
00120   m_pCbItem = (CDummy*) pItem;
00121   m_pCbMeth = (void (CDummy::*)(MySQLResult)) pMethod;
00122 }
00123 #endif
00124 
00125 void MySQLClient::setTimeout(int ms)
00126 {
00127   m_timeout = 1000*ms;
00128 }
00129 
00130 void MySQLClient::poll() //Called by NetServices
00131 {
00132   if(m_closed)
00133   {
00134     return;
00135   }
00136   if(m_watchdog.read_us()>m_timeout)
00137   {
00138     onTimeout();
00139   }
00140 }
00141 
00142 void MySQLClient::resetTimeout()
00143 {
00144   m_watchdog.reset();
00145   m_watchdog.start();
00146 }
00147 
00148 void MySQLClient::init()
00149 {
00150   close(); //Remove previous elements
00151   if(!m_closed) //Already opened
00152     return;
00153   m_state = MYSQL_HANDSHAKE;
00154   m_pTCPSocket = new TCPSocket;
00155   m_pTCPSocket->setOnEvent(this, &MySQLClient::onTCPSocketEvent);
00156   m_closed = false;
00157 }
00158 
00159 void MySQLClient::close()
00160 {
00161   if(m_closed)
00162     return;
00163   m_state = MYSQL_CLOSED;
00164   m_closed = true; //Prevent recursive calling or calling on an object being destructed by someone else
00165   m_watchdog.stop(); //Stop timeout
00166   m_watchdog.reset();
00167   m_pTCPSocket->resetOnEvent();
00168   m_pTCPSocket->close();
00169   delete m_pTCPSocket;
00170   m_pTCPSocket = NULL;
00171   if( m_pDnsReq )
00172   {
00173     m_pDnsReq->close();
00174     delete m_pDnsReq;
00175     m_pDnsReq = NULL;
00176   }
00177 }
00178 
00179 void MySQLClient::setup(Host& host, const string& user, const string& password, const string& db) //Setup connection, make DNS Req if necessary
00180 {
00181   init(); //Initialize client in known state, create socket
00182   resetTimeout();
00183   m_host = host;
00184   if(!host.getPort())
00185     host.setPort( MYSQL_PORT ); //Default port
00186     
00187   m_user = user;
00188   m_password = password;
00189   
00190   m_db = db;
00191   
00192   IpAddr ip;
00193   if( !host.getIp().isNull() )
00194   {
00195     connect();
00196   }
00197   else //Need to do a DNS Query...
00198   {
00199     DBG("\r\nDNS Query...\r\n");
00200     m_pDnsReq = new DNSRequest();
00201     m_pDnsReq->setOnReply(this, &MySQLClient::onDNSReply);
00202     m_pDnsReq->resolve(&m_host);
00203     DBG("\r\nMySQLClient : DNSRequest %p\r\n", m_pDnsReq);
00204   }
00205 }
00206 
00207 void MySQLClient::connect() //Start Connection
00208 {
00209   resetTimeout();
00210   DBG("\r\nConnecting...\r\n");
00211   m_pTCPSocket->connect(m_host);
00212   m_packetId = 0;
00213 }
00214 
00215 void MySQLClient::handleHandshake()
00216 {
00217   readData();
00218   if( ! (( m_len > 1 ) && ( memchr( m_buf + 1, 0, m_len ) != NULL )) )
00219   {
00220     DBG("Connected but could not find pcsz...\n");
00221     onResult(MYSQL_PRTCL);
00222     return;
00223   }
00224   
00225   DBG("Connected to server: %d bytes read ; Protocol version %d, mysql-%s.\n", m_len, m_buf[0], &m_buf[1]);
00226   
00227   m_pPos = (byte*) memchr( (char*)(m_buf + 1), 0, m_len ) + 1;
00228   
00229   sendAuth();
00230 }
00231 
00232 void MySQLClient::sendAuth()
00233 {
00234   if( m_len - (m_pPos - m_buf) != 44)
00235   {
00236     //We only support protocol >= mysql-4.1
00237     DBG("Message after pcsz has wrong len (%d != 44)...\n", m_len - (m_pPos - m_buf));
00238     onResult(MYSQL_PRTCL);
00239     return;
00240   }
00241    
00242   uint16_t serverFlags = *((uint16_t*)&m_pPos[13]);
00243   DBG("Server capabilities are %04X.\n", serverFlags);
00244   
00245   uint32_t clientFlags = CLIENT_CONNECT_WITH_DB | CLIENT_PROTOCOL_41 | CLIENT_SECURE_CONNECTION | CLIENT_INTERACTIVE;;
00246   
00247   //if(serverFlags & CLIENT_LONG_PASSWORD)
00248   
00249   DBG("Using auth 4.1+\n");
00250   //Encrypt pw using scramble
00251   byte scramble[20+20]={0};
00252   memcpy(scramble, m_pPos+4, 8); 
00253   memcpy(scramble+8, m_pPos+31, 12); // *(m_pPos+43) == 0 (zero-terminated char*)
00254 
00255   byte stage1_hash[20] = {0};
00256   sha1( (byte*)m_password.data(), m_password.length(), stage1_hash );
00257 
00258   sha1( stage1_hash, 20, ((byte*)scramble + 20) );
00259 
00260   byte token[20] = {0};
00261   sha1( scramble, 40, token );
00262 
00263   for(int i=0;i<20;i++)
00264     token[i] = token[i] ^ stage1_hash[i];
00265     
00266   clientFlags |= CLIENT_LONG_PASSWORD;
00267   
00268   DBG("Building response\n"); 
00269   //Build response
00270 
00271   //BE
00272   *((uint32_t*)&m_buf[0]) = htonl(clientFlags);
00273   *((uint32_t*)&m_buf[4]) = BUF_SIZE; //Max packets size
00274   m_buf[8] = 8; //latin1 charset
00275   memset((char*)(m_buf+9),0,23);
00276   strcpy((char*)(m_buf+32),m_user.c_str());
00277   m_pPos = m_buf + 32 + m_user.length() + 1;
00278   m_pPos[0] = 20;
00279   memcpy((char*)&m_pPos[1],token,20);
00280   strcpy((char*)(m_pPos+21),m_db.c_str());
00281   m_len = 32 + m_user.length() + 1 + 21 + m_db.length() + 1;
00282   
00283   //Save first part of scramble in case we need it again
00284   memcpy(&m_buf[BUF_SIZE-8], scramble, 8);
00285 
00286   m_state = MYSQL_AUTH;
00287   
00288   DBG("Writing data\n"); 
00289   writeData();
00290 }
00291 
00292 void MySQLClient::handleAuthResult()
00293 {
00294   readData();
00295   if(m_len==1 && *m_buf==0xfe)
00296   {
00297     //Re-send auth using 4.0- auth
00298     sendAuth323();
00299     return;
00300   }
00301   m_watchdog.stop(); //Stop timeout
00302   m_watchdog.reset();
00303   if(m_len<2)
00304   {
00305     DBG("Response too short..\n");
00306     onResult(MYSQL_PRTCL);
00307     return;
00308   }
00309   DBG("RC=%d ",m_buf[0]);
00310   if(m_buf[0]==0)
00311   {
00312     m_buf[m_len] = 0;
00313     m_pPos = m_buf + 1;
00314     m_pPos += m_buf[1] +1;
00315     m_pPos += m_pPos[0];
00316     m_pPos += 1;
00317     DBG("(OK) : Server status %d, Message : %s\n", *((uint16_t*)&m_pPos[0]), m_pPos+4);
00318     onResult(MYSQL_OK);
00319   }
00320   else
00321   {
00322     m_buf[m_len] = 0;
00323     DBG("(Error %d) : %s\n", *((uint16_t*)&m_buf[1]), &m_buf[9]); //LE
00324     onResult(MYSQL_AUTHFAILED);
00325     return;
00326   }
00327   m_state = MYSQL_COMMANDS;
00328 }
00329 
00330 void MySQLClient::sendAuth323()
00331 {
00332   DBG("Using auth 4.0-\n");
00333   byte scramble[8]={0};
00334   
00335   memcpy(scramble, &m_buf[BUF_SIZE-8], 8); //Recover scramble
00336   
00337   //memcpy(scramble+8, m_pPos+31, 12); // *(m_pPos+43) == 0 (zero-terminated char*)
00338   
00339   byte token[9]={0};
00340   
00341   scramble_323((char*)token, (const char*)scramble, m_password.c_str());
00342   
00343   DBG("Building response\n"); 
00344   //Build response
00345   
00346   memcpy((char*)m_buf,token,9);
00347   m_len = 9;
00348 
00349   #if 0
00350   *((uint32_t*)&m_buf[0]) = htonl(clientFlags);
00351   *((uint32_t*)&m_buf[4]) = BUF_SIZE; //Max packets size
00352   m_buf[8] = 8; //latin1 charset
00353   memset((char*)(m_buf+9),0,23);
00354   strcpy((char*)(m_buf+32),m_user.c_str());
00355   m_pPos = m_buf + 32 + m_user.length() + 1;
00356   m_pPos[0] = 8;
00357   memcpy((char*)&m_pPos[1],token+1,8);
00358   strcpy((char*)(m_pPos+9),m_db.c_str());
00359   m_len = 32 + m_user.length() + 1 + 9 + m_db.length() + 1;
00360   #endif
00361   
00362   DBG("Writing data\n"); 
00363   writeData();
00364 }
00365 
00366 void MySQLClient::sendCommand(byte command, byte* arg, int len)
00367 {
00368   DBG("Sending command %d, payload of len %d\n", command, len);
00369   m_packetId=0;//Reset packet ID (New sequence)
00370   m_buf[0] = command;
00371   memcpy(&m_buf[1], arg, len);
00372   m_len = 1 + len;
00373   writeData();
00374   m_watchdog.start();
00375 }
00376 
00377 void MySQLClient::handleCommandResult()
00378 {
00379   readData();
00380   m_watchdog.stop(); //Stop timeout
00381   m_watchdog.reset();
00382   if(m_len<2)
00383   {
00384     DBG("Response too short..\n");
00385     onResult(MYSQL_PRTCL);
00386     return;
00387   }
00388   DBG("RC=%d ",m_buf[0]);
00389   if(m_buf[0]==0)
00390   {
00391     DBG("(OK)\n");
00392     onResult(MYSQL_OK);
00393   }
00394   else
00395   {
00396     m_buf[m_len] = 0;
00397     DBG("(SQL Error %d) : %s\n", *((uint16_t*)&m_buf[1]), &m_buf[9]); //LE
00398     onResult(MYSQL_SQL);
00399     return;
00400   }
00401 }
00402 
00403 void MySQLClient::readData() //Copy to buf
00404 {
00405   byte head[4];
00406   int ret = m_pTCPSocket->recv((char*)head, 4); //Packet header
00407   m_len = *((uint16_t*)&head[0]);
00408   m_packetId = head[3];
00409   DBG("Packet Id %d of length %d\n", head[3], m_len);
00410   m_packetId++;
00411   if(ret>0)
00412     ret = m_pTCPSocket->recv((char*)m_buf, m_len);
00413   if(ret < 0)//Error
00414   {
00415     onResult(MYSQL_CONN);
00416     return;
00417   }
00418   if(ret < m_len)
00419   {
00420     DBG("WARN: Incomplete packet\n");
00421   }
00422   m_len = ret;
00423 }
00424 
00425 void MySQLClient::writeData() //Copy from buf
00426 {
00427   byte head[4] = { 0 };
00428   *((uint16_t*)&head[0]) = m_len;
00429   head[3] = m_packetId;
00430   DBG("Packet Id %d\n", head[3]);
00431   m_packetId++;
00432   int ret = m_pTCPSocket->send((char*)head, 4); //Packet header
00433   if(ret>0)
00434     ret = m_pTCPSocket->send((char*)m_buf, m_len);
00435   if(ret < 0)//Error
00436   {
00437     onResult(MYSQL_CONN);
00438     return;
00439   }
00440   m_len = 0;//FIXME... incomplete packets handling
00441 }
00442 
00443 void MySQLClient::onTCPSocketEvent(TCPSocketEvent e)
00444 {
00445   DBG("\r\nEvent %d in MySQLClient::onTCPSocketEvent()\r\n", e);
00446 
00447   if(m_closed)
00448   {
00449     DBG("\r\nWARN: Discarded\r\n");
00450     return;
00451   }
00452   
00453   switch(e)
00454   {
00455   case TCPSOCKET_READABLE: //Incoming data
00456     resetTimeout();
00457     if(m_state == MYSQL_HANDSHAKE)
00458       handleHandshake();
00459     else if(m_state == MYSQL_AUTH)
00460       handleAuthResult();
00461     else if(m_state == MYSQL_COMMANDS)
00462       handleCommandResult();
00463     break;
00464   case TCPSOCKET_WRITEABLE: //We can send data
00465     resetTimeout();
00466     break;
00467   case TCPSOCKET_CONNECTED: //Connected, wait for handshake packet
00468     resetTimeout();
00469     break;
00470   case TCPSOCKET_CONTIMEOUT:
00471   case TCPSOCKET_CONRST:
00472   case TCPSOCKET_CONABRT:
00473   case TCPSOCKET_ERROR:
00474     DBG("\r\nConnection error.\r\n");
00475     onResult(MYSQL_CONN);
00476   case TCPSOCKET_DISCONNECTED:
00477     //There might still be some data available for reading
00478     //So if we are in a reading state, do not close the socket yet
00479     if(m_state != MYSQL_CLOSED)
00480     {
00481       onResult(MYSQL_CONN);
00482     }
00483     DBG("\r\nConnection closed by remote host.\r\n");
00484     break;
00485   }
00486 }
00487 
00488 void MySQLClient::onDNSReply(DNSReply r)
00489 {
00490   if(m_closed)
00491   {
00492     DBG("\r\nWARN: Discarded\r\n");
00493     return;
00494   }
00495   
00496   if( r != DNS_FOUND )
00497   {
00498     DBG("\r\nCould not resolve hostname.\r\n");
00499     onResult(MYSQL_DNS);
00500     return;
00501   }
00502   
00503   DBG("\r\nDNS Resolved to %d.%d.%d.%d.\r\n",m_host.getIp()[0],m_host.getIp()[1],m_host.getIp()[2],m_host.getIp()[3]);
00504   //If no error, m_host has been updated by m_pDnsReq so we're set to go !
00505   m_pDnsReq->close();
00506   delete m_pDnsReq;
00507   m_pDnsReq = NULL;
00508   connect();
00509 }
00510 
00511 void MySQLClient::onResult(MySQLResult r) //Called when exchange completed or on failure
00512 {
00513   if(m_pCbItem && m_pCbMeth)
00514     (m_pCbItem->*m_pCbMeth)(r);
00515   else if(m_pCb)
00516     m_pCb(r);
00517     
00518   if( (r==MYSQL_DNS) || (r==MYSQL_PRTCL) || (r==MYSQL_AUTHFAILED) || (r==MYSQL_TIMEOUT) || (r==MYSQL_CONN) ) //Fatal error, close connection
00519     close();
00520 }
00521 
00522 void MySQLClient::onTimeout() //Connection has timed out
00523 {
00524   DBG("\r\nTimed out.\n");
00525   onResult(MYSQL_DNS);
00526   close();
00527 }