forked from Silverfish/proton-bridge
We build too many walls and not enough bridges
This commit is contained in:
Jakub
151
internal/imap/cache/cache.go
vendored
Normal file
151
internal/imap/cache/cache.go
vendored
Normal file
@ -0,0 +1,151 @@
|
||||
// Copyright (c) 2020 Proton Technologies AG
|
||||
//
|
||||
// This file is part of ProtonMail Bridge.
|
||||
//
|
||||
// ProtonMail Bridge is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// ProtonMail Bridge is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with ProtonMail Bridge. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
package cache
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
backendMessage "github.com/ProtonMail/proton-bridge/pkg/message"
|
||||
)
|
||||
|
||||
type key struct {
|
||||
ID string
|
||||
Timestamp int64
|
||||
Size int
|
||||
}
|
||||
|
||||
type oldestFirst []key
|
||||
|
||||
func (s oldestFirst) Len() int { return len(s) }
|
||||
func (s oldestFirst) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
|
||||
func (s oldestFirst) Less(i, j int) bool { return s[i].Timestamp < s[j].Timestamp }
|
||||
|
||||
type cachedMessage struct {
|
||||
key
|
||||
data []byte
|
||||
structure backendMessage.BodyStructure
|
||||
}
|
||||
|
||||
//nolint[gochecknoglobals]
|
||||
var (
|
||||
cacheTimeLimit = int64(1 * 60 * 60 * 1000) // milliseconds
|
||||
cacheSizeLimit = 100 * 1000 * 1000 // B - MUST be larger than email max size limit (~ 25 MB)
|
||||
mailCache = make(map[string]cachedMessage)
|
||||
|
||||
// cacheMutex takes care of one single operation, whereas buildMutex takes
|
||||
// care of the whole action doing multiple operations. buildMutex will protect
|
||||
// you from asking server or decrypting or building the same message more
|
||||
// than once. When first request to build the message comes, it will block
|
||||
// all other build requests. When the first one is done, all others are
|
||||
// handled by cache, not doing anything twice. With cacheMutex we are safe
|
||||
// only to not mess up with the cache, but we could end up downloading and
|
||||
// building message twice.
|
||||
cacheMutex = &sync.Mutex{}
|
||||
buildMutex = &sync.Mutex{}
|
||||
buildLocks = map[string]interface{}{}
|
||||
)
|
||||
|
||||
func (m *cachedMessage) isValidOrDel() bool {
|
||||
if m.key.Timestamp+cacheTimeLimit < timestamp() {
|
||||
delete(mailCache, m.key.ID)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func timestamp() int64 {
|
||||
return time.Now().UnixNano() / int64(time.Millisecond)
|
||||
}
|
||||
|
||||
func Clear() {
|
||||
mailCache = make(map[string]cachedMessage)
|
||||
}
|
||||
|
||||
// BuildLock locks per message level, not on global level.
|
||||
// Multiple different messages can be building at once.
|
||||
func BuildLock(messageID string) {
|
||||
for {
|
||||
buildMutex.Lock()
|
||||
if _, ok := buildLocks[messageID]; ok { // if locked, wait
|
||||
buildMutex.Unlock()
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
} else { // if unlocked, lock it
|
||||
buildLocks[messageID] = struct{}{}
|
||||
buildMutex.Unlock()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BuildUnlock(messageID string) {
|
||||
buildMutex.Lock()
|
||||
defer buildMutex.Unlock()
|
||||
delete(buildLocks, messageID)
|
||||
}
|
||||
|
||||
func LoadMail(mID string) (reader *bytes.Reader, structure *backendMessage.BodyStructure) {
|
||||
reader = &bytes.Reader{}
|
||||
cacheMutex.Lock()
|
||||
defer cacheMutex.Unlock()
|
||||
if message, ok := mailCache[mID]; ok && message.isValidOrDel() {
|
||||
reader = bytes.NewReader(message.data)
|
||||
structure = &message.structure
|
||||
|
||||
// Update timestamp to keep emails which are used often.
|
||||
message.Timestamp = timestamp()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func SaveMail(mID string, msg []byte, structure *backendMessage.BodyStructure) {
|
||||
cacheMutex.Lock()
|
||||
defer cacheMutex.Unlock()
|
||||
|
||||
newMessage := cachedMessage{
|
||||
key: key{
|
||||
ID: mID,
|
||||
Timestamp: timestamp(),
|
||||
Size: len(msg),
|
||||
},
|
||||
data: msg,
|
||||
structure: *structure,
|
||||
}
|
||||
|
||||
// Remove old and reduce size.
|
||||
totalSize := 0
|
||||
messageList := []key{}
|
||||
for _, message := range mailCache {
|
||||
if message.isValidOrDel() {
|
||||
messageList = append(messageList, message.key)
|
||||
totalSize += message.key.Size
|
||||
}
|
||||
}
|
||||
sort.Sort(oldestFirst(messageList))
|
||||
var oldest key
|
||||
for totalSize+newMessage.key.Size >= cacheSizeLimit {
|
||||
oldest, messageList = messageList[0], messageList[1:]
|
||||
delete(mailCache, oldest.ID)
|
||||
totalSize -= oldest.Size
|
||||
}
|
||||
|
||||
// Write new.
|
||||
mailCache[mID] = newMessage
|
||||
}
|
||||
99
internal/imap/cache/cache_test.go
vendored
Normal file
99
internal/imap/cache/cache_test.go
vendored
Normal file
@ -0,0 +1,99 @@
|
||||
// Copyright (c) 2020 Proton Technologies AG
|
||||
//
|
||||
// This file is part of ProtonMail Bridge.
|
||||
//
|
||||
// ProtonMail Bridge is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// ProtonMail Bridge is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with ProtonMail Bridge. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
package cache
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
bckMsg "github.com/ProtonMail/proton-bridge/pkg/message"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
var bs = &bckMsg.BodyStructure{} //nolint[gochecknoglobals]
|
||||
const testUID = "testmsg"
|
||||
|
||||
func TestSaveAndLoad(t *testing.T) {
|
||||
msg := []byte("Test message")
|
||||
|
||||
SaveMail(testUID, msg, bs)
|
||||
require.Equal(t, mailCache[testUID].data, msg)
|
||||
|
||||
reader, _ := LoadMail(testUID)
|
||||
require.Equal(t, reader.Len(), len(msg))
|
||||
stored := make([]byte, len(msg))
|
||||
_, _ = reader.Read(stored)
|
||||
require.Equal(t, stored, msg)
|
||||
}
|
||||
|
||||
func TestMissing(t *testing.T) {
|
||||
reader, _ := LoadMail("non-existing")
|
||||
require.Equal(t, reader.Len(), 0)
|
||||
}
|
||||
|
||||
func TestClearOld(t *testing.T) {
|
||||
cacheTimeLimit = 10
|
||||
msg := []byte("Test message")
|
||||
SaveMail(testUID, msg, bs)
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
reader, _ := LoadMail(testUID)
|
||||
require.Equal(t, reader.Len(), 0)
|
||||
}
|
||||
|
||||
func TestClearBig(t *testing.T) {
|
||||
msg := []byte("Test message")
|
||||
|
||||
nSize := 3
|
||||
cacheSizeLimit = nSize*len(msg) + 1
|
||||
cacheTimeLimit = int64(nSize * nSize * 2) // be sure the message will survive
|
||||
|
||||
// It should have more than nSize items.
|
||||
for i := 0; i < nSize*nSize; i++ {
|
||||
time.Sleep(1 * time.Millisecond)
|
||||
SaveMail(fmt.Sprintf("%s%d", testUID, i), msg, bs)
|
||||
if len(mailCache) > nSize {
|
||||
t.Error("Number of items in cache should not be more than", nSize)
|
||||
}
|
||||
}
|
||||
|
||||
// Check that the oldest are deleted first.
|
||||
for i := 0; i < nSize*nSize; i++ {
|
||||
iUID := fmt.Sprintf("%s%d", testUID, i)
|
||||
reader, _ := LoadMail(iUID)
|
||||
if i < nSize*(nSize-1) && reader.Len() != 0 {
|
||||
mail := mailCache[iUID]
|
||||
t.Error("LoadMail should return empty but have:", mail.data, iUID, mail.key.Timestamp)
|
||||
}
|
||||
stored := make([]byte, len(msg))
|
||||
_, _ = reader.Read(stored)
|
||||
|
||||
if i >= nSize*(nSize-1) && !bytes.Equal(stored, msg) {
|
||||
t.Error("LoadMail returned wrong message:", stored, iUID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurency(t *testing.T) {
|
||||
msg := []byte("Test message")
|
||||
for i := 0; i < 10; i++ {
|
||||
go SaveMail(fmt.Sprintf("%s%d", testUID, i), msg, bs)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user