From 240149786f222327625488ed8da82941bcd07099 Mon Sep 17 00:00:00 2001
From: Oleg Raev <ovel87@gmail.com>
Date: Tue, 17 Oct 2017 04:07:42 +0700
Subject: [PATCH] Jail.RemoveCells before logout and switching account (#382)

Called jail.Stop upon SwitchAccount and Logout to ensure all jail cells have been terminated.
---
 e2e/api/api_test.go    | 72 ++++++++++++++++++++++++++++++++++++++++++
 e2e/jail/jail_test.go  | 29 +++++++++++++++++
 geth/api/api.go        |  4 +++
 geth/jail/cell_test.go |  5 +++
 4 files changed, 110 insertions(+)

diff --git a/e2e/api/api_test.go b/e2e/api/api_test.go
index 444a99630..5754d4283 100644
--- a/e2e/api/api_test.go
+++ b/e2e/api/api_test.go
@@ -12,9 +12,14 @@ import (
 	"github.com/status-im/status-go/geth/api"
 	"github.com/status-im/status-go/geth/log"
 	"github.com/status-im/status-go/geth/params"
+	. "github.com/status-im/status-go/testing"
 	"github.com/stretchr/testify/suite"
 )
 
+const (
+	testChatID = "testChat"
+)
+
 func TestAPI(t *testing.T) {
 	suite.Run(t, new(APITestSuite))
 }
@@ -118,3 +123,70 @@ func (s *APITestSuite) TestRaceConditions() {
 	time.Sleep(2 * time.Second) // so that we see some logs
 	s.api.StopNode()            // just in case we have a node running
 }
+
+func (s *APITestSuite) TestCellsRemovedAfterSwitchAccount() {
+	const itersCount = 5
+	var (
+		require   = s.Require()
+		getChatId = func(id int) string {
+			return testChatID + strconv.Itoa(id)
+		}
+	)
+
+	config, err := e2e.MakeTestNodeConfig(params.RopstenNetworkID)
+	require.NoError(err)
+	err = s.api.StartNode(config)
+	require.NoError(err)
+	defer s.api.StopNode()
+
+	address1, _, _, err := s.api.AccountManager().CreateAccount(TestConfig.Account1.Password)
+	require.NoError(err)
+
+	address2, _, _, err := s.api.AccountManager().CreateAccount(TestConfig.Account2.Password)
+	require.NoError(err)
+
+	err = s.api.SelectAccount(address1, TestConfig.Account1.Password)
+	require.NoError(err)
+
+	for i := 0; i < itersCount; i++ {
+		_, err := s.api.JailManager().NewCell(getChatId(i))
+		require.NoError(err)
+	}
+
+	err = s.api.SelectAccount(address2, TestConfig.Account2.Password)
+	require.NoError(err)
+
+	for i := 0; i < itersCount; i++ {
+		_, err := s.api.JailManager().Cell(getChatId(i))
+		require.Error(err)
+	}
+}
+
+// TestLogoutRemovesCells we want be sure that
+// cells will be removed after the API call "Logout"
+func (s *APITestSuite) TestLogoutRemovesCells() {
+	var (
+		err     error
+		require = s.Require()
+	)
+
+	config, err := e2e.MakeTestNodeConfig(params.RopstenNetworkID)
+	require.NoError(err)
+	err = s.api.StartNode(config)
+	require.NoError(err)
+	defer s.api.StopNode()
+
+	address1, _, _, err := s.api.AccountManager().CreateAccount(TestConfig.Account1.Password)
+	require.NoError(err)
+
+	err = s.api.SelectAccount(address1, TestConfig.Account1.Password)
+	require.NoError(err)
+
+	s.api.JailManager().Parse(testChatID, ``)
+
+	err = s.api.Logout()
+	require.NoError(err)
+
+	_, err = s.api.JailManager().Cell(testChatID)
+	require.Error(err, "Expected that cells was removed")
+}
diff --git a/e2e/jail/jail_test.go b/e2e/jail/jail_test.go
index b22036410..7462bb943 100644
--- a/e2e/jail/jail_test.go
+++ b/e2e/jail/jail_test.go
@@ -3,6 +3,7 @@ package jail
 import (
 	"encoding/json"
 	"errors"
+	"strconv"
 	"testing"
 	"time"
 
@@ -157,3 +158,31 @@ func (s *JailTestSuite) TestEventSignal() {
 	expectedResponse := `{"jsonrpc":"2.0","result":true}`
 	s.Equal(expectedResponse, response)
 }
+
+func (s *JailTestSuite) TestJailCellsRemovedAfterStop() {
+	const loopLen = 5
+
+	getTestCellID := func(id int) string {
+		return testChatID + strconv.Itoa(id)
+	}
+	require := s.Require()
+
+	for i := 0; i < loopLen; i++ {
+		s.jail.Parse(getTestCellID(i), "")
+		cell, err := s.jail.Cell(getTestCellID(i))
+		require.NoError(err)
+		_, err = cell.Run(`
+			var counter = 1;
+			setInterval(function(){
+				counter++;
+			}, 1000);
+		`)
+	}
+
+	s.jail.Stop()
+
+	for i := 0; i < loopLen; i++ {
+		_, err := s.jail.Cell(getTestCellID(i))
+		require.Error(err, "Expected cells removing (from Jail) after stop")
+	}
+}
diff --git a/geth/api/api.go b/geth/api/api.go
index 07cdaaa2a..8afcf628c 100644
--- a/geth/api/api.go
+++ b/geth/api/api.go
@@ -146,11 +146,15 @@ func (api *StatusAPI) VerifyAccountPassword(keyStoreDir, address, password strin
 // using provided password. Once verification is done, decrypted key is injected into Whisper (as a single identity,
 // all previous identities are removed).
 func (api *StatusAPI) SelectAccount(address, password string) error {
+	// FIXME(oleg-raev): This method doesn't make stop, it rather resets its cells to an initial state
+	// and should be properly renamed, for example: ResetCells
+	api.b.jailManager.Stop()
 	return api.b.AccountManager().SelectAccount(address, password)
 }
 
 // Logout clears whisper identities
 func (api *StatusAPI) Logout() error {
+	api.b.jailManager.Stop()
 	return api.b.AccountManager().Logout()
 }
 
diff --git a/geth/jail/cell_test.go b/geth/jail/cell_test.go
index 600e5e918..0d8fed8d0 100644
--- a/geth/jail/cell_test.go
+++ b/geth/jail/cell_test.go
@@ -3,6 +3,7 @@ package jail_test
 import (
 	"net/http"
 	"net/http/httptest"
+	"testing"
 	"time"
 
 	"github.com/robertkrimen/otto"
@@ -19,6 +20,10 @@ var (
 	baseStatusJSCode = string(static.MustAsset("testdata/jail/status.js"))
 )
 
+func TestCellTestSuite(t *testing.T) {
+	suite.Run(t, new(CellTestSuite))
+}
+
 type CellTestSuite struct {
 	suite.Suite
 	jail *jail.Jail