1. 程式人生 > >原生socket,實現簡單HTPP請求

原生socket,實現簡單HTPP請求

#include <windows.h>
#include <fstream>
#include <WinSock2.h>

#pragma comment(lib, "ws2_32.lib")

int initWin32Net()
{
	WSADATA wsaData;
	int res = WSAStartup(MAKEWORD(2, 2), &wsaData);

	if (res != 0)
		/* Tell the user that we couldn't find a useable */
		/* winsock.dll.     */
		return -1;

	return res;
}

struct RunOnceInitWin32Net
{
	RunOnceInitWin32Net()
	{
		initWin32Net();
	}

	~RunOnceInitWin32Net()
	{
		WSACleanup();
	}
};

RunOnceInitWin32Net win32NetInitor;

bool makeSockNoblock(SOCKET sockfd)
{
	u_long iMode = 1;
	int iResult = ::ioctlsocket(sockfd, FIONBIO, &iMode);
	if (iResult != NO_ERROR)
		return false;
	return true;
}

int getSockRcvBufLen(SOCKET sockfd)
{
	int iOptVal = 0;
	int iOptLen = sizeof(iOptLen);
	if (getsockopt(sockfd, SOL_SOCKET, SO_RCVBUF, (char*)&iOptVal, &iOptLen) != SOCKET_ERROR)
		return iOptVal;

	return 0;
}

int getSockSndBufLen(SOCKET sockfd)
{
	int iOptVal = 0;
	int iOptLen = sizeof(iOptLen);
	if (getsockopt(sockfd, SOL_SOCKET, SO_SNDBUF, (char*)&iOptVal, &iOptLen) != SOCKET_ERROR)
		return iOptVal;

	return 0;
}

bool selectSendData(SOCKET sockfd, const std::string& header, uint64_t timeout_ms)
{
	int iResult = 0;
	fd_set writefds;
	struct timeval timeout;
	timeout.tv_sec = (long)(timeout_ms / 1000);
	timeout.tv_usec = (long)((timeout_ms % 1000) * 1000);

	const char* szSendData = header.data();
	int iSendLen = header.size();
	while (true)
	{
		FD_ZERO(&writefds);
		FD_SET(sockfd, &writefds);
		int hr = ::select(1, NULL, &writefds, NULL, &timeout);
		if (hr == 0)
			continue;
		else if (hr == -1)
			return false;

		if (FD_ISSET(sockfd, &writefds))
		{
			iResult = ::send(sockfd, szSendData, iSendLen, 0);
			if (iResult >= iSendLen)
				break; // success

			szSendData += iResult;
			iSendLen -= iResult;
		}
	}

	return true;
}

bool selectRecvHttpContext(SOCKET sockfd, std::string& context, uint64_t timeout_ms)
{
	int iReadyBufferLen = getSockRcvBufLen(sockfd);
	if (!iReadyBufferLen)
		iReadyBufferLen = 1024 * 6;

	struct timeval timeout;
	timeout.tv_sec = (long)(timeout_ms / 1000);
	timeout.tv_usec = (long)((timeout_ms % 1000) * 1000);
	fd_set readfds;
	int iResult = 0;

	char* szReadyBuffer = (char*)malloc(iReadyBufferLen);
	int iRecvLen = 0;
	while (true)
	{
		FD_ZERO(&readfds);
		FD_SET(sockfd, &readfds);
		int hr = ::select(1, &readfds, NULL, NULL, &timeout);
		if (hr == 0)
		{
			continue;
		}
		else if (hr == -1)
		{
			free(szReadyBuffer);
			return false;
		}

		if (FD_ISSET(sockfd, &readfds))
		{
			iResult = ::recv(sockfd, szReadyBuffer, iReadyBufferLen, 0);
			if (iResult == -1)
			{
				free(szReadyBuffer);
				return false;
			}
			else if (iResult == 0)
			{
				break; // success
			}

			iRecvLen += iResult;
			context.append(szReadyBuffer, iResult);
			memset(szReadyBuffer, 0, iReadyBufferLen);

			// 判斷是否接收完資料
			size_t nHeaderEndPos = context.find("\r\n\r\n");
			if (nHeaderEndPos != std::string::npos)
			{
				int nContextStartPos = nHeaderEndPos + strlen("\r\n\r\n");
				size_t pos = context.find("Content-Length: ");
				if (pos != std::string::npos)
				{
					int nContextLen = atoi(context.substr(pos + strlen("Content-Length: ")).c_str());
					if (iRecvLen >= nContextLen + nContextStartPos)
					{
						context = context.substr(nContextStartPos, nContextLen);
						break; // success
					}
				}
			}
		}
	}

	free(szReadyBuffer);
	return true;
}

bool httpGet(const std::string& url, const std::string& ctx, std::string* rsp, unsigned int port = 80)
{
	SOCKET sockfd = ::socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (sockfd == INVALID_SOCKET)
	{
		std::cout << WSAGetLastError() << std::endl;
		return false;
	}

	if (!makeSockNoblock(sockfd))
	{
		std::cout << WSAGetLastError() << std::endl;
		return false;
	}

	struct hostent* hostaddr = gethostbyname(url.c_str());
	if (!hostaddr)
	{
		std::cout << WSAGetLastError() << std::endl;
		return false;
	}

	struct sockaddr_in addr;
	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);
	addr.sin_addr.s_addr = *(ULONG*)hostaddr->h_addr; //*(ULONG*)hostaddr->h_addr_list[0];//inet_addr(hostaddr->h_name);
	int iResult = ::connect(sockfd, (const struct sockaddr*)&addr, sizeof(sockaddr));
	if (iResult == SOCKET_ERROR)
	{
		if (WSAGetLastError() != WSAEWOULDBLOCK)
		{
			std::cout << WSAGetLastError() << std::endl;
			::closesocket(sockfd);
			return false;
		}
	}

	std::string reqHeader;
	reqHeader = "GET " + ctx + " HTTP/1.1" + "\r\n";
	reqHeader += "HOST: " + url + "\r\n";
	reqHeader += "\r\n";

	if (!selectSendData(sockfd, reqHeader, 500))
	{
		std::cout << WSAGetLastError() << std::endl;
		::closesocket(sockfd);
		return false;
	}

	if (!rsp)
		return true;

	if (!selectRecvHttpContext(sockfd, *rsp, 500))
	{
		std::cout << WSAGetLastError() << std::endl;
		::closesocket(sockfd);
		return false;
	}

	::closesocket(sockfd);
	return true;
}

bool testHttpGetQunHeadPic()
{
	std::string rsp;
	if (!httpGet("p.qlogo.cn", "/gh/436683351/436683351/100/", &rsp))
		return false;

	std::ofstream ofs;
	ofs.open("qun.jpg", std::ios::binary | std::ios_base::out);
	if (ofs.fail())
		return false;
	ofs.write(rsp.data(), rsp.size());
	ofs.close();
	return true;
}

int main(int argc, char **argv)
{
	testHttpGetQunHeadPic();
        return 0;
}