VolumeLocker.cpp

#define UNICODE
#define _UNICODE
#define WIN32_LEAN_AND_MEAN

#include <windows.h>
#include <shellapi.h>
#include <shlobj.h>
#include <mmdeviceapi.h>
#include <endpointvolume.h>
#include <functiondiscoverykeys_devpkey.h>
#include <commctrl.h>
#include <string>
#include <vector>

#pragma comment(lib, "ole32.lib")
#pragma comment(lib, "comctl32.lib")
#pragma comment(lib, "shell32.lib")

#define WM_TRAYICON (WM_USER + 1)
#define IDM_SHOW 1001
#define IDM_EXIT 1002
#define IDD_MAIN_DIALOG 101
#define IDC_DEVICE_COMBO 1003
#define IDC_LOCK_BUTTON 1004
#define IDC_UNLOCK_BUTTON 1005
#define IDC_STARTUP_CHECK 1006
#define IDC_STATUS_TEXT 1007
#define IDC_STARTMIN_CHECK 1008

// GUID identifying volume changes made by this application.
// Used to break the feedback loop: if OnNotify fires because we set the
// volume ourselves, we skip the correction.
// {6A4C7E32-9F8B-4D5A-B231-C05E87F41D92}
static const GUID GUID_OUR_EVENT_CONTEXT =
    { 0x6a4c7e32, 0x9f8b, 0x4d5a, { 0xb2, 0x31, 0xc0, 0x5e, 0x87, 0xf4, 0x1d, 0x92 } };

// Forward declarations
INT_PTR CALLBACK DialogProc(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam);
void EnumerateDevices();
void LockVolume();
void UnlockVolume();
void UpdateStatus(const wchar_t* status);
void AddToSystemTray();
void RemoveFromSystemTray();
void ShowMainDialog();
void HideMainDialog();
void SetStartup(bool enable);
bool IsStartupEnabled();
void SaveConfig();
void LoadConfig();
std::wstring GetConfigPath();
void AutoLockDevice();

// Global variables
HINSTANCE g_hInst = NULL;
HWND g_hMainDlg = NULL;
NOTIFYICONDATA g_nid = {0};
IMMDeviceEnumerator* g_pEnumerator = NULL;
IAudioEndpointVolume* g_pVolume = NULL;
std::wstring g_lockedDeviceId;
bool g_isLocked = false;
std::wstring g_configPath;
bool g_startMinimized = false;

struct DeviceInfo {
    std::wstring id;
    std::wstring name;
};
std::vector<DeviceInfo> g_devices;

// ---------------------------------------------------------------------------
// IAudioEndpointVolumeCallback implementation
//
// Statically allocated; refcount is fixed at 1 since the object lives for
// the lifetime of the process.  COM requires the full IUnknown surface even
// though we never actually do ref-counted lifetime management here.
// ---------------------------------------------------------------------------
class CVolumeCallback : public IAudioEndpointVolumeCallback {
public:
    // IUnknown
    STDMETHODIMP QueryInterface(REFIID riid, void** ppv) {
        if (riid == IID_IUnknown ||
            riid == __uuidof(IAudioEndpointVolumeCallback)) {
            *ppv = static_cast<IAudioEndpointVolumeCallback*>(this);
            return S_OK;
        }
        *ppv = nullptr;
        return E_NOINTERFACE;
    }
    STDMETHODIMP_(ULONG) AddRef()  { return 1; }
    STDMETHODIMP_(ULONG) Release() { return 1; }

    // IAudioEndpointVolumeCallback
    STDMETHODIMP OnNotify(PAUDIO_VOLUME_NOTIFICATION_DATA pNotify) {
        if (!g_isLocked || !g_pVolume)
            return S_OK;

        // Ignore notifications that originated from our own SetMasterVolumeLevelScalar
        // call — otherwise each correction would trigger another notification.
        if (pNotify->guidEventContext == GUID_OUR_EVENT_CONTEXT)
            return S_OK;

        // An external source lowered (or changed) the volume — restore it.
        if (pNotify->fMasterVolume < 0.99f) {
            g_pVolume->SetMasterVolumeLevelScalar(1.0f, &GUID_OUR_EVENT_CONTEXT);
        }

        return S_OK;
    }
};

// Single static instance — lives for the whole process.
static CVolumeCallback g_volumeCallback;

// ---------------------------------------------------------------------------
// Device enumeration
// ---------------------------------------------------------------------------
void EnumerateDevices() {
    g_devices.clear();

    if (!g_pEnumerator) return;

    IMMDeviceCollection* pCollection = NULL;
    HRESULT hr = g_pEnumerator->EnumAudioEndpoints(eCapture, DEVICE_STATE_ACTIVE, &pCollection);

    if (SUCCEEDED(hr)) {
        UINT count;
        pCollection->GetCount(&count);

        for (UINT i = 0; i < count; i++) {
            IMMDevice* pDevice = NULL;
            if (SUCCEEDED(pCollection->Item(i, &pDevice))) {
                LPWSTR pwszID = NULL;
                pDevice->GetId(&pwszID);

                IPropertyStore* pProps = NULL;
                if (SUCCEEDED(pDevice->OpenPropertyStore(STGM_READ, &pProps))) {
                    PROPVARIANT varName;
                    PropVariantInit(&varName);

                    if (SUCCEEDED(pProps->GetValue(PKEY_Device_FriendlyName, &varName))) {
                        DeviceInfo info;
                        info.id   = pwszID;
                        info.name = varName.pwszVal;
                        g_devices.push_back(info);
                    }

                    PropVariantClear(&varName);
                    pProps->Release();
                }

                CoTaskMemFree(pwszID);
                pDevice->Release();
            }
        }
        pCollection->Release();
    }
}

// ---------------------------------------------------------------------------
// Lock / Unlock
// ---------------------------------------------------------------------------
void LockVolume() {
    if (g_isLocked) return;

    HWND hCombo = GetDlgItem(g_hMainDlg, IDC_DEVICE_COMBO);
    int sel = (int)SendMessage(hCombo, CB_GETCURSEL, 0, 0);

    if (sel == CB_ERR || sel >= (int)g_devices.size()) {
        MessageBox(g_hMainDlg, L"Please select a device first.", L"No Device Selected", MB_OK | MB_ICONWARNING);
        return;
    }

    g_lockedDeviceId = g_devices[sel].id;

    IMMDevice* pDevice = NULL;
    HRESULT hr = g_pEnumerator->GetDevice(g_lockedDeviceId.c_str(), &pDevice);

    if (SUCCEEDED(hr)) {
        hr = pDevice->Activate(__uuidof(IAudioEndpointVolume), CLSCTX_ALL, NULL, (void**)&g_pVolume);

        if (SUCCEEDED(hr)) {
            // Snap to 100% immediately, identifying this change as ours.
            g_pVolume->SetMasterVolumeLevelScalar(1.0f, &GUID_OUR_EVENT_CONTEXT);

            // Register for future change notifications.
            hr = g_pVolume->RegisterControlChangeNotify(&g_volumeCallback);
        }

        pDevice->Release();
    }

    if (SUCCEEDED(hr)) {
        g_isLocked = true;

        std::wstring status = L"Locked: " + g_devices[sel].name;
        UpdateStatus(status.c_str());

        EnableWindow(GetDlgItem(g_hMainDlg, IDC_LOCK_BUTTON), FALSE);
        EnableWindow(GetDlgItem(g_hMainDlg, IDC_UNLOCK_BUTTON), TRUE);
        EnableWindow(hCombo, FALSE);

        SaveConfig();
    } else {
        if (g_pVolume) { g_pVolume->Release(); g_pVolume = NULL; }
        MessageBox(g_hMainDlg, L"Failed to lock device volume.", L"Error", MB_OK | MB_ICONERROR);
    }
}

void UnlockVolume() {
    if (!g_isLocked) return;

    if (g_pVolume) {
        g_pVolume->UnregisterControlChangeNotify(&g_volumeCallback);
        g_pVolume->Release();
        g_pVolume = NULL;
    }

    g_isLocked = false;
    g_lockedDeviceId.clear();

    if (g_hMainDlg) {
        UpdateStatus(L"Not locked");
        EnableWindow(GetDlgItem(g_hMainDlg, IDC_LOCK_BUTTON), TRUE);
        EnableWindow(GetDlgItem(g_hMainDlg, IDC_UNLOCK_BUTTON), FALSE);
        EnableWindow(GetDlgItem(g_hMainDlg, IDC_DEVICE_COMBO), TRUE);
    }

    SaveConfig();
}

// ---------------------------------------------------------------------------
// UI helpers
// ---------------------------------------------------------------------------
void UpdateStatus(const wchar_t* status) {
    SetDlgItemText(g_hMainDlg, IDC_STATUS_TEXT, status);
}

void AddToSystemTray() {
    g_nid.cbSize = sizeof(NOTIFYICONDATA);
    g_nid.hWnd   = g_hMainDlg;
    g_nid.uID    = 1;
    g_nid.uFlags = NIF_ICON | NIF_MESSAGE | NIF_TIP;
    g_nid.uCallbackMessage = WM_TRAYICON;
    g_nid.hIcon  = LoadIcon(NULL, IDI_APPLICATION);
    wcscpy_s(g_nid.szTip, L"Volume Locker");
    Shell_NotifyIcon(NIM_ADD, &g_nid);
}

void RemoveFromSystemTray() {
    Shell_NotifyIcon(NIM_DELETE, &g_nid);
}

void ShowMainDialog() {
    ShowWindow(g_hMainDlg, SW_SHOW);
    SetForegroundWindow(g_hMainDlg);
}

void HideMainDialog() {
    ShowWindow(g_hMainDlg, SW_HIDE);
}

// ---------------------------------------------------------------------------
// Startup registry
// ---------------------------------------------------------------------------
void SetStartup(bool enable) {
    HKEY hKey;
    const wchar_t* keyPath = L"Software\\Microsoft\\Windows\\CurrentVersion\\Run";

    if (RegOpenKeyEx(HKEY_CURRENT_USER, keyPath, 0, KEY_SET_VALUE, &hKey) == ERROR_SUCCESS) {
        if (enable) {
            wchar_t exePath[MAX_PATH];
            GetModuleFileName(NULL, exePath, MAX_PATH);
            RegSetValueEx(hKey, L"VolumeLocker", 0, REG_SZ,
                          (BYTE*)exePath, (DWORD)((wcslen(exePath) + 1) * sizeof(wchar_t)));
        } else {
            RegDeleteValue(hKey, L"VolumeLocker");
        }
        RegCloseKey(hKey);
        SaveConfig();
    }
}

bool IsStartupEnabled() {
    HKEY hKey;
    const wchar_t* keyPath = L"Software\\Microsoft\\Windows\\CurrentVersion\\Run";
    bool exists = false;

    if (RegOpenKeyEx(HKEY_CURRENT_USER, keyPath, 0, KEY_QUERY_VALUE, &hKey) == ERROR_SUCCESS) {
        exists = (RegQueryValueEx(hKey, L"VolumeLocker", NULL, NULL, NULL, NULL) == ERROR_SUCCESS);
        RegCloseKey(hKey);
    }
    return exists;
}

// ---------------------------------------------------------------------------
// Config persistence
// ---------------------------------------------------------------------------
std::wstring GetConfigPath() {
    wchar_t appDataPath[MAX_PATH];
    if (SHGetFolderPath(NULL, CSIDL_APPDATA, NULL, 0, appDataPath) == S_OK) {
        std::wstring path = appDataPath;
        path += L"\\VolumeLocker";
        CreateDirectory(path.c_str(), NULL);
        path += L"\\config.ini";
        return path;
    }
    return L"config.ini";
}

void SaveConfig() {
    if (g_configPath.empty()) g_configPath = GetConfigPath();
    WritePrivateProfileString(L"Settings", L"LockedDeviceID",  g_lockedDeviceId.c_str(),  g_configPath.c_str());
    WritePrivateProfileString(L"Settings", L"IsLocked",        g_isLocked ? L"1" : L"0",  g_configPath.c_str());
    WritePrivateProfileString(L"Settings", L"RunAtStartup",    IsStartupEnabled() ? L"1" : L"0", g_configPath.c_str());
    WritePrivateProfileString(L"Settings", L"StartMinimized",  g_startMinimized ? L"1" : L"0",   g_configPath.c_str());
}

void LoadConfig() {
    if (g_configPath.empty()) g_configPath = GetConfigPath();

    wchar_t buffer[512];
    GetPrivateProfileString(L"Settings", L"LockedDeviceID", L"", buffer, 512, g_configPath.c_str());
    g_lockedDeviceId = buffer;

    g_startMinimized = GetPrivateProfileInt(L"Settings", L"StartMinimized", 0, g_configPath.c_str()) == 1;
}

void AutoLockDevice() {
    if (g_lockedDeviceId.empty()) return;

    int deviceIndex = -1;
    for (size_t i = 0; i < g_devices.size(); i++) {
        if (g_devices[i].id == g_lockedDeviceId) {
            deviceIndex = (int)i;
            break;
        }
    }

    if (deviceIndex == -1) {
        g_lockedDeviceId.clear();
        SaveConfig();
        return;
    }

    SendMessage(GetDlgItem(g_hMainDlg, IDC_DEVICE_COMBO), CB_SETCURSEL, deviceIndex, 0);
    LockVolume();
}

// ---------------------------------------------------------------------------
// Dialog procedure
// ---------------------------------------------------------------------------
INT_PTR CALLBACK DialogProc(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam) {
    switch (message) {
        case WM_INITDIALOG: {
            g_hMainDlg = hDlg;

            if (g_startMinimized) {
                SetWindowLong(hDlg, GWL_EXSTYLE,
                    GetWindowLong(hDlg, GWL_EXSTYLE) | WS_EX_NOACTIVATE);
            }

            LoadConfig();

            // Center the dialog
            RECT rc;
            GetWindowRect(hDlg, &rc);
            int x = (GetSystemMetrics(SM_CXSCREEN) - (rc.right - rc.left)) / 2;
            int y = (GetSystemMetrics(SM_CYSCREEN) - (rc.bottom - rc.top)) / 2;
            SetWindowPos(hDlg, NULL, x, y, 0, 0, SWP_NOSIZE | SWP_NOZORDER);

            EnumerateDevices();
            HWND hCombo = GetDlgItem(hDlg, IDC_DEVICE_COMBO);
            for (const auto& dev : g_devices)
                SendMessage(hCombo, CB_ADDSTRING, 0, (LPARAM)dev.name.c_str());

            if (!g_devices.empty())
                SendMessage(hCombo, CB_SETCURSEL, 0, 0);

            CheckDlgButton(hDlg, IDC_STARTUP_CHECK,  IsStartupEnabled()  ? BST_CHECKED : BST_UNCHECKED);
            CheckDlgButton(hDlg, IDC_STARTMIN_CHECK, g_startMinimized    ? BST_CHECKED : BST_UNCHECKED);

            EnableWindow(GetDlgItem(hDlg, IDC_UNLOCK_BUTTON), FALSE);
            UpdateStatus(L"Not locked");

            AddToSystemTray();
            AutoLockDevice();

            return TRUE;
        }

        case WM_COMMAND: {
            switch (LOWORD(wParam)) {
                case IDC_LOCK_BUTTON:
                    LockVolume();
                    return TRUE;

                case IDC_UNLOCK_BUTTON:
                    UnlockVolume();
                    return TRUE;

                case IDC_STARTUP_CHECK:
                    SetStartup(IsDlgButtonChecked(hDlg, IDC_STARTUP_CHECK) == BST_CHECKED);
                    return TRUE;

                case IDC_STARTMIN_CHECK:
                    g_startMinimized = (IsDlgButtonChecked(hDlg, IDC_STARTMIN_CHECK) == BST_CHECKED);
                    SaveConfig();
                    return TRUE;

                case IDCANCEL:
                    HideMainDialog();
                    return TRUE;

                // Tray context menu items.  Previously handled in the default
                // branch — which was dead code because WM_COMMAND is matched
                // by this case, so the default was never reached.
                case IDM_SHOW:
                    ShowMainDialog();
                    return TRUE;

                case IDM_EXIT:
                    UnlockVolume();
                    RemoveFromSystemTray();
                    DestroyWindow(hDlg);
                    return TRUE;
            }
            break;
        }

        case WM_TRAYICON: {
            if (lParam == WM_LBUTTONDBLCLK) {
                ShowMainDialog();
            } else if (lParam == WM_RBUTTONUP) {
                POINT pt;
                GetCursorPos(&pt);

                HMENU hMenu = CreatePopupMenu();
                AppendMenu(hMenu, MF_STRING, IDM_SHOW, L"Show");
                AppendMenu(hMenu, MF_STRING, IDM_EXIT, L"Exit");

                SetForegroundWindow(hDlg);
                TrackPopupMenu(hMenu, TPM_BOTTOMALIGN | TPM_LEFTALIGN, pt.x, pt.y, 0, hDlg, NULL);
                DestroyMenu(hMenu);
            }
            return TRUE;
        }

        case WM_SYSCOMMAND: {
            if (wParam == SC_MINIMIZE) {
                HideMainDialog();
                return TRUE;
            }
            break;
        }

        case WM_CLOSE: {
            HideMainDialog();
            return TRUE;
        }

        case WM_DESTROY: {
            // UnlockVolume / RemoveFromSystemTray are already called by
            // IDM_EXIT before DestroyWindow.  Guard against double-release.
            UnlockVolume();
            RemoveFromSystemTray();
            PostQuitMessage(0);
            return TRUE;
        }
    }

    return FALSE;
}

// ---------------------------------------------------------------------------
// Entry point
// ---------------------------------------------------------------------------
int WINAPI wWinMain(HINSTANCE hInstance, HINSTANCE /*hPrevInstance*/,
                    LPWSTR /*lpCmdLine*/, int /*nCmdShow*/) {
    g_hInst = hInstance;

    // Apartment-threaded COM is required for Core Audio APIs.
    CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);

    // Load config early so we know the start-minimized preference.
    LoadConfig();

    CoCreateInstance(__uuidof(MMDeviceEnumerator), NULL, CLSCTX_ALL,
                     __uuidof(IMMDeviceEnumerator), (void**)&g_pEnumerator);

    HWND hDlg = CreateDialogParam(hInstance, MAKEINTRESOURCE(IDD_MAIN_DIALOG),
                                  NULL, DialogProc, 0);
    if (!hDlg) {
        MessageBox(NULL, L"Failed to create dialog", L"Error", MB_OK | MB_ICONERROR);
        CoUninitialize();
        return 1;
    }

    if (!g_startMinimized) {
        ShowWindow(hDlg, SW_SHOW);
        UpdateWindow(hDlg);
    }

    // Main thread runs the message loop — no polling thread needed.
    // Core Audio delivers IAudioEndpointVolumeCallback::OnNotify on this
    // apartment thread via the message queue.
    MSG msg;
    while (GetMessage(&msg, NULL, 0, 0)) {
        if (!IsDialogMessage(hDlg, &msg)) {
            TranslateMessage(&msg);
            DispatchMessage(&msg);
        }
    }

    if (g_pEnumerator) { g_pEnumerator->Release(); g_pEnumerator = NULL; }
    CoUninitialize();

    return (int)msg.wParam;
}