245 lines
4.7 KiB
C++
245 lines
4.7 KiB
C++
#include "websocket.h"
|
||
#include "gsock_helper.h"
|
||
#include <vector>
|
||
#include <queue>
|
||
#include "base64.hpp"
|
||
#include "sha1.h"
|
||
#include <WinSock2.h>
|
||
using namespace std;
|
||
|
||
int HandleKey(sock& s, const string& key)
|
||
{
|
||
string server_key = key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
|
||
SHA1 sha;
|
||
sha.update(server_key);
|
||
uint32_t arr[5];
|
||
sha.final(arr);
|
||
for (int i = 0; i < 5; i++) arr[i] = htonl(arr[i]); // ÖØÒª
|
||
string response_key = base64_encode_std((const unsigned char*)arr, sizeof(arr));
|
||
sock_helper sp(s);
|
||
string response_header = string("HTTP/1.1 101 Switching Protocols\r\n") +
|
||
"Connection: upgrade\r\n" +
|
||
"Sec-WebSocket-Accept: " + response_key + "\r\n" +
|
||
"Upgrade: websocket\r\n\r\n";
|
||
if (sp.sendall(response_header) <= 0)
|
||
{
|
||
return -1;
|
||
}
|
||
else
|
||
{
|
||
return 0;
|
||
}
|
||
}
|
||
|
||
int Handshake(sock& s)
|
||
{
|
||
sock_helper sp(s);
|
||
vector<string> lines;
|
||
while (true)
|
||
{
|
||
string str;
|
||
int ret = sp.recvline(str);
|
||
if (ret <= 0) return -1;
|
||
if (str.empty()) break;
|
||
lines.push_back(str);
|
||
// cout << str << endl;
|
||
}
|
||
|
||
for (auto& str : lines)
|
||
{
|
||
string target("Sec-WebSocket-Key");
|
||
if (str.find(target) != string::npos)
|
||
{
|
||
string key = str.substr(str.find(target) + 19, 24);
|
||
if (HandleKey(s, key) == 0)
|
||
{
|
||
// Handshake OK
|
||
return 0;
|
||
}
|
||
else
|
||
{
|
||
// Handshake Failed: Key is not ok
|
||
return -1;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Not websocket protocol
|
||
return -2;
|
||
}
|
||
|
||
|
||
|
||
int ReadFrame(sock& s, WSFrame& f)
|
||
{
|
||
sock_helper sp(s);
|
||
unsigned char c;
|
||
if (s.recv(&c, 1) <= 0) return -1;
|
||
f.fin = c & 0x80; // 1000 0000
|
||
f.rsv1 = c & 0x40; // 0100 0000
|
||
f.rsv2 = c & 0x20; // 0010 0000
|
||
f.rsv3 = c & 0x10; // 0001 0000
|
||
f.opcode = c & 0xF; // 0000 1111
|
||
|
||
if (s.recv(&c, 1) <= 0) return -1;
|
||
f.ismask = c & 0x80; // 1000 0000
|
||
int payload_head = c & 0x7F; // 0111 1111
|
||
if (payload_head < 126)
|
||
{
|
||
f.len = payload_head;
|
||
}
|
||
else if (payload_head == 126)
|
||
{
|
||
uint16_t x;
|
||
if (sp.recvall(&x, sizeof(x)) <= 0) return -1;
|
||
f.len = ntohs(x);
|
||
}
|
||
else if (payload_head == 127)
|
||
{
|
||
uint64_t x;
|
||
if (sp.recvall(&x, sizeof(x)) <= 0) return -1;
|
||
f.len = ntohll(x);
|
||
}
|
||
if (f.ismask)
|
||
{
|
||
if (sp.recvall(f.mask, sizeof(f.mask)) <= 0) return -1;
|
||
}
|
||
unique_ptr<char[]> xp(new char[f.len]);
|
||
if (xp.get() == nullptr) return -2;
|
||
memset(xp.get(), 0, f.len);
|
||
if (sp.recvall(xp.get(), f.len) <= 0) return -1;
|
||
if (f.ismask)
|
||
{
|
||
// ´¦ÀíÑÚÂëÎÊÌâ
|
||
for (unsigned long long i = 0; i < f.len; i++)
|
||
{
|
||
xp[i] ^= f.mask[i % 4];
|
||
}
|
||
}
|
||
f.data = string(xp.get(), f.len);
|
||
|
||
return 0;
|
||
}
|
||
|
||
int SendFrame(sock& s, const WSFrame& f)
|
||
{
|
||
sock_helper sp(s);
|
||
unsigned char c = 0;
|
||
if (f.fin) c = 0x80; // 1000 0000
|
||
if (f.rsv1) c |= 0x40; // 0100 0000
|
||
if (f.rsv2) c |= 0x20; // 0010 0000
|
||
if (f.rsv3) c |= 0x10; // 0001 0000
|
||
c |= (f.opcode & 0xF); // 0000 1111
|
||
if (s.send(&c, 1) <= 0) return -1;
|
||
|
||
if (f.ismask) c = 0x80; // 1000 0000
|
||
else c = 0;
|
||
|
||
if (f.len < 126)
|
||
{
|
||
c |= (f.len & 0x7F);
|
||
if (s.send(&c, 1) <= 0) return -1;
|
||
}
|
||
else if (f.len < (2 << 16))
|
||
{
|
||
c |= 126;
|
||
if (s.send(&c, 1) <= 0) return -1;
|
||
uint16_t len = f.len;
|
||
if (sp.sendall(&len, sizeof(len)) <= 0) return -1;
|
||
}
|
||
else
|
||
{
|
||
c |= 127;
|
||
if (s.send(&c, 1) <= 0) return -1;
|
||
uint64_t len = f.len;
|
||
if (sp.sendall(&len, sizeof(len)) <= 0) return -1;
|
||
}
|
||
if (f.ismask)
|
||
{
|
||
if (sp.sendall(&(f.mask), sizeof(f.mask)) <= 0) return -1;
|
||
vector<char> vec(f.data.begin(), f.data.end());
|
||
for (uint64_t i = 0; i < f.len; i++)
|
||
{
|
||
vec[i] ^= f.mask[i % 4];
|
||
}
|
||
if (sp.sendall(vec.data(), vec.size()) <= 0) return -1;
|
||
}
|
||
else
|
||
{
|
||
if (sp.sendall(f.data) <= 0) return -1;
|
||
}
|
||
|
||
return 0;
|
||
}
|
||
|
||
int SendPong(sock& s, const WSFrame& ping)
|
||
{
|
||
cout << "Sending pong to " << (&s) << endl;
|
||
|
||
WSFrame f;
|
||
f.fin = true;
|
||
f.rsv1 = false;
|
||
f.rsv2 = false;
|
||
f.rsv3 = false;
|
||
f.ismask = false;
|
||
f.len = ping.len;
|
||
f.data = ping.data;
|
||
f.opcode = 0xA;
|
||
|
||
return SendFrame(s, f);
|
||
}
|
||
|
||
int ReadMsg(sock& s, string& out_data)
|
||
{
|
||
string data;
|
||
WSFrame f;
|
||
int pack_cnt = 0;
|
||
while (true)
|
||
{
|
||
int ret = ReadFrame(s, f);
|
||
if (ret < 0) return -1;
|
||
if (f.opcode == 0x9) // ping°üÔò»Ø¸´Ò»¸öpong°ü
|
||
{
|
||
cout << "Received ping from " << (&s) << endl;
|
||
if (SendPong(s, f) < 0) return -2;
|
||
continue;
|
||
}
|
||
else if (f.opcode == 0xA) // pong°üÔòʲô¶¼²»×ö
|
||
{
|
||
cout << "Received pong from " << (&s) << endl;
|
||
continue;
|
||
}
|
||
else if (f.opcode == 0x8) // Á¬½ÓÖ÷¶¯¹Ø±Õ
|
||
{
|
||
cout << "Websocket is closing " << (&s) << endl;
|
||
return 0;
|
||
}
|
||
pack_cnt++;
|
||
data.append(f.data);
|
||
if (f.fin) break;
|
||
}
|
||
out_data = data;
|
||
return pack_cnt;
|
||
}
|
||
|
||
int SendMsg(sock& s, const string& data, bool is_text)
|
||
{
|
||
WSFrame f;
|
||
f.fin = true;
|
||
f.rsv1 = false;
|
||
f.rsv2 = false;
|
||
f.rsv3 = false;
|
||
f.ismask = false;
|
||
f.len = data.size();
|
||
f.data = data;
|
||
f.opcode = is_text ? 0x1 : 0x2;
|
||
|
||
if (SendFrame(s, f) < 0)
|
||
{
|
||
return -1;
|
||
}
|
||
else
|
||
{
|
||
return 1;
|
||
}
|
||
} |