13 files changed, 662 insertions, 71 deletions
diff --git a/DcsBml/ComponentName.c b/DcsBml/ComponentName.c
new file mode 100644
index 0000000..1971328
--- /dev/null
+++ b/DcsBml/ComponentName.c
@@ -0,0 +1,136 @@
+/** @file
+This is DCS boot menu lock application
+
+Copyright (c) 2016. Disk Cryptography Services for EFI (DCS), Alex Kolotnikov
+
+This program and the accompanying materials
+are licensed and made available under the terms and conditions
+of the GNU Lesser General Public License, version 3.0 (LGPL-3.0).
+
+The full text of the license may be found at
+https://opensource.org/licenses/LGPL-3.0
+**/
+
+#include "DcsBml.h"
+
+///
+/// Component Name Protocol instance
+///
+GLOBAL_REMOVE_IF_UNREFERENCED 
+EFI_COMPONENT_NAME_PROTOCOL  gDcsBmlComponentName = {
+  (EFI_COMPONENT_NAME_GET_DRIVER_NAME)    DcsBmlComponentNameGetDriverName,
+  (EFI_COMPONENT_NAME_GET_CONTROLLER_NAME)DcsBmlComponentNameGetControllerName,
+  "eng"
+};
+
+///
+/// Component Name 2 Protocol instance
+///
+GLOBAL_REMOVE_IF_UNREFERENCED 
+EFI_COMPONENT_NAME2_PROTOCOL  gDcsBmlComponentName2 = {
+  DcsBmlComponentNameGetDriverName,
+  DcsBmlComponentNameGetControllerName,
+  "en"
+};
+
+///
+/// Table of driver names
+///
+GLOBAL_REMOVE_IF_UNREFERENCED 
+EFI_UNICODE_STRING_TABLE mDcsBmlDriverNameTable[] = {
+  { "eng;en", (CHAR16 *)L"DcsBml" },
+  { NULL, NULL }
+};
+
+/**
+  Retrieves a Unicode string that is the user-readable name of the EFI Driver.
+
+  @param  This       A pointer to the EFI_COMPONENT_NAME_PROTOCOL instance.
+  @param  Language   A pointer to a three-character ISO 639-2 language identifier.
+                     This is the language of the driver name that that the caller
+                     is requesting, and it must match one of the languages specified
+                     in SupportedLanguages.  The number of languages supported by a
+                     driver is up to the driver writer.
+  @param  DriverName A pointer to the Unicode string to return.  This Unicode string
+                     is the name of the driver specified by This in the language
+                     specified by Language.
+
+  @retval EFI_SUCCESS           The Unicode string for the Driver specified by This
+                                and the language specified by Language was returned
+                                in DriverName.
+  @retval EFI_INVALID_PARAMETER Language is NULL.
+  @retval EFI_INVALID_PARAMETER DriverName is NULL.
+  @retval EFI_UNSUPPORTED       The driver specified by This does not support the
+                                language specified by Language.
+
+**/
+EFI_STATUS
+EFIAPI
+DcsBmlComponentNameGetDriverName (
+  IN EFI_COMPONENT_NAME2_PROTOCOL  *This,
+  IN  CHAR8                        *Language,
+  OUT CHAR16                       **DriverName
+  )
+{
+  return LookupUnicodeString2 (
+           Language,
+           This->SupportedLanguages,
+           mDcsBmlDriverNameTable,
+           DriverName,
+           (BOOLEAN)(This != &gDcsBmlComponentName2)
+           );
+}
+
+/**
+  Retrieves a Unicode string that is the user readable name of the controller
+  that is being managed by an EFI Driver.
+
+  @param  This             A pointer to the EFI_COMPONENT_NAME_PROTOCOL instance.
+  @param  ControllerHandle The handle of a controller that the driver specified by
+                           This is managing.  This handle specifies the controller
+                           whose name is to be returned.
+  @param  ChildHandle      The handle of the child controller to retrieve the name
+                           of.  This is an optional parameter that may be NULL.  It
+                           will be NULL for device drivers.  It will also be NULL
+                           for a bus drivers that wish to retrieve the name of the
+                           bus controller.  It will not be NULL for a bus driver
+                           that wishes to retrieve the name of a child controller.
+  @param  Language         A pointer to a three character ISO 639-2 language
+                           identifier.  This is the language of the controller name
+                           that the caller is requesting, and it must match one
+                           of the languages specified in SupportedLanguages.  The
+                           number of languages supported by a driver is up to the
+                           driver writer.
+  @param  ControllerName   A pointer to the Unicode string to return.  This Unicode
+                           string is the name of the controller specified by
+                           ControllerHandle and ChildHandle in the language specified
+                           by Language, from the point of view of the driver specified
+                           by This.
+
+  @retval EFI_SUCCESS           The Unicode string for the user-readable name in the
+                                language specified by Language for the driver
+                                specified by This was returned in DriverName.
+  @retval EFI_INVALID_PARAMETER ControllerHandle is NULL.
+  @retval EFI_INVALID_PARAMETER ChildHandle is not NULL and it is not a valid EFI_HANDLE.
+  @retval EFI_INVALID_PARAMETER Language is NULL.
+  @retval EFI_INVALID_PARAMETER ControllerName is NULL.
+  @retval EFI_UNSUPPORTED       The driver specified by This is not currently managing
+                                the controller specified by ControllerHandle and
+                                ChildHandle.
+  @retval EFI_UNSUPPORTED       The driver specified by This does not support the
+                                language specified by Language.
+
+**/
+EFI_STATUS
+EFIAPI
+DcsBmlComponentNameGetControllerName (
+  IN  EFI_COMPONENT_NAME2_PROTOCOL  *This,
+  IN  EFI_HANDLE                    ControllerHandle,
+  IN  EFI_HANDLE                    ChildHandle        OPTIONAL,
+  IN  CHAR8                         *Language,
+  OUT CHAR16                        **ControllerName
+  )
+{
+	EFI_STATUS  Status = EFI_SUCCESS;
+  return Status;
+}
diff --git a/DcsBml/ComponentName.h b/DcsBml/ComponentName.h
new file mode 100644
index 0000000..8813300
--- /dev/null
+++ b/DcsBml/ComponentName.h
@@ -0,0 +1,92 @@
+/** @file
+This is DCS boot menu lock application
+
+Copyright (c) 2016. Disk Cryptography Services for EFI (DCS), Alex Kolotnikov
+
+This program and the accompanying materials
+are licensed and made available under the terms and conditions
+of the GNU Lesser General Public License, version 3.0 (LGPL-3.0).
+
+The full text of the license may be found at
+https://opensource.org/licenses/LGPL-3.0
+**/
+
+/**
+  Retrieves a Unicode string that is the user-readable name of the EFI Driver.
+
+  @param  This       A pointer to the EFI_COMPONENT_NAME_PROTOCOL instance.
+  @param  Language   A pointer to a three-character ISO 639-2 language identifier.
+                     This is the language of the driver name that that the caller
+                     is requesting, and it must match one of the languages specified
+                     in SupportedLanguages.  The number of languages supported by a
+                     driver is up to the driver writer.
+  @param  DriverName A pointer to the Unicode string to return.  This Unicode string
+                     is the name of the driver specified by This in the language
+                     specified by Language.
+
+  @retval EFI_SUCCESS           The Unicode string for the Driver specified by This
+                                and the language specified by Language was returned
+                                in DriverName.
+  @retval EFI_INVALID_PARAMETER Language is NULL.
+  @retval EFI_INVALID_PARAMETER DriverName is NULL.
+  @retval EFI_UNSUPPORTED       The driver specified by This does not support the
+                                language specified by Language.
+
+**/
+EFI_STATUS
+EFIAPI
+DcsBmlComponentNameGetDriverName (
+  IN EFI_COMPONENT_NAME2_PROTOCOL  *This,
+  IN  CHAR8                        *Language,
+  OUT CHAR16                       **DriverName
+  );
+
+/**
+  Retrieves a Unicode string that is the user readable name of the controller
+  that is being managed by an EFI Driver.
+
+  @param  This             A pointer to the EFI_COMPONENT_NAME_PROTOCOL instance.
+  @param  ControllerHandle The handle of a controller that the driver specified by
+                           This is managing.  This handle specifies the controller
+                           whose name is to be returned.
+  @param  ChildHandle      The handle of the child controller to retrieve the name
+                           of.  This is an optional parameter that may be NULL.  It
+                           will be NULL for device drivers.  It will also be NULL
+                           for a bus drivers that wish to retrieve the name of the
+                           bus controller.  It will not be NULL for a bus driver
+                           that wishes to retrieve the name of a child controller.
+  @param  Language         A pointer to a three character ISO 639-2 language
+                           identifier.  This is the language of the controller name
+                           that the caller is requesting, and it must match one
+                           of the languages specified in SupportedLanguages.  The
+                           number of languages supported by a driver is up to the
+                           driver writer.
+  @param  ControllerName   A pointer to the Unicode string to return.  This Unicode
+                           string is the name of the controller specified by
+                           ControllerHandle and ChildHandle in the language specified
+                           by Language, from the point of view of the driver specified
+                           by This.
+
+  @retval EFI_SUCCESS           The Unicode string for the user-readable name in the
+                                language specified by Language for the driver
+                                specified by This was returned in DriverName.
+  @retval EFI_INVALID_PARAMETER ControllerHandle is NULL.
+  @retval EFI_INVALID_PARAMETER ChildHandle is not NULL and it is not a valid EFI_HANDLE.
+  @retval EFI_INVALID_PARAMETER Language is NULL.
+  @retval EFI_INVALID_PARAMETER ControllerName is NULL.
+  @retval EFI_UNSUPPORTED       The driver specified by This is not currently managing
+                                the controller specified by ControllerHandle and
+                                ChildHandle.
+  @retval EFI_UNSUPPORTED       The driver specified by This does not support the
+                                language specified by Language.
+
+**/
+EFI_STATUS
+EFIAPI
+DcsBmlComponentNameGetControllerName (
+  IN  EFI_COMPONENT_NAME2_PROTOCOL  *This,
+  IN  EFI_HANDLE                    ControllerHandle,
+  IN  EFI_HANDLE                    ChildHandle        OPTIONAL,
+  IN  CHAR8                         *Language,
+  OUT CHAR16                        **ControllerName
+  );
diff --git a/DcsBml/DcsBml.c b/DcsBml/DcsBml.c
index 3376936..f5aaf3a 100644
--- a/DcsBml/DcsBml.c
+++ b/DcsBml/DcsBml.c
@@ -13,18 +13,27 @@ https://opensource.org/licenses/LGPL-3.0
 
 #include <Uefi.h>
 #include <Guid/EventGroup.h>
+#include <Guid/GlobalVariable.h>
 #include <Library/UefiBootServicesTableLib.h>
 #include <Library/UefiRuntimeLib.h>
 #include <Library/BaseLib.h>
 #include <Library/UefiLib.h>
 
+#include <Library/CommonLib.h>
+
+#include <Protocol/DcsBmlProto.h>
+#include "DcsBml.h"
+
+//////////////////////////////////////////////////////////////////////////
+// Runtime data to lock
+//////////////////////////////////////////////////////////////////////////
 typedef struct _BML_GLOBALS {
 	UINT64		Signature;
 	UINTN			size;
 } BML_GLOBALS, *PBML_GLOBALS;
 
 STATIC PBML_GLOBALS   gBmlData = NULL;
-STATIC BOOLEAN        BootMenuLocked = TRUE;
+STATIC BOOLEAN        BootMenuLocked = FALSE;
 EFI_EVENT             mBmlVirtualAddrChangeEvent;
 EFI_SET_VARIABLE      orgSetVariable = NULL;
 
@@ -71,6 +80,65 @@ BmlVirtualNotifyEvent(
 	return;
 }
 
+//////////////////////////////////////////////////////////////////////////
+// DcsBml protocol to control lock in BS mode
+//////////////////////////////////////////////////////////////////////////
+CHAR16* sDcsBootEfi = L"EFI\\VeraCrypt\\DcsBoot.efi";
+CHAR16* sDcsBootEfiDesc = L"VeraCrypt(DCS) loader";
+
+GUID gEfiDcsBmlProtocolGuid = EFI_DCSBML_INTERFACE_PROTOCOL_GUID;
+EFI_DCSBML_PROTOCOL gEfiDcsBmlProtocol = {
+    BootMenuLock
+};
+
+EFI_STATUS
+BootMenuLock(
+    IN EFI_DCSBML_PROTOCOL                *This,
+    IN     BOOLEAN                        Lock
+    ) {
+    BootMenuLocked = Lock;
+    return EFI_SUCCESS;
+}
+
+//////////////////////////////////////////////////////////////////////////
+// Driver
+//////////////////////////////////////////////////////////////////////////
+
+/**
+Unloads an image.
+
+@param  ImageHandle           Handle that identifies the image to be unloaded.
+
+@retval EFI_SUCCESS           The image has been unloaded.
+@retval EFI_INVALID_PARAMETER ImageHandle is not a valid image handle.
+
+**/
+EFI_STATUS
+EFIAPI
+DcsBmlUnload(
+    IN EFI_HANDLE  ImageHandle
+    )
+{
+    EFI_STATUS  res;
+
+    res = EFI_SUCCESS;
+    //
+    // Uninstall Driver Supported EFI Version Protocol onto ImageHandle
+    //
+    res = gBS->UninstallMultipleProtocolInterfaces(
+        ImageHandle,
+        &gEfiDcsBmlProtocolGuid, &gEfiDcsBmlProtocol,
+        NULL
+        );
+
+    if (EFI_ERROR(res)) {
+        return res;
+    }
+    // Clean up
+    return EFI_SUCCESS;
+}
+
+
 /**
 The actual entry point for the application.
 
@@ -89,7 +157,27 @@ DcsBmlMain(
    )
 {
    EFI_STATUS          res;
+   // Check multiple execution of DcsBml
+   if (!EFI_ERROR(InitBml())) {
+       return EFI_ACCESS_DENIED;
+   }
+
+   //
+   // Install DcsBml protocol onto ImageHandle
+   //
+   res = gBS->InstallMultipleProtocolInterfaces(
+       &ImageHandle,
+       &gEfiDcsBmlProtocolGuid, &gEfiDcsBmlProtocol,
+       NULL
+       );
+   ASSERT_EFI_ERROR(res);
+
+   if (EFI_ERROR(res)) {
+       Print(L"Install protocol %r\n", res);
+       return res;
+   }
 
+    // runtime lock
 	res = gBS->AllocatePool(
 		EfiRuntimeServicesData,
 		(UINTN) sizeof(BML_GLOBALS),
@@ -120,5 +208,23 @@ DcsBmlMain(
 
 	orgSetVariable = gST->RuntimeServices->SetVariable;
 	gST->RuntimeServices->SetVariable = BmlSetVaribale;
+
+    // select boot next
+    {
+        UINT16  DcsBootNum = 0x0DC5B;
+        UINTN               len;
+        UINT32              attr;
+        CHAR16*             tmp = NULL;
+        res = EfiGetVar(L"BootDC5B", &gEfiGlobalVariableGuid, &tmp, &len, &attr);
+        if (EFI_ERROR(res)) {
+            InitFS();
+            res = BootMenuItemCreate(L"BootDC5B", sDcsBootEfiDesc, gFileRootHandle, sDcsBootEfi, TRUE);
+            res = BootOrderInsert(L"BootOrder", 0, 0x0DC5B);
+        }
+        res = EfiSetVar(L"BootNext", &gEfiGlobalVariableGuid, &DcsBootNum, sizeof(DcsBootNum), EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_BOOTSERVICE_ACCESS);
+        MEM_FREE(tmp);
+    }
+
+    // Prepare BootDC5B
 	return EFI_SUCCESS;
 }
diff --git a/DcsBml/DcsBml.h b/DcsBml/DcsBml.h
new file mode 100644
index 0000000..df639fc
--- /dev/null
+++ b/DcsBml/DcsBml.h
@@ -0,0 +1,65 @@
+/** @file
+This is DCS boot menu lock application
+
+Copyright (c) 2016. Disk Cryptography Services for EFI (DCS), Alex Kolotnikov
+
+This program and the accompanying materials
+are licensed and made available under the terms and conditions
+of the GNU Lesser General Public License, version 3.0 (LGPL-3.0).
+
+The full text of the license may be found at
+https://opensource.org/licenses/LGPL-3.0
+**/
+
+#ifndef __EFI_DCSBML_H__
+#define __EFI_DCSBML_H__
+
+#include <Uefi.h>
+
+//
+// Libraries
+//
+#include <Library/UefiBootServicesTableLib.h>
+#include <Library/MemoryAllocationLib.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/BaseLib.h>
+#include <Library/UefiLib.h>
+#include <Library/DevicePathLib.h>
+#include <Library/DebugLib.h>
+
+//
+// UEFI Driver Model Protocols
+//
+#include <Protocol/ComponentName2.h>
+#include <Protocol/ComponentName.h>
+
+//
+// Consumed Protocols
+//
+
+//
+// Produced Protocols
+//
+#include <Protocol/DcsBmlProto.h>
+
+
+//
+// Protocol instances
+//
+extern EFI_COMPONENT_NAME2_PROTOCOL  gDcsBmlComponentName2;
+extern EFI_COMPONENT_NAME_PROTOCOL  gDcsBmlComponentName;
+extern EFI_DCSBML_PROTOCOL gEfiDcsBmlProtocol;
+
+//
+// Include files with function prototypes
+//
+#include "ComponentName.h"
+
+EFI_STATUS
+BootMenuLock(
+    IN EFI_DCSBML_PROTOCOL                *This,
+    IN     BOOLEAN                        Lock
+    );
+
+
+#endif
diff --git a/DcsBml/DcsBml.inf b/DcsBml/DcsBml.inf
index 10bdfda..f0b3109 100644
--- a/DcsBml/DcsBml.inf
+++ b/DcsBml/DcsBml.inf
@@ -28,6 +28,9 @@
 
 [Sources]
   DcsBml.c
+  DcsBml.h
+  ComponentName.c
+  ComponentName.h
 
 [Packages]
   MdePkg/MdePkg.dec
@@ -41,14 +44,18 @@
   UefiLib
   UefiRuntimeLib
   UefiBootServicesTableLib
+  CommonLib
 
 [Guids]
   gEfiGlobalVariableGuid
   gEfiDcsVariableGuid
   gEfiEventVirtualAddressChangeGuid
+  gEfiFileInfoGuid
 
 [Protocols]
   gEfiBlockIoProtocolGuid
+  gEfiComponentName2ProtocolGuid
+  gEfiComponentNameProtocolGuid
 
 [BuildOptions.IA32]
 RELEASE_VS2010x86_IA32_CC_FLAGS    =  /arch:IA32 /FAcs
diff --git a/DcsBoot/DcsBoot.c b/DcsBoot/DcsBoot.c
index 244eb62..90dcb5d 100644
--- a/DcsBoot/DcsBoot.c
+++ b/DcsBoot/DcsBoot.c
@@ -48,8 +48,50 @@ DoExecCmd()
 	return res;
 }
 
-CHAR16* sDcsBootEfi = L"EFI\\VeraCrypt\\DcsBoot.efi";
-CHAR16* sDcsDriverEfiDesc = L"VeraCrypt(DCS) driver";
+//////////////////////////////////////////////////////////////////////////
+// BML
+//////////////////////////////////////////////////////////////////////////
+CHAR16* sDcsBmlEfi = L"EFI\\VeraCrypt\\DcsBml.dcs";
+CHAR16* sDcsBmlEfiDesc = L"VeraCrypt(DcsBml) driver";
+CHAR16* sDcsBmlDriverVar = L"DriverDC5B";
+UINT16  sDcsBmlDriverNum = 0x0DC5B;
+
+VOID
+UpdateDriverBmlStart() {
+    EFI_STATUS          res;
+    UINTN               len;
+    UINT32              attr;
+    int                 drvInst;
+    CHAR16*             tmp = NULL;
+
+    // Driver load selected?
+    drvInst = ConfigReadInt("DcsBmlDriver", 1);
+    if (drvInst) {
+        res = EfiGetVar(sDcsBmlDriverVar, &gEfiGlobalVariableGuid, &tmp, &len, &attr);
+        // Driver installed?
+        if (EFI_ERROR(res)) {
+            // No -> install
+            res = BootMenuItemCreate(sDcsBmlDriverVar, sDcsBmlEfiDesc, gFileRootHandle, sDcsBmlEfi, FALSE);
+            if (!EFI_ERROR(res)) {
+                len = 0;
+                res = EfiGetVar(L"DriverOrder", &gEfiGlobalVariableGuid, &tmp, &len, &attr);
+                if (!EFI_ERROR(res)) len = len / 2;
+                res = BootOrderInsert(L"DriverOrder", len, sDcsBmlDriverNum);
+            }
+        }
+        MEM_FREE(tmp);
+    }
+    else {
+        // uninstall driver
+        res = EfiGetVar(sDcsBmlDriverVar, &gEfiGlobalVariableGuid, &tmp, &len, &attr);
+        if (!EFI_ERROR(res)) {
+            BootMenuItemRemove(sDcsBmlDriverVar);
+            BootOrderRemove(L"DriverOrder", sDcsBmlDriverNum);
+        }
+    }
+    MEM_FREE(tmp);
+}
+
 /**
 The actual entry point for the application.
 
@@ -67,62 +109,25 @@ DcsBootMain(
    IN EFI_SYSTEM_TABLE  *SystemTable
    )
 {
-   EFI_STATUS          res;
+   EFI_STATUS           res;
 	UINTN               len;
 	UINT32              attr;
-	int                 drvInst;
 	BOOLEAN             searchOnESP = FALSE;
-	EFI_INPUT_KEY       key;
+//	EFI_INPUT_KEY       key;
 
 	InitBio();
    res = InitFS();
    if (EFI_ERROR(res)) {
       ERR_PRINT(L"InitFS %r\n", res);
    }
-	// Check multiple execution
-	res = EfiGetVar(L"DcsExecPartGuid", NULL, &gEfiExecPartGuid, &len, &attr);
-	if (!EFI_ERROR(res)) {
-		// DcsBoot executed already.
-		ERR_PRINT(L"Multiple execution of DcsBoot\n");
-		MEM_FREE(gEfiExecPartGuid);
-		return EFI_INVALID_PARAMETER;
-	}
 
-	// Driver load selected?
-	drvInst = ConfigReadInt("DcsDriver", 0);
-	if (drvInst) {
-		CHAR16* tmp = NULL;
-		// Driver installed?
-		res = EfiGetVar(L"DriverDC5B", &gEfiGlobalVariableGuid, &tmp, &len, &attr);
-		if (EFI_ERROR(res)) {
-			// No - install and reboot.
-			res = BootMenuItemCreate(L"DriverDC5B", sDcsDriverEfiDesc, gFileRootHandle, sDcsBootEfi, FALSE);
-			if (!EFI_ERROR(res)) {
-				len = 0;
-				res = EfiGetVar(L"DriverOrder", &gEfiGlobalVariableGuid, &tmp, &len, &attr);
-				if (!EFI_ERROR(res)) len = len / 2;
-				res = BootOrderInsert(L"DriverOrder", len, 0x0DC5B);
-				OUT_PRINT(L"DcsBoot driver installed, %r\n", res);
-				key = KeyWait(L"%2d   \r", 10, 0, 0);
-				gST->RuntimeServices->ResetSystem(EfiResetCold, EFI_SUCCESS, 0, NULL);
-				return res;
-			}
-			ERR_PRINT(L"Failed to install DcsBoot driver. %r\n", res);
-			key = KeyWait(L"%2d   \r", 10, 0, 0);
-		}
-		MEM_FREE(tmp);
-	}	else {
-		CHAR16* tmp = NULL;
-		// Try uninstall driver
-		res = EfiGetVar(L"DriverDC5B", &gEfiGlobalVariableGuid, &tmp, &len, &attr);
-		if (!EFI_ERROR(res)) {
-			BootMenuItemRemove(L"DriverDC5B");
-			BootOrderRemove(L"DriverOrder", 0x0DC5B);
-			OUT_PRINT(L"DcsBoot driver uninstalled\n");
-			key = KeyWait(L"%2d   \r", 10, 0, 0);
-			gST->RuntimeServices->ResetSystem(EfiResetCold, EFI_SUCCESS, 0, NULL);
-		}
-	}
+   // BML installed?
+   if (EFI_ERROR(InitBml())) {
+       // if not -> execute
+       EfiExec(NULL, sDcsBmlEfi);
+   }
+
+   UpdateDriverBmlStart();
 
 	// Try platform info
 	if (EFI_ERROR(FileExist(NULL, L"\\EFI\\VeraCrypt\\PlatformInfo")) &&
@@ -165,12 +170,12 @@ DcsBootMain(
 	searchOnESP = CompareGuid(gEfiExecPartGuid, &ImagePartGuid) &&
 		EFI_ERROR(FileExist(NULL, gEfiExecCmd));
 
-	// Clear DcsExecPartGuid before execute OS to avoid problem in VirtualBox with reboot.
-	EfiSetVar(L"DcsExecPartGuid", NULL, NULL, 0, EFI_VARIABLE_BOOTSERVICE_ACCESS);
-	EfiSetVar(L"DcsExecCmd", NULL, NULL, 0, EFI_VARIABLE_BOOTSERVICE_ACCESS);
+    // Clear DcsExecPartGuid before execute OS to avoid problem in VirtualBox with reboot.
+    EfiSetVar(L"DcsExecPartGuid", NULL, NULL, 0, EFI_VARIABLE_BOOTSERVICE_ACCESS);
+    EfiSetVar(L"DcsExecCmd", NULL, NULL, 0, EFI_VARIABLE_BOOTSERVICE_ACCESS);
 
 	// Find new start partition
-   ConnectAllEfi();
+    ConnectAllEfi();
 	InitBio();
 	res = InitFS();
 
diff --git a/DcsInt/DcsInt.c b/DcsInt/DcsInt.c
index efa3f82..4b84256 100644
--- a/DcsInt/DcsInt.c
+++ b/DcsInt/DcsInt.c
@@ -83,12 +83,32 @@ UINTN                   SecRegionSize = 0;
 UINTN                   SecRegionOffset = 0;
 PCRYPTO_INFO            SecRegionCryptInfo = NULL;
 
+VOID
+CleanSensitiveData()
+{
+	if (SecRegionCryptInfo != NULL) {
+		MEM_BURN(SecRegionCryptInfo, sizeof(*SecRegionCryptInfo));
+	}
+
+	if (gRnd != NULL) {
+		MEM_BURN(gRnd, sizeof(*gRnd));
+	}
+
+	if (SecRegionData != NULL) {
+		MEM_BURN(SecRegionData, SecRegionSize);
+	}
+
+	if (gAutoPassword != NULL) {
+		MEM_BURN(gAutoPassword, MAX_PASSWORD);
+	}
+}
+
 void HaltPrint(const CHAR16* Msg)
 {
+	CleanSensitiveData();
 	Print(L"%s - system Halted\n", Msg);
 	EfiCpuHalt();
 }
-
 //////////////////////////////////////////////////////////////////////////
 // Boot params memory
 //////////////////////////////////////////////////////////////////////////
@@ -592,7 +612,18 @@ SecRegionChangePwd() {
 	}
 	CopyMem(&gAuthPassword, &newPassword, sizeof(gAuthPassword));
 	CopyMem(SecRegionData + SecRegionOffset, Header, 512);
+
 	ERR_PRINT(L"Update (%r)\n", Status);
+	if (!EFI_ERROR(Status)) {
+		EFI_INPUT_KEY key;
+		key = KeyWait(L"Boot OS in %2d ('r' to reset)   \r", 5, 0, 0);
+		if (key.UnicodeChar == 'r') {
+			MEM_BURN(&newPassword, sizeof(newPassword));
+			MEM_BURN(&confirmPassword, sizeof(confirmPassword));
+			CleanSensitiveData();
+			gST->RuntimeServices->ResetSystem(EfiResetCold, EFI_SUCCESS, 0, NULL);
+		}
+	}
 
 ret:
 	MEM_BURN(&newPassword, sizeof(newPassword));
@@ -935,21 +966,7 @@ VirtualNotifyEvent(
 	)
 {
 	// Clean all sensible info and keys before transfer to OS
-	if (SecRegionCryptInfo != NULL) {
-		MEM_BURN(SecRegionCryptInfo, sizeof(*SecRegionCryptInfo));
-	}
-
-	if (gRnd != NULL) {
-		MEM_BURN(gRnd, sizeof(*gRnd));
-	}
-
-	if (SecRegionData != NULL) {
-		MEM_BURN(SecRegionData, SecRegionSize);
-	}
-
-	if (gAutoPassword != NULL) {
-		MEM_BURN(gAutoPassword, MAX_PASSWORD);
-	}
+	CleanSensitiveData();
 }
 
 //////////////////////////////////////////////////////////////////////////
@@ -1137,7 +1154,8 @@ UefiMain(
 	}
 
 	// Lock EFI boot variables
-	EfiExec(NULL, L"EFI\\VeraCrypt\\DcsBml.dcs");
+    InitBml();
+    BmlLock(TRUE);
 
 	// Install decrypt
 	res = EfiLibInstallDriverBindingComponentName2(
diff --git a/DcsPkg.dec b/DcsPkg.dec
index 539fda8..0a27256 100644
--- a/DcsPkg.dec
+++ b/DcsPkg.dec
@@ -27,3 +27,6 @@
   # Include/CommonLib.h
   # {101F8560-D73A-4FF7-89F6-8170F6615587}
   gEfiDcsVariableGuid         = { 0x101f8560, 0xd73a, 0x4ff7, { 0x89, 0xf6, 0x81, 0x70, 0xf6, 0x61, 0x55, 0x87 } }
+  # Include/Protocol/DcsBmlProto.h
+  #// {7FB6D090-8755-43FC-84B5-6E297F9EC1CD}
+  gEfiDcsBmlProtocolGuid      = { 0x7fb6d090, 0x8755, 0x43fc, { 0x84, 0xb5, 0x6e, 0x29, 0x7f, 0x9e, 0xc1, 0xcd } }
diff --git a/Include/Library/CommonLib.h b/Include/Library/CommonLib.h
index 2b46632..d64ae41 100644
--- a/Include/Library/CommonLib.h
+++ b/Include/Library/CommonLib.h
@@ -15,8 +15,9 @@ https://opensource.org/licenses/LGPL-3.0
 #define __COMMONLIB_H__
 
 #include <Uefi.h>
-#include <Protocol/BlockIo.h>
 #include <Library/UefiLib.h>
+#include <Protocol/BlockIo.h>
+#include <Protocol/DcsBmlProto.h>
 #include <Protocol/SimpleFileSystem.h>
 #include <Protocol/UsbIo.h>
 #include <Protocol/AbsolutePointer.h>
@@ -551,6 +552,23 @@ SpeakerSelect(
 	);
 
 //////////////////////////////////////////////////////////////////////////
+// BML
+//////////////////////////////////////////////////////////////////////////
+extern EFI_HANDLE*                gBmlHandles;
+extern UINTN                      gBmlCount;
+extern EFI_DCSBML_PROTOCOL*	      gBml;
+extern EFI_GUID                   gBmlGuid;
+
+EFI_STATUS
+InitBml();
+
+EFI_STATUS
+BmlLock(
+    IN BOOLEAN lock
+    );
+
+
+//////////////////////////////////////////////////////////////////////////
 // Efi variables
 //////////////////////////////////////////////////////////////////////////
 
@@ -589,6 +607,11 @@ BootOrderRemove(
 	);
 
 EFI_STATUS
+BootOrderPresent(
+    IN CHAR16 *OrderVarName,
+    UINT16   value);
+
+EFI_STATUS
 BootMenuItemCreate(
 	IN CHAR16     *VarName,
 	IN CHAR16     *Desc,
diff --git a/Include/Protocol/DcsBmlProto.h b/Include/Protocol/DcsBmlProto.h
new file mode 100644
index 0000000..9fafa97
--- /dev/null
+++ b/Include/Protocol/DcsBmlProto.h
@@ -0,0 +1,51 @@
+/** @file
+This is DCS boot menu lock protocol
+
+Copyright (c) 2016. Disk Cryptography Services for EFI (DCS), Alex Kolotnikov
+
+This program and the accompanying materials
+are licensed and made available under the terms and conditions
+of the GNU Lesser General Public License, version 3.0 (LGPL-3.0).
+
+The full text of the license may be found at
+https://opensource.org/licenses/LGPL-3.0
+**/
+
+#ifndef _EFI_DCSBMLPROTO_H
+#define _EFI_DCSBMLPROTO_H
+
+#include <Uefi.h>
+#include <ProcessorBind.h>
+#include <Base.h>
+
+//
+// Global Id for DcsBml Interface
+// {7FB6D090-8755-43FC-84B5-6E297F9EC1CD}
+//
+#define EFI_DCSBML_INTERFACE_PROTOCOL_GUID \
+  { \
+    0x7fb6d090, 0x8755, 0x43fc, 0x84, 0xb5, 0x6e, 0x29, 0x7f, 0x9e, 0xc1, 0xcd \
+  }
+
+typedef struct _EFI_DCSBML_PROTOCOL  EFI_DCSBML_PROTOCOL;
+
+//
+// Lock boot menu
+//
+typedef
+EFI_STATUS
+(EFIAPI *EFI_BOOT_MENU_LOCK) (
+  IN EFI_DCSBML_PROTOCOL                *This,
+  IN     BOOLEAN                        Lock
+  );
+
+
+//
+// Protocol definition
+//
+struct _EFI_DCSBML_PROTOCOL {
+    EFI_BOOT_MENU_LOCK BootMenuLock;
+} ;
+
+extern EFI_GUID gEfiDcsBmlProtocolGuid;
+#endif
diff --git a/Library/CommonLib/CommonLib.inf b/Library/CommonLib/CommonLib.inf
index dde5c61..ec658d7 100644
--- a/Library/CommonLib/CommonLib.inf
+++ b/Library/CommonLib/CommonLib.inf
@@ -39,6 +39,7 @@
   EfiBluetooth.c
   EfiTpm.c
   GptRead.c
+  EfiBml.c
 
 [Sources.IA32]
   IA32/EfiCpuHalt.asm
diff --git a/Library/CommonLib/EfiBml.c b/Library/CommonLib/EfiBml.c
new file mode 100644
index 0000000..184ca75
--- /dev/null
+++ b/Library/CommonLib/EfiBml.c
@@ -0,0 +1,57 @@
+/** @file
+EFI DCS BML helpers routines/wrappers
+
+Copyright (c) 2016. Disk Cryptography Services for EFI (DCS), Alex Kolotnikov
+Copyright (c) 2016. VeraCrypt, Mounir IDRASSI 
+
+This program and the accompanying materials are licensed and made available
+under the terms and conditions of the GNU Lesser General Public License, version 3.0 (LGPL-3.0).
+
+The full text of the license may be found at
+https://opensource.org/licenses/LGPL-3.0
+**/
+
+#include <Library/CommonLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+#include <Protocol/DcsBmlProto.h>
+
+//////////////////////////////////////////////////////////////////////////
+// BML
+//////////////////////////////////////////////////////////////////////////
+EFI_HANDLE*                gBmlHandles = NULL;
+UINTN                      gBmlCount = 0;
+EFI_DCSBML_PROTOCOL*	   gBml = NULL;
+EFI_GUID                   gBmlGuid = EFI_DCSBML_INTERFACE_PROTOCOL_GUID;
+
+EFI_STATUS
+BmlSelect(
+    IN UINTN index) {
+    if (index < gBmlCount) {
+        return gBS->HandleProtocol(gBmlHandles[index], &gBmlGuid, (VOID**)&gBml);
+    }
+    return EFI_NOT_FOUND;
+}
+
+EFI_STATUS
+InitBml() {
+	EFI_STATUS	res;
+	// BML control if supported
+	res = EfiGetHandles(ByProtocol, &gBmlGuid, 0, &gBmlHandles, &gBmlCount);
+	if (gBmlCount > 0) {
+		return BmlSelect(gBmlCount - 1);
+    }
+    return EFI_NOT_FOUND;
+}
+
+
+EFI_STATUS
+BmlLock(
+	IN BOOLEAN lock
+	)
+{
+	if (gBml != NULL) {
+        return gBml->BootMenuLock(gBml, lock);
+	}
+	return EFI_UNSUPPORTED;
+}
+
diff --git a/Library/CommonLib/EfiVar.c b/Library/CommonLib/EfiVar.c
index fad6eb5..e9c1082 100644
--- a/Library/CommonLib/EfiVar.c
+++ b/Library/CommonLib/EfiVar.c
@@ -146,6 +146,33 @@ BootOrderRemove(
 }
 
 EFI_STATUS
+BootOrderPresent(
+    IN CHAR16 *OrderVarName,
+    UINT16   value)
+{
+    EFI_STATUS res = EFI_NOT_READY;
+    UINT16*   varBootOrder;
+    UINTN     varBootOrderSize;
+    UINT32    varBootOrderAttr;
+    UINTN     BootOrderCount;
+    UINTN     i;
+    UINTN     j;
+
+    res = EfiGetVar(OrderVarName, &gEfiGlobalVariableGuid, &varBootOrder, &varBootOrderSize, &varBootOrderAttr);
+    if (EFI_ERROR(res)) return res;
+    BootOrderCount = varBootOrderSize / sizeof(UINT16);
+    res = EFI_NOT_FOUND;
+    for (j = 0, i = 0; i < BootOrderCount; ++i) {
+        if (varBootOrder[i] == value) {
+            MEM_FREE(varBootOrder);
+            res = EFI_SUCCESS;
+            break;
+        }
+    }
+    return res;
+}
+
+EFI_STATUS
 BootMenuItemRemove(
 	IN CHAR16     *VarName
 	)