Websocket.cpp

00001 #include "Websocket.h"
00002 #include "DNSClient.h"
00003 
00004 #define MAX_TRY_WRITE 20
00005 #define MAX_TRY_READ 10
00006 
00007 //Debug is disabled by default
00008 #if 1
00009 #define DBG(x, ...) std::printf("[WebSocket : DBG]"x"\r\n", ##__VA_ARGS__); 
00010 #define WARN(x, ...) std::printf("[WebSocket : WARN]"x"\r\n", ##__VA_ARGS__); 
00011 #define ERR(x, ...) std::printf("[WebSocket : ERR]"x"\r\n", ##__VA_ARGS__); 
00012 #else
00013 #define DBG(x, ...) 
00014 #define WARN(x, ...)
00015 #define ERR(x, ...) 
00016 #endif
00017 
00018 #define INFO(x, ...) printf("[WebSocket : INFO]"x"\r\n", ##__VA_ARGS__); 
00019 
00020 Websocket::Websocket(NetworkStack *ns, char * url) {
00021     m_ns = ns;
00022     fillFields(url);
00023     socket.open(ns);
00024     socket.set_blocking(false);
00025     socket.set_timeout(400);
00026 }
00027 
00028 void Websocket::fillFields(char * url) {
00029   int ret = parseURL(url, scheme, sizeof(scheme), host, sizeof(host), &port, path, sizeof(path));
00030   if(ret)
00031   {
00032     ERR("URL parsing failed; please use: \"ws://ip-or-domain[:port]/path\"");
00033     return;
00034   }
00035 
00036   if(port == 0) //TODO do handle WSS->443
00037   {
00038     port = 80;
00039   }
00040   
00041   if(strcmp(scheme, "ws"))
00042   {
00043     ERR("Wrong scheme, please use \"ws\" instead");
00044   }
00045 }
00046 
00047 int Websocket::parseURL(const char* url, char* scheme, size_t maxSchemeLen, char* host, size_t maxHostLen, uint16_t* port, char* path, size_t maxPathLen) //Parse URL
00048 {
00049   char* schemePtr = (char*) url;
00050   char* hostPtr = (char*) strstr(url, "://");
00051   if(hostPtr == NULL)
00052   {
00053     WARN("Could not find host");
00054     return -1; //URL is invalid
00055   }
00056 
00057   if( maxSchemeLen < hostPtr - schemePtr + 1 ) //including NULL-terminating char
00058   {
00059     WARN("Scheme str is too small (%d >= %d)", maxSchemeLen, hostPtr - schemePtr + 1);
00060     return -1;
00061   }
00062   memcpy(scheme, schemePtr, hostPtr - schemePtr);
00063   scheme[hostPtr - schemePtr] = '\0';
00064 
00065   hostPtr+=3;
00066 
00067   size_t hostLen = 0;
00068 
00069   char* portPtr = strchr(hostPtr, ':');
00070   if( portPtr != NULL )
00071   {
00072     hostLen = portPtr - hostPtr;
00073     portPtr++;
00074     if( sscanf(portPtr, "%hu", port) != 1)
00075     {
00076       WARN("Could not find port");
00077       return -1;
00078     }
00079   }
00080   else
00081   {
00082     *port=0;
00083   }
00084   char* pathPtr = strchr(hostPtr, '/');
00085   if( hostLen == 0 )
00086   {
00087     hostLen = pathPtr - hostPtr;
00088   }
00089 
00090   if( maxHostLen < hostLen + 1 ) //including NULL-terminating char
00091   {
00092     WARN("Host str is too small (%d >= %d)", maxHostLen, hostLen + 1);
00093     return -1;
00094   }
00095   memcpy(host, hostPtr, hostLen);
00096   host[hostLen] = '\0';
00097 
00098   size_t pathLen;
00099   char* fragmentPtr = strchr(hostPtr, '#');
00100   if(fragmentPtr != NULL)
00101   {
00102     pathLen = fragmentPtr - pathPtr;
00103   }
00104   else
00105   {
00106     pathLen = strlen(pathPtr);
00107   }
00108 
00109   if( maxPathLen < pathLen + 1 ) //including NULL-terminating char
00110   {
00111     WARN("Path str is too small (%d >= %d)", maxPathLen, pathLen + 1);
00112     return -1;
00113   }
00114   memcpy(path, pathPtr, pathLen);
00115   path[pathLen] = '\0';
00116 
00117   return 0;
00118 }
00119 
00120 
00121 bool Websocket::connect() {
00122     DNSClient dns(m_ns);
00123     char cmd[200];
00124     
00125     if (!dns.lookup(host)) {
00126         ERR("Websocket: Unable to resolve hostname (%s)", host);
00127         wait(0.2);
00128         connected = false;
00129         return false;
00130     }
00131     
00132     while (socket.connect(dns.get_ip_address(), port) < 0) {
00133         ERR("Websocket: Unable to connect to (%s) on port (%d)", host, port);
00134         wait(0.2);
00135         connected = false;
00136         return false;
00137     }
00138     
00139     connected = true;
00140 
00141     // sent http header to upgrade to the ws protocol
00142     sprintf(cmd, "GET %s HTTP/1.1\r\n", path);
00143     write(cmd, strlen(cmd));
00144     
00145     sprintf(cmd, "Host: %s:%d\r\n", host, port);
00146     write(cmd, strlen(cmd));
00147 
00148     sprintf(cmd, "Upgrade: WebSocket\r\n");
00149     write(cmd, strlen(cmd));
00150 
00151     sprintf(cmd, "Connection: Upgrade\r\n");
00152     write(cmd, strlen(cmd));
00153 
00154     sprintf(cmd, "Sec-WebSocket-Key: L159VM0TWUzyDxwJEIEzjw==\r\n");
00155     write(cmd, strlen(cmd));
00156 
00157     sprintf(cmd, "Sec-WebSocket-Version: 13\r\n\r\n");
00158     int ret = write(cmd, strlen(cmd));
00159     if (ret != strlen(cmd)) {
00160         close();
00161         ERR("Could not send request");
00162         return false;
00163     }
00164 
00165     ret = read(cmd, 200, 100);
00166     if (ret < 0) {
00167         close();
00168         ERR("Could not receive answer\r\n");
00169         return false;
00170     }
00171 
00172     cmd[ret] = '\0';
00173     DBG("recv: %s\r\n", cmd);
00174 
00175     if ( strstr(cmd, "DdLWT/1JcX+nQFHebYP+rqEx5xI=") == NULL ) {
00176         ERR("Wrong answer from server, got \"%s\" instead\r\n", cmd);
00177         do {
00178             ret = read(cmd, 200, 100);
00179             if (ret < 0) {
00180                 ERR("Could not receive answer\r\n");
00181                 return false;
00182             }
00183             cmd[ret] = '\0';
00184             printf("%s",cmd);
00185         } while (ret > 0);
00186         close();
00187         return false;
00188     }
00189 
00190     INFO("\r\nhost: %s\r\npath: %s\r\nport: %d\r\n\r\n", host, path, port);
00191     return true;
00192 }
00193 
00194 int Websocket::sendLength(uint32_t len, char * msg) {
00195 
00196     if (len < 126) {
00197         msg[0] = len | (1<<7);
00198         return 1;
00199     } else if (len < 65535) {
00200         msg[0] = 126 | (1<<7);
00201         msg[1] = (len >> 8) & 0xff;
00202         msg[2] = len & 0xff;
00203         return 3;
00204     } else {
00205         msg[0] = 127 | (1<<7);
00206         for (int i = 0; i < 8; i++) {
00207             msg[i+1] = (len >> i*8) & 0xff;
00208         }
00209         return 9;
00210     }
00211 }
00212 
00213 int Websocket::readChar(char * pC, bool block) {
00214     return read(pC, 1, 1);
00215 }
00216 
00217 int Websocket::sendOpcode(uint8_t opcode, char * msg) {
00218     msg[0] = 0x80 | (opcode & 0x0f);
00219     return 1;
00220 }
00221 
00222 int Websocket::sendMask(char * msg) {
00223     for (int i = 0; i < 4; i++) {
00224         msg[i] = 0;
00225     }
00226     return 4;
00227 }
00228 
00229 int Websocket::send(char * str) {
00230     char msg[strlen(str) + 15];
00231     int idx = 0;
00232     idx = sendOpcode(0x01, msg);
00233     idx += sendLength(strlen(str), msg + idx);
00234     idx += sendMask(msg + idx);
00235     memcpy(msg+idx, str, strlen(str));
00236     int res = write(msg, idx + strlen(str));
00237     return res;
00238 }
00239 
00240 
00241 bool Websocket::read(char * message) {
00242     int i = 0;
00243     uint32_t len_msg;
00244     char opcode = 0;
00245     char c;
00246     char mask[4] = {0, 0, 0, 0};
00247     bool is_masked = false;
00248     Timer tmr;
00249 
00250     // read the opcode
00251     tmr.start();
00252     while (true) {
00253         if (tmr.read() > 3) {
00254             DBG("timeout ws\r\n");
00255             return false;
00256         }
00257         
00258         if(!connected)
00259         {
00260             WARN("Connection was closed by server");
00261             return false;
00262         }
00263 
00264         socket.set_blocking(false);
00265         socket.set_timeout(1);
00266         if (socket.recv(&opcode, 1) != 1) {
00267             socket.set_blocking(false);
00268             socket.set_timeout(2000);
00269             return false;
00270         }
00271 
00272         socket.set_blocking(false);
00273         socket.set_timeout(2000);
00274 
00275         if (opcode == 0x81)
00276             break;
00277     }
00278     DBG("opcode: 0x%X\r\n", opcode);
00279 
00280     readChar(&c);
00281     len_msg = c & 0x7f;
00282     is_masked = c & 0x80;
00283     if (len_msg == 126) {
00284         readChar(&c);
00285         len_msg = c << 8;
00286         readChar(&c);
00287         len_msg += c;
00288     } else if (len_msg == 127) {
00289         len_msg = 0;
00290         for (int i = 0; i < 8; i++) {
00291             readChar(&c);
00292             len_msg += (c << (7-i)*8);
00293         }
00294     }
00295 
00296     if (len_msg == 0) {
00297         return false;
00298     }
00299     DBG("length: %d\r\n", len_msg);
00300     
00301     if (is_masked) {
00302         for (i = 0; i < 4; i++)
00303             readChar(&c);
00304         mask[i] = c;
00305     }
00306 
00307     int nb = read(message, len_msg, len_msg);
00308     if (nb != len_msg)
00309         return false;
00310 
00311     for (i = 0; i < len_msg; i++) {
00312         message[i] = message[i] ^ mask[i % 4];
00313     }
00314 
00315     message[len_msg] = '\0';
00316 
00317     return true;
00318 }
00319 
00320 bool Websocket::close() {
00321     if (!is_connected())
00322         return false;
00323 
00324     int ret = socket.close();
00325     if (ret < 0) {
00326         ERR("Could not disconnect");
00327         return false;
00328     }
00329     return true;
00330 }
00331 
00332 bool Websocket::is_connected() {
00333     return connected;
00334 }
00335 
00336 char* Websocket::getPath() {
00337     return path;
00338 }
00339 
00340 int Websocket::write(char * str, int len) {
00341     int res = 0, idx = 0;
00342     
00343     for (int j = 0; j < MAX_TRY_WRITE; j++) {
00344     
00345         if(!connected)
00346         {
00347             WARN("Connection was closed by server");
00348             break;
00349         }
00350 
00351         if ((res = socket.send(str + idx, len - idx)) == -1)
00352             continue;
00353 
00354         idx += res;
00355         
00356         if (idx == len)
00357             return len;
00358     }
00359     
00360     return (idx == 0) ? -1 : idx;
00361 }
00362 
00363 int Websocket::read(char * str, int len, int min_len) {
00364     int res = 0, idx = 0;
00365     
00366     for (int j = 0; j < MAX_TRY_WRITE; j++) {
00367 
00368         if ((res = socket.recv(str + idx, len - idx)) == -1)
00369             continue;
00370 
00371         idx += res;
00372         
00373         if (idx == len || (min_len != -1 && idx > min_len))
00374             return idx;
00375     }
00376     
00377     return (idx == 0) ? -1 : idx;
00378 }