SSPI идентификация

// Client.cpp : Defines the entry point for the console application. 
// 
#define SECURITY_WIN32 
 
#include "stdafx.h" 
#include <stdexcept> 
#include <windows.h> 
#include <winsock.h> 
#include <string> 
#include <iostream> 
#include <iomanip> 
#include <Sspi.h> 
#include <Security.h> 
 
 
#define PORT 9999 
#define BUFFER_SIZE 4096 
 
typedef void* DATA; 
 
//TCHAR errorBuf[256]; 
 
#ifndef UNICODE   
  #define String std::string 
  #define tcout std::cout 
  #define toString std::to_string 
#else 
  #define String std::wstring 
  #define tcout std::wcout 
  #define toString std::to_wstring 
#endif 
 
String errorStr(DWORD code) { 
	LPTSTR errorBuf = NULL; 
	int len = FormatMessage( 
		FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS | FORMAT_MESSAGE_ALLOCATE_BUFFER,  
		NULL, 
		code, 0, (LPTSTR)&errorBuf, 0, NULL); 
	String res = errorBuf; 
	LocalFree((HLOCAL)errorBuf); 
	return res; 
} 
 
void throwSystemError(int error) { 
	tcout << errorStr(error) << std::flush; 
	std::exit(1); 
} 
 
void throwSocketError() { 
	throwSystemError(WSAGetLastError()); 
} 
 
void win32Check(BOOL res) { 
	if (!res) 
		throwSystemError(GetLastError()); 
} 
 
int sockCheck(int res) { 
	if (res == SOCKET_ERROR) 
		throwSocketError(); 
	return res; 
} 
 
void secCheck(SECURITY_STATUS status) { 
	if (status != SEC_E_OK) 
		throwSystemError(status); 
} 
 
class CustomAuth { 
private: 
	CredHandle credHandle; 
	PCtxtHandle contextPtr; 
	CtxtHandle context; 
 
public: 
	CustomAuth() { 
		secCheck(AcquireCredentialsHandle(NULL, NEGOSSP_NAME, SECPKG_CRED_BOTH, 
			NULL, NULL, NULL, NULL, &credHandle, NULL)); 
		memset(&context, 0, sizeof(context)); 
		contextPtr = NULL;  
	} 
 
	~CustomAuth() { 
		reset(); 
		FreeCredentialsHandle(&credHandle); 
	} 
 
	void reset() { 
		if (contextPtr)	{ 
			secCheck(DeleteSecurityContext(contextPtr)); 
			contextPtr = NULL; 
		} 
	} 
 
	SECURITY_STATUS auth(void * data, int &len) { 
		PSecBufferDesc inBufPtr = NULL; 
		SecBufferDesc inBufDesc; 
		SecBuffer inBuf; 
		if (len) { 
			inBufDesc.ulVersion = SECBUFFER_VERSION; 
			inBufDesc.cBuffers = 1; 
			inBufDesc.pBuffers = &inBuf; 
 
			inBuf.BufferType = SECBUFFER_TOKEN; 
			inBuf.cbBuffer = len; 
			inBuf.pvBuffer = data; 
 
			inBufPtr = &inBufDesc; 
		} 
 
		SecBufferDesc outBufDesc; 
		SecBuffer outBuf; 
		outBufDesc.ulVersion = SECBUFFER_VERSION; 
		outBufDesc.cBuffers = 1; 
		outBufDesc.pBuffers = &outBuf; 
		memset(&outBuf, 0, sizeof(outBuf)); 
		outBuf.BufferType = SECBUFFER_TOKEN; 
 
		SECURITY_STATUS res = initContext(&context, inBufPtr, &outBufDesc); 
		contextPtr = &context; 
		switch (res) { 
			case SEC_E_OK: 
			case SEC_I_CONTINUE_NEEDED: 
				len = outBuf.cbBuffer; 
				break; 
			case SEC_E_INCOMPLETE_MESSAGE: 
				len = 0; 
				break; 
			case SEC_I_COMPLETE_NEEDED: 
				len = 0; 
				res = SEC_E_OK; 
			case SEC_I_COMPLETE_AND_CONTINUE: 
				secCheck(CompleteAuthToken(contextPtr, &outBufDesc)); 
				len = outBuf.cbBuffer; 
				break; 
			default: 
				throwSystemError(res); 
		} 
		memcpy(data, outBuf.pvBuffer, len); 
		secCheck(FreeContextBuffer(outBuf.pvBuffer)); 
		return res; 
	} 
 
	HANDLE getUserToken() { 
		SecPkgContext_AccessToken token; 
		secCheck(QueryContextAttributes(contextPtr, SECPKG_ATTR_ACCESS_TOKEN, &token)); 
		return token.AccessToken; 
	} 
 
	static String statusName(SECURITY_STATUS status) { 
		switch (status)	{ 
			case SEC_E_OK: return _T("SEC_E_OK"); 
			case SEC_I_COMPLETE_AND_CONTINUE: return _T("SEC_I_COMPLETE_AND_CONTINUE"); 
			case SEC_I_COMPLETE_NEEDED: return _T("SEC_I_COMPLETE_NEEDED"); 
			case SEC_I_CONTINUE_NEEDED: return _T("SEC_I_CONTINUE_NEEDED"); 
			case SEC_I_INCOMPLETE_CREDENTIALS: return _T("SEC_I_INCOMPLETE_CREDENTIALS"); 
			case SEC_E_INCOMPLETE_MESSAGE: return _T("SEC_E_INCOMPLETE_MESSAGE"); 
		default: 
			return toString(status); 
		} 
	} 
 
protected: 
	virtual SECURITY_STATUS initContext(PCtxtHandle context, PSecBufferDesc inBuf, PSecBufferDesc outBuf) = 0; 
 
	PCredHandle getCredHandle() { 
		return &credHandle; 
	} 
 
	PCtxtHandle getContext() { 
		return contextPtr; 
	} 
}; 
 
class ServerAuth: public CustomAuth { 
protected: 
	SECURITY_STATUS initContext(PCtxtHandle context, PSecBufferDesc inBuf, PSecBufferDesc outBuf) { 
		DWORD attr; 
		return AcceptSecurityContext(getCredHandle(), getContext(), inBuf, 
			ISC_REQ_CONFIDENTIALITY | ISC_REQ_CONNECTION | ISC_REQ_ALLOCATE_MEMORY, 
			SECURITY_NETWORK_DREP, context, outBuf, &attr, NULL 
			); 
	} 
}; 
 
class ClientAuth: public CustomAuth { 
protected: 
	SECURITY_STATUS initContext(PCtxtHandle context, PSecBufferDesc inBuf, PSecBufferDesc outBuf) { 
		DWORD attr; 
		return InitializeSecurityContext(getCredHandle(), getContext(), NULL, 
			ISC_REQ_CONFIDENTIALITY | ISC_REQ_CONNECTION | ISC_REQ_ALLOCATE_MEMORY, 
			0, SECURITY_NETWORK_DREP,  
			inBuf, 0, context, outBuf, &attr, NULL 
			); 
	} 
}; 
 
SOCKET init_socket() { 
	WSADATA data; 
	int code = WSAStartup(MAKEWORD(2, 2), &data); 
	if (code) 
		throwSystemError(code); 
	SOCKET res = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP); 
	if (res == INVALID_SOCKET) 
		throwSocketError(); 
	return res; 
} 
 
void connect_socket(SOCKET socket, char * host) { 
	sockaddr_in addr; 
	addr.sin_family = AF_INET; 
	addr.sin_port = htons(PORT); 
	hostent * hostEnt = gethostbyname(host); 
	if (!hostEnt) 
		throwSocketError(); 
	addr.sin_addr.S_un =  ((in_addr *)hostEnt->h_addr_list[0])->S_un; 
	memset(addr.sin_zero, 0, sizeof(addr.sin_zero)); 
	sockCheck(connect(socket, (sockaddr*)&addr, sizeof(addr))); 
} 
 
SOCKET accept_socket(SOCKET socket) { 
	sockaddr_in addr; 
	addr.sin_family = AF_INET; 
	addr.sin_port = htons(PORT); 
	addr.sin_addr.S_un.S_addr =  0; 
	memset(addr.sin_zero, 0, sizeof(addr.sin_zero)); 
	sockCheck(bind(socket, (sockaddr*)&addr, sizeof(addr))); 
	sockCheck(listen(socket, SOMAXCONN)); 
	SOCKET res = accept(socket, NULL, NULL); 
	if (res == INVALID_SOCKET) 
		throwSocketError(); 
	return res; 
} 
 
void sendData(SOCKET socket, char * data, int len) { 
	sockCheck(send(socket, (char*)&len, sizeof(len), 0)); 
	if (len) 
		sockCheck(send(socket, data, len, 0)); 
	tcout << __T("Sended ") << len << _T(" bytes") << std::endl; 
} 
 
void _recvData(SOCKET socket, char * data, int len) { 
	int read = sockCheck(recv(socket, data, len, 0)); 
	if (read != len) 
		throwSystemError(ERROR_NO_MORE_FILES); 
} 
 
void recvData(SOCKET socket, char * data, int &len) { 
	_recvData(socket, (char*)&len, sizeof(len)); 
	if (len) 
		_recvData(socket, data, len); 
	tcout << __T("Received ") << len << _T(" bytes") << std::endl; 
} 
 
String getUserName(HANDLE token) { 
	char data[BUFFER_SIZE]; 
	DWORD len; 
	win32Check(GetTokenInformation(token, TokenUser, data, BUFFER_SIZE, &len)); 
	PSID sid = ((PTOKEN_USER)data)->User.Sid; 
	win32Check(IsValidSid(sid)); 
	TCHAR login[BUFFER_SIZE]; 
	TCHAR domain[BUFFER_SIZE]; 
	len = BUFFER_SIZE; 
	DWORD domLen = BUFFER_SIZE; 
	SID_NAME_USE use; 
	win32Check(LookupAccountSid(NULL, sid, login, &len, domain, &domLen, &use)); 
	String res = String(domain) + String(_T("\\")) + String(login); 
	return res; 
} 
 
String getCurrentUser() { 
	TCHAR buf[BUFFER_SIZE]; 
	DWORD len = BUFFER_SIZE; 
	win32Check(GetUserNameEx(NameSamCompatible, buf, &len)); 
	return buf; 
} 
 
int main(int argc, char * argv[]) 
{ 
	std::locale::global(std::locale("")); 
	bool isServer = argc < 2; 
	SOCKET socket = init_socket(); 
	CustomAuth * auth; 
    tcout << _T("Current user: ") << getCurrentUser() << std::endl; 
	char data[BUFFER_SIZE]; 
	int len = 0; 
	if (isServer) { 
		tcout << _T("Wait clients") << std::endl; 
		SOCKET acc = accept_socket(socket); 
		closesocket(socket); 
		tcout << _T("Client accepted") << std::endl; 
		socket = acc; 
		auth = new ServerAuth(); 
		recvData(socket, data, len); 
	} else { 
		connect_socket(socket, argv[1]); 
		auth = new ClientAuth(); 
	} 
	SECURITY_STATUS status = auth->auth(data, len); 
	while (status != SEC_E_OK || len > 0) { 
		tcout << _T("ClientAuthCode: ") << auth->statusName(status) << std::endl; 
		if (len) { 
			sendData(socket, data, len); 
			len = 0; 
		} 
		if (status != SEC_E_OK) { 
			recvData(socket, data, len); 
			status = auth->auth(data, len); 
		} 
	} 
	tcout << _T("End authentification. Status: ") << auth->statusName(status) << std::endl; 
	if (isServer) 
		tcout << _T("Connected user: ") << getUserName(auth->getUserToken()) << std::endl; 
 
	delete auth; 
	closesocket(socket); 
 
	return 0; 
}