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 45000
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   if( !host.getIp().isNull() )
00193   {
00194     connect();
00195   }
00196   else //Need to do a DNS Query...
00197   {
00198     DBG("DNS Query...\n");
00199     m_pDnsReq = new DNSRequest();
00200     m_pDnsReq->setOnReply(this, &MySQLClient::onDNSReply);
00201     m_pDnsReq->resolve(&m_host);
00202     DBG("MySQLClient : DNSRequest %p\n", m_pDnsReq);
00203   }
00204 }
00205 
00206 void MySQLClient::connect() //Start Connection
00207 {
00208   resetTimeout();
00209   DBG("Connecting...\n");
00210   m_pTCPSocket->connect(m_host);
00211   m_packetId = 0;
00212 }
00213 
00214 void MySQLClient::handleHandshake()
00215 {
00216   readData();
00217   if( ! (( m_len > 1 ) && ( memchr( m_buf + 1, 0, m_len ) != NULL )) )
00218   {
00219     DBG("Connected but could not find pcsz...\n");
00220     onResult(MYSQL_PRTCL);
00221     return;
00222   }
00223   
00224   DBG("Connected to server: %d bytes read ; Protocol version %d, mysql-%s.\n", m_len, m_buf[0], &m_buf[1]);
00225   
00226   m_pPos = (byte*) memchr( (char*)(m_buf + 1), 0, m_len ) + 1;
00227   
00228   sendAuth();
00229 }
00230 
00231 void MySQLClient::sendAuth()
00232 {
00233   if( m_len - (m_pPos - m_buf) != 44)
00234   {
00235     //We only support protocol >= mysql-4.1
00236     DBG("Message after pcsz has wrong len (%d != 44)...\n", m_len - (m_pPos - m_buf));
00237     onResult(MYSQL_PRTCL);
00238     return;
00239   }
00240    
00241   uint16_t serverFlags = *((uint16_t*)&m_pPos[13]);
00242   DBG("Server capabilities are %04X.\n", serverFlags);
00243   
00244   uint32_t clientFlags = CLIENT_CONNECT_WITH_DB | CLIENT_PROTOCOL_41 | CLIENT_SECURE_CONNECTION | CLIENT_INTERACTIVE;;
00245   
00246   //if(serverFlags & CLIENT_LONG_PASSWORD)
00247   
00248   DBG("Using auth 4.1+\n");
00249   //Encrypt pw using scramble
00250   byte scramble[20+20]={0};
00251   memcpy(scramble, m_pPos+4, 8); 
00252   memcpy(scramble+8, m_pPos+31, 12); // *(m_pPos+43) == 0 (zero-terminated char*)
00253 
00254   byte stage1_hash[20] = {0};
00255   sha1( (byte*)m_password.data(), m_password.length(), stage1_hash );
00256 
00257   sha1( stage1_hash, 20, ((byte*)scramble + 20) );
00258 
00259   byte token[20] = {0};
00260   sha1( scramble, 40, token );
00261 
00262   for(int i=0;i<20;i++)
00263     token[i] = token[i] ^ stage1_hash[i];
00264     
00265   clientFlags |= CLIENT_LONG_PASSWORD;
00266   
00267   DBG("Building response\n"); 
00268   //Build response
00269 
00270   //BE
00271   *((uint32_t*)&m_buf[0]) = htonl(clientFlags);
00272   *((uint32_t*)&m_buf[4]) = BUF_SIZE; //Max packets size
00273   m_buf[8] = 8; //latin1 charset
00274   memset((char*)(m_buf+9),0,23);
00275   strcpy((char*)(m_buf+32),m_user.c_str());
00276   m_pPos = m_buf + 32 + m_user.length() + 1;
00277   m_pPos[0] = 20;
00278   memcpy((char*)&m_pPos[1],token,20);
00279   strcpy((char*)(m_pPos+21),m_db.c_str());
00280   m_len = 32 + m_user.length() + 1 + 21 + m_db.length() + 1;
00281   
00282   //Save first part of scramble in case we need it again
00283   memcpy(&m_buf[BUF_SIZE-8], scramble, 8);
00284 
00285   m_state = MYSQL_AUTH;
00286   
00287   DBG("Writing data\n"); 
00288   writeData();
00289 }
00290 
00291 void MySQLClient::handleAuthResult()
00292 {
00293   readData();
00294   if(m_len==1 && *m_buf==0xfe)
00295   {
00296     //Re-send auth using 4.0- auth
00297     sendAuth323();
00298     return;
00299   }
00300   m_watchdog.stop(); //Stop timeout
00301   m_watchdog.reset();
00302   if(m_len<2)
00303   {
00304     DBG("Response too short..\n");
00305     onResult(MYSQL_PRTCL);
00306     return;
00307   }
00308   DBG("RC=%d ",m_buf[0]);
00309   if(m_buf[0]==0)
00310   {
00311     m_buf[m_len] = 0;
00312     m_pPos = m_buf + 1;
00313     m_pPos += m_buf[1] +1;
00314     m_pPos += m_pPos[0];
00315     m_pPos += 1;
00316     DBG("(OK) : Server status %d, Message : %s\n", *((uint16_t*)&m_pPos[0]), m_pPos+4);
00317     onResult(MYSQL_OK);
00318   }
00319   else
00320   {
00321     m_buf[m_len] = 0;
00322     DBG("(Error %d) : %s\n", *((uint16_t*)&m_buf[1]), &m_buf[9]); //LE
00323     onResult(MYSQL_AUTHFAILED);
00324     return;
00325   }
00326   m_state = MYSQL_COMMANDS;
00327 }
00328 
00329 void MySQLClient::sendAuth323()
00330 {
00331   DBG("Using auth 4.0-\n");
00332   byte scramble[8]={0};
00333   
00334   memcpy(scramble, &m_buf[BUF_SIZE-8], 8); //Recover scramble
00335   
00336   //memcpy(scramble+8, m_pPos+31, 12); // *(m_pPos+43) == 0 (zero-terminated char*)
00337   
00338   byte token[9]={0};
00339   
00340   scramble_323((char*)token, (const char*)scramble, m_password.c_str());
00341   
00342   DBG("Building response\n"); 
00343   //Build response
00344   
00345   memcpy((char*)m_buf,token,9);
00346   m_len = 9;
00347 
00348   #if 0
00349   *((uint32_t*)&m_buf[0]) = htonl(clientFlags);
00350   *((uint32_t*)&m_buf[4]) = BUF_SIZE; //Max packets size
00351   m_buf[8] = 8; //latin1 charset
00352   memset((char*)(m_buf+9),0,23);
00353   strcpy((char*)(m_buf+32),m_user.c_str());
00354   m_pPos = m_buf + 32 + m_user.length() + 1;
00355   m_pPos[0] = 8;
00356   memcpy((char*)&m_pPos[1],token+1,8);
00357   strcpy((char*)(m_pPos+9),m_db.c_str());
00358   m_len = 32 + m_user.length() + 1 + 9 + m_db.length() + 1;
00359   #endif
00360   
00361   DBG("Writing data\n"); 
00362   writeData();
00363 }
00364 
00365 void MySQLClient::sendCommand(byte command, byte* arg, int len)
00366 {
00367   DBG("Sending command %d, payload of len %d\n", command, len);
00368   m_packetId=0;//Reset packet ID (New sequence)
00369   m_buf[0] = command;
00370   memcpy(&m_buf[1], arg, len);
00371   m_len = 1 + len;
00372   writeData();
00373   m_watchdog.start();
00374 }
00375 
00376 void MySQLClient::handleCommandResult()
00377 {
00378   readData();
00379   m_watchdog.stop(); //Stop timeout
00380   m_watchdog.reset();
00381   if(m_len<2)
00382   {
00383     DBG("Response too short..\n");
00384     onResult(MYSQL_PRTCL);
00385     return;
00386   }
00387   DBG("RC=%d ",m_buf[0]);
00388   if(m_buf[0]==0)
00389   {
00390     DBG("(OK)\n");
00391     onResult(MYSQL_OK);
00392   }
00393   else
00394   {
00395     m_buf[m_len] = 0;
00396     DBG("(SQL Error %d) : %s\n", *((uint16_t*)&m_buf[1]), &m_buf[9]); //LE
00397     onResult(MYSQL_SQL);
00398     return;
00399   }
00400 }
00401 
00402 void MySQLClient::readData() //Copy to buf
00403 {
00404   byte head[4];
00405   int ret = m_pTCPSocket->recv((char*)head, 4); //Packet header
00406   m_len = *((uint16_t*)&head[0]);
00407   m_packetId = head[3];
00408   DBG("Packet Id %d of length %d\n", head[3], m_len);
00409   m_packetId++;
00410   if(ret>0)
00411     ret = m_pTCPSocket->recv((char*)m_buf, m_len);
00412   if(ret < 0)//Error
00413   {
00414     onResult(MYSQL_CONN);
00415     return;
00416   }
00417   if(ret < m_len)
00418   {
00419     DBG("WARN: Incomplete packet\n");
00420   }
00421   m_len = ret;
00422 }
00423 
00424 void MySQLClient::writeData() //Copy from buf
00425 {
00426   byte head[4] = { 0 };
00427   *((uint16_t*)&head[0]) = m_len;
00428   head[3] = m_packetId;
00429   DBG("Packet Id %d\n", head[3]);
00430   m_packetId++;
00431   int ret = m_pTCPSocket->send((char*)head, 4); //Packet header
00432   if(ret>0)
00433     ret = m_pTCPSocket->send((char*)m_buf, m_len);
00434   if(ret < 0)//Error
00435   {
00436     onResult(MYSQL_CONN);
00437     return;
00438   }
00439   m_len = 0;//FIXME... incomplete packets handling
00440 }
00441 
00442 void MySQLClient::onTCPSocketEvent(TCPSocketEvent e)
00443 {
00444   DBG("Event %d in MySQLClient::onTCPSocketEvent()\n", e);
00445 
00446   if(m_closed)
00447   {
00448     DBG("WARN: Discarded\n");
00449     return;
00450   }
00451   
00452   switch(e)
00453   {
00454   case TCPSOCKET_READABLE: //Incoming data
00455     resetTimeout();
00456     if(m_state == MYSQL_HANDSHAKE)
00457       handleHandshake();
00458     else if(m_state == MYSQL_AUTH)
00459       handleAuthResult();
00460     else if(m_state == MYSQL_COMMANDS)
00461       handleCommandResult();
00462     break;
00463   case TCPSOCKET_WRITEABLE: //We can send data
00464     resetTimeout();
00465     break;
00466   case TCPSOCKET_CONNECTED: //Connected, wait for handshake packet
00467     resetTimeout();
00468     break;
00469   case TCPSOCKET_CONTIMEOUT:
00470   case TCPSOCKET_CONRST:
00471   case TCPSOCKET_CONABRT:
00472   case TCPSOCKET_ERROR:
00473     DBG("Connection error.\n");
00474     onResult(MYSQL_CONN);
00475   case TCPSOCKET_DISCONNECTED:
00476     //There might still be some data available for reading
00477     //So if we are in a reading state, do not close the socket yet
00478     if(m_state != MYSQL_CLOSED)
00479     {
00480       onResult(MYSQL_CONN);
00481     }
00482     DBG("Connection closed by remote host.\n");
00483     break;
00484   }
00485 }
00486 
00487 void MySQLClient::onDNSReply(DNSReply r)
00488 {
00489   if(m_closed)
00490   {
00491     DBG("WARN: Discarded\n");
00492     return;
00493   }
00494   
00495   if( r != DNS_FOUND )
00496   {
00497     DBG("Could not resolve hostname.\n");
00498     onResult(MYSQL_DNS);
00499     return;
00500   }
00501   
00502   DBG("DNS Resolved to %d.%d.%d.%d.\n",m_host.getIp()[0],m_host.getIp()[1],m_host.getIp()[2],m_host.getIp()[3]);
00503   //If no error, m_host has been updated by m_pDnsReq so we're set to go !
00504   m_pDnsReq->close();
00505   delete m_pDnsReq;
00506   m_pDnsReq = NULL;
00507   connect();
00508 }
00509 
00510 void MySQLClient::onResult(MySQLResult r) //Called when exchange completed or on failure
00511 {
00512   if(m_pCbItem && m_pCbMeth)
00513     (m_pCbItem->*m_pCbMeth)(r);
00514   else if(m_pCb)
00515     m_pCb(r);
00516     
00517   if( (r==MYSQL_DNS) || (r==MYSQL_PRTCL) || (r==MYSQL_AUTHFAILED) || (r==MYSQL_TIMEOUT) || (r==MYSQL_CONN) ) //Fatal error, close connection
00518     close();
00519 }
00520 
00521 void MySQLClient::onTimeout() //Connection has timed out
00522 {
00523   DBG("Timed out.\n");
00524   onResult(MYSQL_TIMEOUT);
00525   close();
00526 }