From 3c50523025bfe276d3828bc9fa7920c842ec72eb Mon Sep 17 00:00:00 2001
From: Thomas Kriechbaumer <thomas@kriechbaumer.name>
Date: Mon, 16 Nov 2020 22:38:34 +0100
Subject: [PATCH] command-history: fail-safe file handling

---
 CHANGELOG.rst                                 |  3 +-
 mitmproxy/addons/command_history.py           | 18 +++++--
 test/mitmproxy/addons/test_command_history.py | 49 ++++++++++++++++---
 3 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index a5e45c264..b365fc291 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -4,7 +4,8 @@ Release History
 Unreleased: mitmproxy next
 ==========================
 
-* Fix query parameters in asgiapp addon (@jpstotz) 
+* Fix query parameters in asgiapp addon (@jpstotz)
+* Fix command history failing on file IO errors (@Kriechi)
 * --- TODO: add new PRs above this line ---
 
 * ... and various other fixes, documentation improvements, dependency version bumps, etc.
diff --git a/mitmproxy/addons/command_history.py b/mitmproxy/addons/command_history.py
index d5aaa9280..a54a34462 100644
--- a/mitmproxy/addons/command_history.py
+++ b/mitmproxy/addons/command_history.py
@@ -39,7 +39,10 @@ class CommandHistory:
         if ctx.options.command_history and len(self.history) >= self.VACUUM_SIZE:
             # vacuum history so that it doesn't grow indefinitely.
             history_str = "\n".join(self.history[-self.VACUUM_SIZE // 2:]) + "\n"
-            self.history_file.write_text(history_str)
+            try:
+                self.history_file.write_text(history_str)
+            except Exception as e:
+                ctx.log.alert(f"Failed writing to {self.history_file}: {e}")
 
     @command.command("commands.history.add")
     def add_command(self, command: str) -> None:
@@ -48,9 +51,11 @@ class CommandHistory:
 
         self.history.append(command)
         if ctx.options.command_history:
-            with self.history_file.open("a") as f:
-                f.write(f"{command}\n")
-                f.close()
+            try:
+                with self.history_file.open("a") as f:
+                    f.write(f"{command}\n")
+            except Exception as e:
+                ctx.log.alert(f"Failed writing to {self.history_file}: {e}")
 
         self.set_filter('')
 
@@ -62,7 +67,10 @@ class CommandHistory:
     @command.command("commands.history.clear")
     def clear_history(self):
         if self.history_file.exists():
-            self.history_file.unlink()
+            try:
+                self.history_file.unlink()
+            except Exception as e:
+                ctx.log.alert(f"Failed deleting {self.history_file}: {e}")
         self.history = []
         self.set_filter('')
 
diff --git a/test/mitmproxy/addons/test_command_history.py b/test/mitmproxy/addons/test_command_history.py
index 245bbc261..735d44094 100644
--- a/test/mitmproxy/addons/test_command_history.py
+++ b/test/mitmproxy/addons/test_command_history.py
@@ -1,4 +1,8 @@
 import os
+from unittest.mock import patch
+from pathlib import Path
+
+import pytest
 
 from mitmproxy.addons import command_history
 from mitmproxy.test import taddons
@@ -22,16 +26,35 @@ class TestCommandHistory:
         with open(history_file, "r") as f:
             assert f.read() == "cmd3\ncmd4\n"
 
+    @pytest.mark.asyncio
+    async def test_done_writing_failed(self):
+        ch = command_history.CommandHistory()
+        ch.VACUUM_SIZE = 1
+        with taddons.context(ch) as tctx:
+            ch.history.append('cmd1')
+            ch.history.append('cmd2')
+            ch.history.append('cmd3')
+            tctx.options.confdir = '/non/existent/path/foobar1234/'
+            ch.done()
+            assert await tctx.master.await_log(f"Failed writing to {ch.history_file}")
+
     def test_add_command(self):
-        history = command_history.CommandHistory()
+        ch = command_history.CommandHistory()
 
-        history.add_command('cmd1')
-        history.add_command('cmd2')
+        ch.add_command('cmd1')
+        ch.add_command('cmd2')
+        assert ch.history == ['cmd1', 'cmd2']
 
-        assert history.history == ['cmd1', 'cmd2']
+        ch.add_command('')
+        assert ch.history == ['cmd1', 'cmd2']
 
-        history.add_command('')
-        assert history.history == ['cmd1', 'cmd2']
+    @pytest.mark.asyncio
+    async def test_add_command_failed(self):
+        ch = command_history.CommandHistory()
+        with taddons.context(ch) as tctx:
+            tctx.options.confdir = '/non/existent/path/foobar1234/'
+            ch.add_command('cmd1')
+            assert await tctx.master.await_log(f"Failed writing to {ch.history_file}")
 
     def test_get_next_and_prev(self, tmpdir):
         ch = command_history.CommandHistory()
@@ -133,6 +156,20 @@ class TestCommandHistory:
 
             ch.clear_history()
 
+    @pytest.mark.asyncio
+    async def test_clear_failed(self, monkeypatch):
+        ch = command_history.CommandHistory()
+
+        with taddons.context(ch) as tctx:
+            tctx.options.confdir = '/non/existent/path/foobar1234/'
+
+            with patch.object(Path, 'exists') as mock_exists:
+                mock_exists.return_value = True
+                with patch.object(Path, 'unlink') as mock_unlink:
+                    mock_unlink.side_effect = IOError()
+                    ch.clear_history()
+            assert await tctx.master.await_log(f"Failed deleting {ch.history_file}")
+
     def test_filter(self, tmpdir):
         ch = command_history.CommandHistory()