SimpleRemoter/server/2015Remote/ConnectionsDlg.cpp

// ConnectionsDlg.cpp : 实现文件
//

#include "stdafx.h"
#include "2015Remote.h"
#include "ConnectionsDlg.h"
#include "afxdialogex.h"
#include <ws2tcpip.h>
#include <utility>


// CConnectionsDlg 对话框

IMPLEMENT_DYNAMIC(CConnectionsDlg, CDialog)

// 连接记录排序辅助：保存完整 ConnectionInfo（供右键强断）与各列排序键。
// 数字列零填充、地址列转定宽十六进制，保证字符串比较即得正确数值顺序。
typedef struct ConnItemData {
    ConnectionInfo info;
    CString key[8];
    CString GetKey(int index) const { return key[index]; }
} ConnItemData;

CConnectionsDlg::CConnectionsDlg(CWnd* pParent, Server* IOCPServer, CONTEXT_OBJECT *ContextObject)
    : DialogBase(CConnectionsDlg::IDD, pParent, IOCPServer, ContextObject, IDI_CONNECTION)
    , m_nRClickItem(-1)
    , m_nSortColumn(-1)
    , m_bSortAscending(true)
{
}

CConnectionsDlg::~CConnectionsDlg()
{
}

void CConnectionsDlg::DoDataExchange(CDataExchange* pDX)
{
    __super::DoDataExchange(pDX);
    DDX_Control(pDX, IDC_LIST, m_ControlList);
    DDX_Control(pDX, IDC_STATIC_COUNT, m_Count);
}

BEGIN_MESSAGE_MAP(CConnectionsDlg, CDialog)
    ON_WM_CLOSE()
    ON_WM_SIZE()
    ON_COMMAND(ID_CONNECTION_KILL, &CConnectionsDlg::OnConnectionKill)
    ON_COMMAND(ID_CONNECTION_REFRESH, &CConnectionsDlg::OnConnectionRefresh)
    ON_NOTIFY(NM_RCLICK, IDC_LIST, &CConnectionsDlg::OnNMRClickList)
    ON_NOTIFY(HDN_ITEMCLICK, 0, &CConnectionsDlg::OnHdnItemclickList)
END_MESSAGE_MAP()


// CConnectionsDlg 消息处理程序

static CString FormatAddr(BYTE af, const uint8_t* addr)
{
    char buf[INET6_ADDRSTRLEN] = { 0 };
    if (af == AF_INET) {
        inet_ntop(AF_INET, addr, buf, INET_ADDRSTRLEN);
    } else {
        inet_ntop(AF_INET6, addr, buf, INET6_ADDRSTRLEN);
    }
    return CA2T(buf);
}

static CString TcpStateName(DWORD state)
{
    switch (state) {
    case 1:  return _T("CLOSED");
    case 2:  return _T("LISTEN");
    case 3:  return _T("SYN_SENT");
    case 4:  return _T("SYN_RCVD");
    case 5:  return _T("ESTABLISHED");
    case 6:  return _T("FIN_WAIT1");
    case 7:  return _T("FIN_WAIT2");
    case 8:  return _T("CLOSE_WAIT");
    case 9:  return _T("CLOSING");
    case 10: return _T("LAST_ACK");
    case 11: return _T("TIME_WAIT");
    case 12: return _T("DELETE_TCB");
    default: return _T("-");
    }
}

// 地址排序键：IPv4/IPv6 统一转成 16 字节定宽十六进制（前缀 4/6 区分协议族），
// 使字符串比较等价于按地址数值排序。
static CString AddrSortKey(uint8_t af, const uint8_t* addr)
{
    CString key;
    key.Format(_T("%d:"), (af == AF_INET) ? 4 : 6);
    for (int i = 0; i < 16; ++i)
        key.AppendFormat(_T("%02X"), addr[i]);
    return key;
}

BOOL CConnectionsDlg::OnInitDialog()
{
    __super::OnInitDialog();

    SetIcon(m_hIcon, TRUE);
    SetIcon(m_hIcon, FALSE);
    CString strString;
    strString.FormatL("%s - 连接管理", m_IPAddress);
    SetWindowText(strString);

    m_ControlList.SetExtendedStyle(LVS_EX_FULLROWSELECT);
    m_ControlList.InsertColumnL(0, "协议", LVCFMT_LEFT, 50);
    m_ControlList.InsertColumnL(1, "本地地址", LVCFMT_LEFT, 120);
    m_ControlList.InsertColumnL(2, "本地端口", LVCFMT_LEFT, 70);
    m_ControlList.InsertColumnL(3, "远端地址", LVCFMT_LEFT, 120);
    m_ControlList.InsertColumnL(4, "远端端口", LVCFMT_LEFT, 70);
    m_ControlList.InsertColumnL(5, "状态", LVCFMT_LEFT, 100);
    m_ControlList.InsertColumnL(6, "PID", LVCFMT_LEFT, 60);
    m_ControlList.InsertColumnL(7, "进程名", LVCFMT_LEFT, 180);

    ShowConnectionList();
    return TRUE;
}

int CConnectionsDlg::ShowConnectionList(void)
{
    Buffer tmp = m_ContextObject->InDeCompressedBuffer.GetMyBuffer(0);
    int len = tmp.GetBufferLength();
    DeleteAllItems();

    if (len < 3) {
        CString strEmpty;
        strEmpty.FormatL("连接数:%d", 0);
        m_Count.SetWindowText(strEmpty);
        return 0;
    }

    const BYTE* p = tmp.GetBuffer(0);
    uint16_t count = 0;
    memcpy(&count, p + 1, 2);
    int avail = (len - 3) / (int)sizeof(ConnectionInfo);
    if ((int)count > avail) count = (uint16_t)avail;

    for (int i = 0; i < (int)count; ++i) {
        const ConnectionInfo* info = (const ConnectionInfo*)(p + 3 + i * sizeof(ConnectionInfo));

        CString strProto = (info->proto == 1) ? _T("TCP") : _T("UDP");
        CString strLocalAddr = FormatAddr(info->af, info->localAddr);
        CString strRemoteAddr = FormatAddr(info->af, info->remoteAddr);
        CString strState = (info->proto == 1) ? TcpStateName(info->state) : _T("-");

        CString strLocalPort, strRemotePort, strPid, strProcName;
        strLocalPort.Format(_T("%u"), info->localPort);
        strRemotePort.Format(_T("%u"), info->remotePort);
        strPid.Format(_T("%u"), info->pid);
        strProcName = CA2T(info->procName, CP_UTF8);

        auto* data = new ConnItemData;
        data->info = *info;
        data->key[0] = strProto;
        data->key[1] = AddrSortKey(info->af, info->localAddr);
        data->key[2].Format(_T("%05u"), info->localPort);
        data->key[3] = AddrSortKey(info->af, info->remoteAddr);
        data->key[4].Format(_T("%05u"), info->remotePort);
        data->key[5].Format(_T("%03u"), info->state);
        data->key[6].Format(_T("%010u"), info->pid);
        data->key[7] = strProcName;

        m_ControlList.InsertItem(i, strProto);
        m_ControlList.SetItemText(i, 1, strLocalAddr);
        m_ControlList.SetItemText(i, 2, strLocalPort);
        m_ControlList.SetItemText(i, 3, strRemoteAddr);
        m_ControlList.SetItemText(i, 4, strRemotePort);
        m_ControlList.SetItemText(i, 5, strState);
        m_ControlList.SetItemText(i, 6, strPid);
        m_ControlList.SetItemText(i, 7, strProcName);
        m_ControlList.SetItemData(i, (DWORD_PTR)data);
    }

    CString strCount;
    strCount.FormatL("连接数:%d", (int)count);
    m_Count.SetWindowText(strCount);

    return 0;
}

void CConnectionsDlg::OnClose()
{
    CancelIO();
    if (IsProcessing()) {
        ShowWindow(SW_HIDE);
        return;
    }

    DeleteAllItems();

    DialogBase::OnClose();
}

void CConnectionsDlg::DeleteAllItems()
{
    for (int i = 0; i < m_ControlList.GetItemCount(); i++) {
        auto data = (ConnItemData*)m_ControlList.GetItemData(i);
        if (data) delete data;
    }
    m_ControlList.DeleteAllItems();
}

void CConnectionsDlg::OnNMRClickList(NMHDR *pNMHDR, LRESULT *pResult)
{
    LPNMITEMACTIVATE pAct = reinterpret_cast<LPNMITEMACTIVATE>(pNMHDR);
    m_nRClickItem = pAct->iItem;

    bool canKill = false;
    if (m_nRClickItem >= 0) {
        auto* data = (ConnItemData*)m_ControlList.GetItemData(m_nRClickItem);
        canKill = (data && data->info.proto == 1 && data->info.af == AF_INET);
    }

    CMenu Menu;
    Menu.LoadMenu(IDR_MENU_CONNECTIONS);
    TranslateMenu(&Menu);
    CMenu* SubMenu = Menu.GetSubMenu(0);
    if (SubMenu && !canKill) {
        SubMenu->EnableMenuItem(ID_CONNECTION_KILL, MF_BYCOMMAND | MF_GRAYED);
    }
    CPoint Point;
    GetCursorPos(&Point);
    SubMenu->TrackPopupMenu(TPM_LEFTALIGN | TPM_LEFTBUTTON, Point.x, Point.y, this);
    *pResult = 0;
}

void CConnectionsDlg::OnConnectionKill()
{
    if (m_nRClickItem < 0) return;
    auto* data = (ConnItemData*)m_ControlList.GetItemData(m_nRClickItem);
    if (!data || data->info.proto != 1 || data->info.af != AF_INET) return;

    LPBYTE szBuffer = (LPBYTE)LocalAlloc(LPTR, 1 + sizeof(ConnectionInfo));
    if (szBuffer == NULL) return;
    szBuffer[0] = COMMAND_KILL_CONNECTION;
    memcpy(szBuffer + 1, &data->info, sizeof(ConnectionInfo));
    m_ContextObject->Send2Client(szBuffer, LocalSize(szBuffer));
    LocalFree(szBuffer);
}

void CConnectionsDlg::OnConnectionRefresh()
{
    BYTE bToken = COMMAND_CONNECTION;
    m_ContextObject->Send2Client(&bToken, 1);
}

void CConnectionsDlg::OnReceiveComplete(void)
{
    switch (m_ContextObject->InDeCompressedBuffer.GetBYTE(0)) {
    case TOKEN_CONNECTION_LIST:
        ShowConnectionList();
        break;
    default:
        break;
    }
}

void CConnectionsDlg::OnSize(UINT nType, int cx, int cy)
{
    __super::OnSize(nType, cx, cy);

    if (!m_ControlList.GetSafeHwnd()) return;

    CRect rc;
    m_ControlList.GetWindowRect(&rc);
    ScreenToClient(&rc);

    m_ControlList.MoveWindow(0, 0, cx, cy, TRUE);
}

int CALLBACK CConnectionsDlg::CompareFunction(LPARAM lParam1, LPARAM lParam2, LPARAM lParamSort)
{
    auto* pSortInfo = reinterpret_cast<std::pair<int, bool>*>(lParamSort);
    int nColumn = pSortInfo->first;
    bool bAscending = pSortInfo->second;

    auto* context1 = (ConnItemData*)lParam1;
    auto* context2 = (ConnItemData*)lParam2;
    CString s1 = context1->GetKey(nColumn);
    CString s2 = context2->GetKey(nColumn);

    int result = s1.Compare(s2);
    return bAscending ? result : -result;
}

void CConnectionsDlg::SortByColumn(int nColumn)
{
    if (nColumn == m_nSortColumn) {
        m_bSortAscending = !m_bSortAscending;
    } else {
        m_nSortColumn = nColumn;
        m_bSortAscending = true;
    }

    std::pair<int, bool> sortInfo(m_nSortColumn, m_bSortAscending);
    m_ControlList.SortItems(CompareFunction, reinterpret_cast<LPARAM>(&sortInfo));
}

void CConnectionsDlg::OnHdnItemclickList(NMHDR* pNMHDR, LRESULT* pResult)
{
    LPNMHEADER pNMHeader = reinterpret_cast<LPNMHEADER>(pNMHDR);
    SortByColumn(pNMHeader->iItem);
    *pResult = 0;
}