c++ IOCP websocket 代码-Go服务器开发

以下是一个使用 C++ 和 IOCP 实现的简单 WebSocket 服务器示例代码：

#include <iostream>
#include <vector>
#include <string>
#include <winsock2.h>
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")

// 定义 WebSocket 帧类型
enum Opcode {
    CONTINUATION_FRAME = 0x00,
    TEXT_FRAME = 0x01,
    BINARY_FRAME = 0x02,
    CLOSE_FRAME = 0x08,
    PING_FRAME = 0x09,
    PONG_FRAME = 0x0A
};

// 定义 WebSocket 帧头部结构体
struct FrameHeader {
    bool fin;
    Opcode opcode;
    bool mask;
    uint64_t payloadLength;
};

// 解析 WebSocket 帧头部信息
FrameHeader ParseFrameHeader(const char* buffer, size_t length) {
    FrameHeader header;

    // 判断 FIN 标志位
    header.fin = (buffer[0] & 0x80) != 0;

	// 获取操作码
	header.opcode = static_cast<Opcode>(buffer[0] & 0x0F);

	// 判断是否使用掩码
	header.mask = (buffer[1] & 0x80) != 0;

	// 获取负载数据长度
	uint8_t payloadLenByte = buffer[1] & 0x7F;

	if (payloadLenByte == 126) {
	    if (length < sizeof(FrameHeader) + sizeof(uint16_t)) {
	        throw std::runtime_error("Invalid frame: insufficient data");
	    }
	    
	    const uint16_t* pPayloadLength16Bit =
	        reinterpret_cast<const uint16_t*>(buffer + sizeof(FrameHeader));
	        
	    header.payloadLength = ntohs(*pPayloadLength16Bit);
	} else if (payloadLenByte == 127) {
	    if (length < sizeof(FrameHeader) + sizeof(uint64_t)) {
	        throw std::runtime_error("Invalid frame: insufficient data");
	    }
	    
	    const uint64_t* pPayloadLength64Bit =
	        reinterpret_cast<const uint64_t*>(buffer + sizeof(FrameHeader));
	        
	    header.payloadLength = be64toh(*pPayloadLength64Bit);
	} else {
	    header.payloadLength = payloadLenByte;
	}

    return header;
}

// 解析 WebSocket 帧负载数据
std::vector<char> ParseFramePayload(const char* buffer, size_t length, const FrameHeader& header) {
    std::vector<char> payload;

    // 获取掩码
    const char* maskKey = nullptr;
    
	if (header.mask) {
	    maskKey = buffer + sizeof(FrameHeader);
	}
	
	// 计算负载数据起始位置和长度
	size_t payloadOffset = sizeof(FrameHeader);
	size_t payloadLength = static_cast<size_t>(header.payloadLength);

	if (header.mask) {
	    payloadOffset += 4;
	    payload.resize(payloadLength);
	    
	    for (size_t i = 0; i < payloadLength; ++i) {
	        size_t j = i % 4;
	        payload[i] = buffer[payloadOffset + i] ^ maskKey[j];
	    }
	} else {
        // 如果没有使用掩码，直接拷贝到 vector 中
        // 在实际应用中，不推荐这种做法，因为可能存在安全隐患
        // 推荐使用异或操作解密掩码后的数据
        payload.assign(buffer + payloadOffset, buffer + length);
	}

    return payload;
}

// WebSocket 服务器的 IOCP 回调函数
VOID CALLBACK ServerIocpCallback(DWORD errorCode,
                                  DWORD numberOfBytesTransferred,
                                  LPOVERLAPPED overlapped) {
    // 获取当前连接的 SOCKET 和缓冲区指针
    SOCKET clientSocket = reinterpret_cast<SOCKET>(overlapped->hEvent);
    char* recvBuffer = reinterpret_cast<char*>(overlapped->Pointer);

    if (errorCode != 0 || numberOfBytesTransferred == 0) {
        // 如果发生错误或者客户端主动关闭连接，关闭 SOCKET 并释放缓冲区内存
        closesocket(clientSocket);
        delete[] recvBuffer;

        std::cout << "Client disconnected" << std::endl;
        return;
    }

    try {
        // 解析 WebSocket 帧头部信息
        FrameHeader header = ParseFrameHeader(recvBuffer, numberOfBytesTransferred);

        // 解析负载数据
        std::vector<char> payload = ParseFramePayload(recvBuffer, numberOfBytesTransferred, header);

        // 处理收到的数据，这里只是简单地将接收到的数据返回给客户端
        std::string response(payload.begin(), payload.end());
        
	    const char* responseBuffer = response.c_str();
	    size_t responseLength = response.size();

	    send(clientSocket, responseBuffer, static_cast<int>(responseLength), 0);

	} catch (const std::exception& e) {
	    std::cerr << e.what() << std::endl;

	    // 如果解析失败，关闭 SOCKET 并释放缓冲区内存
	    closesocket(clientSocket);
	    delete[] recvBuffer;
	}

    // 重新投递接收请求
    WSABUF recvWsaBuf;
    recvWsaBuf.buf = recvBuffer;
    recvWsaBuf.len = static_cast<ULONG>(numberOfBytesTransferred);

	OVERLAPPED* pOverlapped = new OVERLAPPED;
	memset(pOverlapped, 0, sizeof(OVERLAPPED));
	pOverlapped->hEvent = reinterpret_cast<HANDLE>(clientSocket);
	pOverlapped->Pointer = recvBuffer;

	WSARecv(clientSocket,
	        &recvWsaBuf,
	        1,
	        nullptr,
	        reinterpret_cast<DWORD*>(pOverlapped),
	        nullptr);
}

int main() {
    // 初始化 Winsock2 库
    WSADATA wsaData;
    WSAStartup(MAKEWORD(2, 2), &wsaData);

	// 创建 SOCKET，绑定地址和端口，并开始监听连接请求
	SOCKET listenSocket = socket(AF_INET6, SOCK_STREAM, IPPROTO_TCP);
	
	sockaddr_in6 serverAddr{};
	serverAddr.sin6_family = AF_INET6;
	serverAddr.sin6_addr = in6addr_any;
	serverAddr.sin6_port = htons(12345);
	
	bind(listenSocket, reinterpret_cast<SOCKADDR*>(&serverAddr), sizeof(serverAddr));
	listen(listenSocket, SOMAXCONN);

	std::cout << "Listening on port 12345" << std::endl;

	// 创建完成端口对象，并将 listenSocket 加入 IOCP 中等待连接请求
	HANDLE iocpHandle =
	    CreateIoCompletionPort(INVALID_HANDLE_VALUE, nullptr, 0, 0);

	CreateIoCompletionPort(reinterpret_cast<HANDLE>(listenSocket),
	                       iocpHandle,
	                       reinterpret_cast<ULONG_PTR>(&listenSocket),
	                       0);

	while (true) {
	    // 接收连接请求
	    sockaddr_in6 clientAddr{};
	    int clientAddrLen = sizeof(clientAddr);

	    SOCKET clientSocket =
	        accept(listenSocket,
	               reinterpret_cast<SOCKADDR*>(&clientAddr),
	               &clientAddrLen);

	    if (clientSocket == INVALID_SOCKET) {
	        std::cerr << "Failed to accept connection request" << std::endl;
	        continue;
	    }

	    // 将新连接的 SOCKET 加入 IOCP 中等待数据接收
	    char* recvBuffer = new char[4096];

	    WSABUF recvWsaBuf;
	    recvWsaBuf.buf = recvBuffer;
	    recvWsaBuf.len = 4096;

		OVERLAPPED* pOverlapped = new OVERLAPPED;
		memset(pOverlapped, 0, sizeof(OVERLAPPED));
		pOverlapped->hEvent = reinterpret_cast<HANDLE>(clientSocket);
		pOverlapped->Pointer = recvBuffer;

		if (CreateIoCompletionPort(reinterpret_cast<HANDLE>(clientSocket),
		                           iocpHandle,
		                           reinterpret_cast<ULONG_PTR>(&clientSocket),
		                           0) == nullptr ||
		    WSARecv(clientSocket,
		            &recvWsaBuf,
		            1,
		            nullptr,
		            reinterpret_cast<DWORD*>(pOverlapped),
		            nullptr) != 0) {
		    std::cerr << "Failed to receive data from client" << std::endl;

		    closesocket(clientSocket);
		    delete[] recvBuffer;
		    delete pOverlapped;
		    continue;
		}

        // 输出客户端连接信息
        char addrStr[INET6_ADDRSTRLEN];
        inet_ntop(AF_INET6, &(clientAddr.sin6_addr), addrStr, INET6_ADDRSTRLEN);
        
        std::cout << "Client connected: " << addrStr << std::endl;

        // 在 IOCP 上等待数据接收完成，并调用回调函数处理数据
		DWORD numberOfBytesTransferred = 0;
		ULONG_PTR completionKey = 0;
		OVERLAPPED* pCompletedOverlapped = nullptr;

		GetQueuedCompletionStatus(iocpHandle,
		                          &numberOfBytesTransferred,
		                          &completionKey,
		                          &pCompletedOverlapped,
		                          INFINITE);

		ServerIocpCallback(0, numberOfBytesTransferred, pCompletedOverlapped);
	}

    // 关闭 Winsock2 库
    WSACleanup();

	return 0;
}

上述示例代码中，使用了 ParseFrameHeader 函数和 ParseFramePayload 函数来分别解析 WebSocket 帧头部信息和负载数据。

在主函数中，首先创建一个监听 socket，并将其加入到 IOCP 中等待连接请求。然后进入循环，在每次循环中接收连接请求并将新连接的 socket 加入到 IOCP 中等待数据接收。每当有数据接收完成时，就会触发回调函数 ServerIocpCallback 处理数据。

在 ServerIocpCallback 回调函数中，首先判断是否发生错误或者客户端主动关闭连接。如果没有，则解析帧头部信息和负载数据，并进行相应处理。最后重新投递接收请求以继续等待下一次数据的到达。

需要注意的是，在实际应用中，WebSocket 协议还涉及握手过程、心跳包等问题，上述示例代码并未完整地实现所有功能。