Websocket.cpp

00001 #include "Websocket.h"
00002 #include <string>
00003 
00004 #define MAX_TRY_WRITE 50
00005 #define MAX_TRY_READ 30
00006 
00007 //Debug is disabled by default
00008 #if 0
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(char * url) {
00021     fillFields(url);
00022     socket.set_blocking(false, 400);
00023 }
00024 
00025 void Websocket::fillFields(char * url) {
00026     char *res = NULL;
00027     char *res1 = NULL;
00028 
00029     char buf[50];
00030     strcpy(buf, url);
00031 
00032     res = strtok(buf, ":");
00033     if (strcmp(res, "ws")) {
00034         DBG("\r\nFormat printfor: please use: \"ws://ip-or-domain[:port]/path\"\r\n\r\n");
00035     } else {
00036         //ip_domain and port
00037         res = strtok(NULL, "/");
00038 
00039         //path
00040         res1 = strtok(NULL, " ");
00041         if (res1 != NULL) {
00042             path = res1;
00043         }
00044 
00045         //ip_domain
00046         res = strtok(res, ":");
00047 
00048         //port
00049         res1 = strtok(NULL, " ");
00050         if (res1 != NULL) {
00051             port = res1;
00052         } else {
00053             port = "80";
00054         }
00055 
00056         if (res != NULL) {
00057             ip_domain = res;
00058         }
00059     }
00060 }
00061 
00062 
00063 bool Websocket::connect() {
00064     char cmd[200];
00065 
00066     while (socket.connect(ip_domain.c_str(), atoi(port.c_str())) < 0) {
00067         ERR("Unable to connect to (%s) on port (%d)\r\n", ip_domain.c_str(), atoi(port.c_str()));
00068         wait(0.2);
00069     }
00070 
00071     // sent http header to upgrade to the ws protocol
00072     sprintf(cmd, "GET /%s HTTP/1.1\r\n", path.c_str());
00073     write(cmd, strlen(cmd));
00074 
00075     sprintf(cmd, "Host: %s:%s\r\n", ip_domain.c_str(), port.c_str());
00076     write(cmd, strlen(cmd));
00077 
00078     sprintf(cmd, "Upgrade: WebSocket\r\n");
00079     write(cmd, strlen(cmd));
00080 
00081     sprintf(cmd, "Connection: Upgrade\r\n");
00082     write(cmd, strlen(cmd));
00083 
00084     sprintf(cmd, "Sec-WebSocket-Key: L159VM0TWUzyDxwJEIEzjw==\r\n");
00085     write(cmd, strlen(cmd));
00086 
00087     sprintf(cmd, "Sec-WebSocket-Version: 13\r\n\r\n");
00088     int ret = write(cmd, strlen(cmd));
00089     if (ret != strlen(cmd)) {
00090         close();
00091         ERR("Could not send request");
00092         return false;
00093     }
00094 
00095     ret = read(cmd, 200, 100);
00096     if (ret < 0) {
00097         close();
00098         ERR("Could not receive answer\r\n");
00099         return false;
00100     }
00101 
00102     cmd[ret] = '\0';
00103     DBG("recv: %s\r\n", cmd);
00104 
00105     if ( strstr(cmd, "DdLWT/1JcX+nQFHebYP+rqEx5xI=") == NULL ) {
00106         ERR("Wrong answer from server, got \"%s\" instead\r\n", cmd);
00107         do {
00108             ret = read(cmd, 200, 100);
00109             if (ret < 0) {
00110                 ERR("Could not receive answer\r\n");
00111                 return false;
00112             }
00113             cmd[ret] = '\0';
00114             printf("%s",cmd);
00115         } while (ret > 0);
00116         close();
00117         return false;
00118     }
00119 
00120     INFO("\r\nip_domain: %s\r\npath: /%s\r\nport: %s\r\n\r\n", ip_domain.c_str(), path.c_str(), port.c_str());
00121     return true;
00122 }
00123 
00124 int Websocket::sendLength(uint32_t len, char * msg) {
00125 
00126     if (len < 126) {
00127         msg[0] = len | (1<<7);
00128         return 1;
00129     } else if (len < 65535) {
00130         msg[0] = 126 | (1<<7);
00131         msg[1] = (len >> 8) & 0xff;
00132         msg[2] = len & 0xff;
00133         return 3;
00134     } else {
00135         msg[0] = 127 | (1<<7);
00136         for (int i = 0; i < 8; i++) {
00137             msg[i+1] = (len >> i*8) & 0xff;
00138         }
00139         return 9;
00140     }
00141 }
00142 
00143 int Websocket::readChar(char * pC, bool block) {
00144     return read(pC, 1, 1);
00145 }
00146 
00147 int Websocket::sendOpcode(uint8_t opcode, char * msg) {
00148     msg[0] = 0x80 | (opcode & 0x0f);
00149     return 1;
00150 }
00151 
00152 int Websocket::sendMask(char * msg) {
00153     for (int i = 0; i < 4; i++) {
00154         msg[i] = 0;
00155     }
00156     return 4;
00157 }
00158 
00159 int Websocket::send(char * str) {
00160     char msg[strlen(str) + 15];
00161     int idx = 0;
00162     idx = sendOpcode(0x01, msg);
00163     idx += sendLength(strlen(str), msg + idx);
00164     idx += sendMask(msg + idx);
00165     memcpy(msg+idx, str, strlen(str));
00166     int res = write(msg, idx + strlen(str));
00167     return res;
00168 }
00169 
00170 
00171 bool Websocket::read(char * message) {
00172     int i = 0;
00173     uint32_t len_msg;
00174     char opcode = 0;
00175     char c;
00176     char mask[4] = {0, 0, 0, 0};
00177     bool is_masked = false;
00178     Timer tmr;
00179 
00180     // read the opcode
00181     tmr.start();
00182     while (true) {
00183         if (tmr.read() > 3) {
00184             DBG("timeout ws\r\n");
00185             return false;
00186         }
00187 
00188         socket.set_blocking(false, 1);
00189         if (socket.receive(&opcode, 1) != 1) {
00190             socket.set_blocking(false, 2000);
00191             return false;
00192         }
00193 
00194         socket.set_blocking(false, 2000);
00195 
00196         if (opcode == 0x81)
00197             break;
00198     }
00199     DBG("opcode: 0x%X\r\n", opcode);
00200 
00201     readChar(&c);
00202     len_msg = c & 0x7f;
00203     is_masked = c & 0x80;
00204     if (len_msg == 126) {
00205         readChar(&c);
00206         len_msg = c << 8;
00207         readChar(&c);
00208         len_msg += c;
00209     } else if (len_msg == 127) {
00210         len_msg = 0;
00211         for (int i = 0; i < 8; i++) {
00212             readChar(&c);
00213             len_msg += (c << (7-i)*8);
00214         }
00215     }
00216 
00217     if (len_msg == 0) {
00218         return false;
00219     }
00220     DBG("length: %d\r\n", len_msg);
00221     
00222     if (is_masked) {
00223         for (i = 0; i < 4; i++)
00224             readChar(&c);
00225         mask[i] = c;
00226     }
00227 
00228     int nb = read(message, len_msg, len_msg);
00229     if (nb != len_msg)
00230         return false;
00231 
00232     for (i = 0; i < len_msg; i++) {
00233         message[i] = message[i] ^ mask[i % 4];
00234     }
00235 
00236     message[len_msg] = '\0';
00237 
00238     return true;
00239 }
00240 
00241 bool Websocket::close() {
00242     if (!is_connected())
00243         return false;
00244 
00245     int ret = socket.close();
00246     if (ret < 0) {
00247         ERR("Could not disconnect");
00248         return false;
00249     }
00250     return true;
00251 }
00252 
00253 bool Websocket::is_connected() {
00254     return socket.is_connected();
00255 }
00256 
00257 std::string Websocket::getPath() {
00258     return path;
00259 }
00260 
00261 int Websocket::write(char * str, int len) {
00262     int res = 0, idx = 0;
00263     
00264     for (int j = 0; j < MAX_TRY_WRITE; j++) {
00265 
00266         if ((res = socket.send_all(str + idx, len - idx)) == -1)
00267             continue;
00268 
00269         idx += res;
00270         
00271         if (idx == len)
00272             return len;
00273     }
00274     
00275     return (idx == 0) ? -1 : idx;
00276 }
00277 
00278 int Websocket::read(char * str, int len, int min_len) {
00279     int res = 0, idx = 0;
00280     
00281     for (int j = 0; j < MAX_TRY_WRITE; j++) {
00282 
00283         if ((res = socket.receive_all(str + idx, len - idx)) == -1)
00284             continue;
00285 
00286         idx += res;
00287         
00288         if (idx == len || (min_len != -1 && idx > min_len))
00289             return idx;
00290     }
00291     
00292     return (idx == 0) ? -1 : idx;
00293 }