// 
// EnumProt.c - Display WSAPROTOCOL_INFO information
//			    for all available protocols

// This define causes windows.h
// to include winsock2.h instead
// of winsock.h
#define _WIN32_WINNT 0x0400

#include <windows.h>
#include <winsock2.h>
#include <commctrl.h>
#include <stdio.h>
#include <stdarg.h>
#include <string.h>

#include "resource.h"
#include "EnumProt.h"

////////////////////////////////////////////////////////////

// Application name
char gszAppName[]  = {"Enumerate Protocols"};

// Window class name
char gszAppClass[] = {"EnumProtClass"};

// Instance handle
HINSTANCE ghInst;

////////////////////////////////////////////////////////////

//
// WinMain()
//
int APIENTRY WinMain(
					  HINSTANCE hInstance,
					  HINSTANCE hPrevInstance,
					  LPSTR lpCmdLine,
					  int nCmdShow
	                 )
{

	// WINSOCK_VERSION is currently defined
	// as MAKEWORD(2,2) in winsock2.h
	WORD wVersionRequested = WINSOCK_VERSION;
	WSADATA wsaData;
	MSG msg;         
	int nRet;

	//
	// Look for WinSock2              
	//
	nRet = WSAStartup(wVersionRequested, &wsaData);
	if (nRet || (wsaData.wVersion != wVersionRequested))
	{
		MessageBox(NULL, "WinSock2 not available", 
						 gszAppName,
						 MB_OK|MB_ICONERROR);
		WSACleanup();
		return 1;
	}

	//
	// Load common control DLL
	//
	InitCommonControls();

	//
	// Register Window class
	//
	if (!InitApplication(hInstance))
		return (FALSE);     

	//
	// Create the ListView window
	//
	if (!InitInstance(hInstance, nCmdShow))
		return (FALSE);

	//
	// Dispatch messages
	//
	while (GetMessage(&msg, NULL, 0, 0))
	{
		TranslateMessage(&msg);
		DispatchMessage(&msg); 
	}

	//
	// Release WinSock
	//
	WSACleanup();
	return (msg.wParam);  
}

////////////////////////////////////////////////////////////

//
// InitApplication()
// Register the window class
//
BOOL InitApplication(HANDLE hInstance)
{
	WNDCLASS  wc;
	
	wc.style = 0;                     
	wc.lpfnWndProc = (WNDPROC)MainWndProc; 
	wc.cbClsExtra = 0;              
	wc.cbWndExtra = 0;              
	wc.hInstance = hInstance;       
	wc.hIcon = LoadIcon(hInstance, 
						MAKEINTRESOURCE(IDI_ICON));
	wc.hCursor = LoadCursor(NULL, IDC_ARROW);
	wc.hbrBackground = GetStockObject(WHITE_BRUSH); 
	wc.lpszMenuName =  MAKEINTRESOURCE(IDR_MENU);
	wc.lpszClassName = gszAppClass;
	return (RegisterClass(&wc));
}

////////////////////////////////////////////////////////////

//
// InitInstance()
// Create main window
//
BOOL InitInstance(HANDLE hInstance, int nCmdShow) 
{
	HWND hWnd;

	ghInst = hInstance;

	hWnd = CreateWindow(
						gszAppClass,           
						gszAppName,
						WS_OVERLAPPEDWINDOW,
						CW_USEDEFAULT, 
						CW_USEDEFAULT, 
						CW_USEDEFAULT, 
						CW_USEDEFAULT,
						NULL,               
						NULL,               
						hInstance,          
						NULL);

	if (!hWnd)
		return (FALSE);

	ShowWindow(hWnd, nCmdShow);
	UpdateWindow(hWnd); 
	return (TRUE);      

}

////////////////////////////////////////////////////////////

//
// MainWndProc()
// Process Windows messages
//
LONG APIENTRY MainWndProc(HWND hWnd, 
						  UINT message, 
						  UINT wParam,
						  LONG lParam
						  )
{
	// Handle to the tree control
	static HWND hWndTree;		

	switch (message) 
	{
		case WM_CREATE:
			//
			// Create the tree view control
			//
			hWndTree = CreateTreeControl(hWnd);
			if (hWndTree == NULL)
				MessageBox(hWnd, 
						   "Error creating TreeView",
						   gszAppName,
						   MB_OK|MB_ICONERROR);
			//
			// Add all protocols to TreeView
			//
			AddAllProtocols(hWndTree);
			break;          

        case WM_SIZE:
			//
			// Resize the TreeView along with the window
			//
            MoveWindow(hWndTree, 
					   0, 
					   0, 
					   LOWORD(lParam), 
					   HIWORD(lParam), 
					   TRUE);
            break;

		case WM_COMMAND:
			switch( LOWORD( wParam ))
			{
				case IDM_EXIT:
					PostQuitMessage(0);
					break;
				case IDM_ABOUT:
					DialogBox(ghInst, 
							  MAKEINTRESOURCE(IDD_ABOUT), 
							  hWnd, 
							  (DLGPROC)About);
					break;
				default:
					return (DefWindowProc(hWnd, 
										  message, 
										  wParam, 
										  lParam));
			}
			break;

		case WM_DESTROY:
			PostQuitMessage(0);
			break;

		default:
			return (DefWindowProc(hWnd, message, wParam, lParam));
	}
	return (0);
}

////////////////////////////////////////////////////////////

//
// CreateTreeControl()
// Setup TreeView common control
//
HWND CreateTreeControl(HWND hWndParent)
{
	HWND hwndTree;
	RECT rect;

	// Get the size and position of the parent window.
	GetClientRect(hWndParent, &rect);

	// Create the tree view control
	hwndTree = CreateWindowEx(0L,
							  WC_TREEVIEW,
							  "",
							  WS_VISIBLE | WS_CHILD |	 
							  WS_BORDER | TVS_HASLINES | 
							  TVS_HASBUTTONS |			 
							  TVS_LINESATROOT,
							  0, 
							  0,
							  rect.right - rect.left, 
							  rect.bottom - rect.top - 15,
							  hWndParent,
							  (HMENU)ID_TREEVIEW,
							  ghInst,
							  NULL);
	return (hwndTree);
}

////////////////////////////////////////////////////////////

//
// AddAllProtocols()
// Call WSAEnumProtocols()
// and add info to TreeControl
//
void AddAllProtocols(HWND hWndTree)
{
	LPBYTE pBuf;
	DWORD dwLen;
	int nRet;
	int nCount;
	LPWSAPROTOCOL_INFO pInfo;
	HTREEITEM hParent;
	HTREEITEM hParent2;
	HTREEITEM hParent3;

	//
	// Determine needed buffer size by
	// intentionally generating an error.
	//
	dwLen = 0;
	nRet = WSAEnumProtocols(NULL,
						   NULL,
						   &dwLen);
	if (nRet == SOCKET_ERROR)
	{
		// Look for the expected error
		if (WSAGetLastError() != WSAENOBUFS)
		{
			ShowWinsockError(WSAGetLastError());
			return;
		}
	}

	//
	// dwLen should now contain the needed buffer size
	// Check to see that it's at least the
	// size of 1 WSAPROTOCOL_INFO structure
	//
	if (dwLen < sizeof(WSAPROTOCOL_INFO))
	{
		MessageBox(NULL, "Internal error",
						 gszAppName,
						 MB_OK|MB_ICONERROR);
		return;
	}

	// Add 1 byte just to be paranoid
	dwLen++;
	pBuf = malloc(dwLen);
	if (pBuf == NULL)
	{
		MessageBox(NULL, 
				   "Couldn't allocate protocol buffer",
				   gszAppName,
				   MB_OK|MB_ICONERROR);
		return;
	}

	//
	// Make the "real" call
	//
	nRet = WSAEnumProtocols(NULL, 
						   (LPWSAPROTOCOL_INFO)pBuf, 
						   &dwLen);
	if (nRet == SOCKET_ERROR)
	{
		free(pBuf);
		ShowWinsockError(WSAGetLastError());
		return;
	}

	//
	// Loop through the protocols
	// nRet contains the number of 
	// protocols returned
	//
	pInfo = (LPWSAPROTOCOL_INFO)pBuf;	
	for(nCount = 0; nCount < nRet; nCount++)
	{
		// Each protocol begins at the root
		// of the tree view control
		hParent = AddTreeItem(hWndTree, 
							  TVI_ROOT, 
							  pInfo->szProtocol);

		// Service flags are added one step down
		hParent2 = AddTreeItem(hWndTree, 
							   hParent, 
							   "Service Flags");

		//
		// Helper macro for adding service flags
		//
		#define ADDSF(f, s1, s2)				\
			AddTreeItem(hWndTree,				\
				hParent2,						\
				(pInfo->dwServiceFlags1 & f) ?	\
				s1 : s2)

		ADDSF(XP1_CONNECTIONLESS, 
			"Connectionless",
			"Connection-oriented");