diff --git a/public/javascripts/note_editor.js b/public/javascripts/note_editor.js
index c2e328a96..bb75a2889 100644
--- a/public/javascripts/note_editor.js
+++ b/public/javascripts/note_editor.js
@@ -79,18 +79,11 @@ const noteEditor = (function() {
     }
 
     function setNoteBackgroundIfProtected(note) {
-        if (note.detail.is_protected) {
-            $("#note-detail-wrapper").addClass("protected");
-            protectButton.hide();
-            unprotectButton.show();
-        }
-        else {
-            $("#note-detail-wrapper").removeClass("protected");
-            protectButton.show();
-            unprotectButton.hide();
-        }
+        const isProtected = !!note.detail.is_protected;
 
-        noteTree.setCurrentNoteTreeBasedOnProtectedStatus();
+        noteDetailWrapperEl.toggleClass("protected", isProtected);
+        protectButton.toggle(!isProtected);
+        unprotectButton.toggle(isProtected);
     }
 
     let isNewNoteCreated = false;
@@ -129,6 +122,7 @@ const noteEditor = (function() {
         noteChangeDisabled = false;
 
         setNoteBackgroundIfProtected(currentNote);
+        noteTree.setNoteTreeBackgroundBasedOnProtectedStatus(noteId);
 
         showAppIfHidden();
     }
diff --git a/public/javascripts/note_tree.js b/public/javascripts/note_tree.js
index c47a19ef7..417d8ff0a 100644
--- a/public/javascripts/note_tree.js
+++ b/public/javascripts/note_tree.js
@@ -617,8 +617,14 @@ const noteTree = (function() {
         }
     }
 
-    function setCurrentNoteTreeBasedOnProtectedStatus() {
-        getCurrentClones().map(node => node.toggleClass("protected", !!node.data.is_protected));
+    function setNoteTreeBackgroundBasedOnProtectedStatus(noteId) {
+        getNodesByNoteId(noteId).map(node => node.toggleClass("protected", !!node.data.is_protected));
+    }
+
+    function setProtected(noteId, isProtected) {
+        getNodesByNoteId(noteId).map(node => node.data.is_protected = isProtected);
+
+        setNoteTreeBackgroundBasedOnProtectedStatus(noteId);
     }
 
     function getAutocompleteItems(parentNoteId, notePath, titlePath) {
@@ -776,7 +782,8 @@ const noteTree = (function() {
         reload,
         collapseTree,
         scrollToCurrentNote,
-        setCurrentNoteTreeBasedOnProtectedStatus,
+        setNoteTreeBackgroundBasedOnProtectedStatus,
+        setProtected,
         getCurrentNode,
         activateNode,
         getCurrentNotePath,
diff --git a/public/javascripts/protected_session.js b/public/javascripts/protected_session.js
index dfa1ead8d..ae4e469bd 100644
--- a/public/javascripts/protected_session.js
+++ b/public/javascripts/protected_session.js
@@ -115,6 +115,8 @@ const protected_session = (function() {
 
         await noteEditor.saveNoteToServer(note);
 
+        noteTree.setProtected(note.detail.note_id, note.detail.is_protected);
+
         noteEditor.setNoteBackgroundIfProtected(note);
     }
 
@@ -129,6 +131,8 @@ const protected_session = (function() {
 
         await noteEditor.saveNoteToServer(note);
 
+        noteTree.setProtected(note.detail.note_id, note.detail.is_protected);
+
         noteEditor.setNoteBackgroundIfProtected(note);
     }
 
diff --git a/views/index.ejs b/views/index.ejs
index 3158e4f96..260226a58 100644
--- a/views/index.ejs
+++ b/views/index.ejs
@@ -171,9 +171,10 @@
       <form id="protected-session-password-form">
         <div class="form-group">
           <label for="protected-session-password">To proceed with requested action you need to enter protected session by entering password:</label>
-          <input id="protected-session-password" style="width: 250px;" type="password">
-          <button class="btn btn-sm">Show</button>
+          <input id="protected-session-password" class="form-control" type="password">
         </div>
+
+        <button class="btn btn-sm">Enter protected session</button>
       </form>
     </div>