Skip to content

Commit 19da986

Browse files
annatarannatar
committed
feat(BRIDGE-464): add the ability to limit the amount of IMAP connections via client
1 parent d2a0b8c commit 19da986

9 files changed

Lines changed: 432 additions & 2 deletions

File tree

‎builder.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"github.com/ProtonMail/gluon/db"
1212
"github.com/ProtonMail/gluon/imap"
1313
"github.com/ProtonMail/gluon/imap/connectioncounter"
14+
"github.com/ProtonMail/gluon/imap/connectionlimiter"
1415
"github.com/ProtonMail/gluon/internal/backend"
1516
"github.com/ProtonMail/gluon/internal/db_impl/sqlite3"
1617
"github.com/ProtonMail/gluon/internal/session"
@@ -46,6 +47,7 @@ type serverBuilder struct {
4647
observabilitySender observability.Sender
4748
featureFlagProvider unleash.FeatureFlagValueProvider
4849
connectionRollingCounter *connectioncounter.RollingCounter
50+
connectionLimiter connectionlimiter.ConnectionLimiter
4951
}
5052

5153
func newBuilder() (*serverBuilder, error) {
@@ -139,6 +141,7 @@ func (builder *serverBuilder) build() (*Server, error) {
139141
observabilitySender: builder.observabilitySender,
140142
connectionRollingCounter: builder.connectionRollingCounter,
141143
featureFlagProvider: builder.featureFlagProvider,
144+
clientLimiter: builder.connectionLimiter,
142145
}
143146

144147
return s, nil

‎imap/connectioncounter/connectioncounter.go‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -138,10 +138,21 @@ func (rc *RollingCounter) GetRollingCount() int {
138138
rc.bucketLock.Lock()
139139
defer rc.bucketLock.Unlock()
140140

141-
var rollingCount int
141+
return rc.getRollingCountUnlocked()
142+
}
143+
144+
func (rc *RollingCounter) OverThreshold() bool {
145+
rc.bucketLock.Lock()
146+
defer rc.bucketLock.Unlock()
147+
148+
return rc.getRollingCountUnlocked() >= rc.newConnectionThreshold
149+
}
150+
151+
func (rc *RollingCounter) getRollingCountUnlocked() int {
152+
var rollingCount int
142153
for _, count := range rc.buckets {
143154
rollingCount += count
144155
}
145156

146157
return rollingCount
147-
}
158+
}

‎imap/connectionlimiter/client.go‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
package connectionlimiter
2+
3+
import (
4+
"strings"
5+
6+
"github.com/ProtonMail/gluon/imap"
7+
)
8+
9+
type Client string
10+
11+
const (
12+
ClientAppleMail Client = "apple-mail"
13+
ClientOutlook Client = "outlook"
14+
ClientThunderbird Client = "thunderbird"
15+
ClientUnknown Client = "unknown"
16+
)
17+
18+
func normalizeClientKey(id imap.IMAPID) Client {
19+
name := strings.TrimSpace(strings.ToLower(id.Name))
20+
switch {
21+
case strings.Contains(name, "mac") || strings.Contains(name, "os x") || strings.Contains(name, "apple"):
22+
return ClientAppleMail
23+
case strings.Contains(name, "outlook"):
24+
return ClientOutlook
25+
case strings.Contains(name, "thunderbird"):
26+
return ClientThunderbird
27+
default:
28+
return ClientUnknown
29+
}
30+
}

‎imap/connectionlimiter/limiter.go‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
package connectionlimiter
2+
3+
import (
4+
"sync"
5+
6+
"github.com/ProtonMail/gluon/imap"
7+
"github.com/sirupsen/logrus"
8+
)
9+
10+
type ConnectionLimiter interface {
11+
//TryBind tries to bind a given sessionID to a client.
12+
TryBind(sessionID int, id imap.IMAPID) (allowed bool, key Client, current int, max int)
13+
14+
//Unbind unbinds a given sessionID from a client.
15+
Unbind(sessionID int)
16+
}
17+
18+
type limiter struct {
19+
mu sync.Mutex
20+
limits Limits
21+
22+
//sessionID mapping to normalized client key
23+
sessionClient map[int]Client
24+
25+
//normalized key mapping to current open sessions
26+
clientCount map[Client]int
27+
28+
log *logrus.Entry
29+
}
30+
31+
func NewConnectionLimiter(limits Limits) ConnectionLimiter {
32+
return newLimiter(limits)
33+
}
34+
35+
func newLimiter(limits Limits) *limiter {
36+
log := logrus.WithFields(logrus.Fields{
37+
"pkg": "gluon/connectionlimiter",
38+
"limits": limits,
39+
})
40+
41+
return &limiter{
42+
limits: limits,
43+
sessionClient: make(map[int]Client),
44+
clientCount: make(map[Client]int),
45+
log: log,
46+
}
47+
}
48+
49+
func (l *limiter) TryBind(sessionID int, id imap.IMAPID) (allowed bool, key Client, current int, max int) {
50+
l.mu.Lock()
51+
defer l.mu.Unlock()
52+
53+
key = normalizeClientKey(id)
54+
55+
// already bound to this client, no-op allow
56+
if prev, ok := l.sessionClient[sessionID]; ok && prev == key {
57+
maxUsages := l.maxForKey(key)
58+
l.log.WithFields(logrus.Fields{
59+
"sessionID": sessionID,
60+
"client": key,
61+
"current": l.clientCount[key],
62+
"max": max,
63+
}).Info("Already bound to this client, no-op allow")
64+
return true, key, l.clientCount[key], maxUsages
65+
}
66+
67+
// if rebind, release the old key first
68+
if prev, ok := l.sessionClient[sessionID]; ok {
69+
if c := l.clientCount[prev]; c > 0 {
70+
l.log.WithFields(logrus.Fields{
71+
"sessionID": sessionID,
72+
"client": prev,
73+
"current": c,
74+
}).Info("Releasing old client")
75+
l.clientCount[prev] = c - 1
76+
}
77+
78+
}
79+
80+
max = l.maxForKey(key)
81+
cur := l.clientCount[key]
82+
83+
if max > 0 && cur >= max {
84+
delete(l.sessionClient, sessionID)
85+
return false, key, cur, max
86+
}
87+
88+
l.clientCount[key] = cur + 1
89+
l.sessionClient[sessionID] = key
90+
91+
l.log.WithFields(logrus.Fields{
92+
"sessionID": sessionID,
93+
"client": key,
94+
"current": l.clientCount[key],
95+
"max": max,
96+
}).Debug("Binding session to client")
97+
98+
return true, key, l.clientCount[key], max
99+
}
100+
101+
func (l *limiter) Unbind(sessionID int) {
102+
l.mu.Lock()
103+
defer l.mu.Unlock()
104+
105+
key, ok := l.sessionClient[sessionID]
106+
if !ok {
107+
return
108+
}
109+
delete(l.sessionClient, sessionID)
110+
111+
if c := l.clientCount[key]; c > 1 {
112+
l.clientCount[key] = c - 1
113+
114+
l.log.WithFields(logrus.Fields{
115+
"sessionID": sessionID,
116+
"client": key,
117+
"current": l.clientCount[key],
118+
}).Debug("Unbinding session from client")
119+
120+
} else {
121+
delete(l.clientCount, key)
122+
123+
l.log.WithFields(logrus.Fields{
124+
"sessionID": sessionID,
125+
"client": key,
126+
}).Debug("Unbinding session from client")
127+
}
128+
}
129+
130+
func (l *limiter) maxForKey(key Client) int {
131+
if limit, ok := l.limits.PerClient[key]; ok {
132+
return limit
133+
}
134+
135+
l.log.WithFields(logrus.Fields{
136+
"client": key,
137+
"limit": l.limits.UnknownLimit,
138+
}).Debug("Client key not found in limits, using unknown limit")
139+
140+
return l.limits.UnknownLimit
141+
}

‎imap/connectionlimiter/limiter_test.go‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
package connectionlimiter
2+
3+
import (
4+
"testing"
5+
6+
"github.com/ProtonMail/gluon/imap"
7+
"github.com/stretchr/testify/assert"
8+
"github.com/stretchr/testify/require"
9+
)
10+
11+
func imapID(name string) imap.IMAPID {
12+
return imap.IMAPID{Name: name}
13+
}
14+
15+
func TestTryBind_FirstConnectionAllowed(t *testing.T) {
16+
limits := NewLimits(map[Client]int{ClientAppleMail: 3}, 1)
17+
l := NewConnectionLimiter(limits)
18+
allowed, key, current, max := l.TryBind(1, imapID("Apple Mail"))
19+
require.True(t, allowed)
20+
assert.Equal(t, ClientAppleMail, key)
21+
assert.Equal(t, 1, current)
22+
assert.Equal(t, 3, max)
23+
}
24+
25+
func TestTryBind_OverLimitDenied(t *testing.T) {
26+
limits := NewLimits(map[Client]int{ClientAppleMail: 2}, 1)
27+
l := NewConnectionLimiter(limits)
28+
l.TryBind(1, imapID("Apple Mail"))
29+
l.TryBind(2, imapID("Apple Mail"))
30+
allowed, key, current, max := l.TryBind(3, imapID("Apple Mail"))
31+
require.False(t, allowed)
32+
assert.Equal(t, ClientAppleMail, key)
33+
assert.Equal(t, 2, current)
34+
assert.Equal(t, 2, max)
35+
}
36+
37+
func TestUnbind_FreesSlot(t *testing.T) {
38+
limits := NewLimits(map[Client]int{ClientAppleMail: 2}, 1)
39+
l := NewConnectionLimiter(limits)
40+
l.TryBind(1, imapID("Apple Mail"))
41+
l.TryBind(2, imapID("Apple Mail"))
42+
l.Unbind(1)
43+
allowed, _, current, _ := l.TryBind(3, imapID("Apple Mail"))
44+
require.True(t, allowed)
45+
assert.Equal(t, 2, current)
46+
}
47+
48+
func TestUnbind_UnknownSessionNoop(t *testing.T) {
49+
limits := NewLimits(map[Client]int{ClientAppleMail: 2}, 1)
50+
l := NewConnectionLimiter(limits)
51+
l.TryBind(1, imapID("Apple Mail"))
52+
l.Unbind(999) // never bound
53+
_, _, current, _ := l.TryBind(2, imapID("Apple Mail"))
54+
assert.Equal(t, 2, current)
55+
}
56+
57+
func TestTryBind_SameSessionSameClientNoop(t *testing.T) {
58+
limits := NewLimits(map[Client]int{ClientAppleMail: 2}, 1)
59+
l := NewConnectionLimiter(limits)
60+
l.TryBind(1, imapID("Apple Mail"))
61+
allowed, key, current, max := l.TryBind(1, imapID("Apple Mail"))
62+
require.True(t, allowed)
63+
assert.Equal(t, ClientAppleMail, key)
64+
assert.Equal(t, 1, current)
65+
assert.Equal(t, 2, max)
66+
}
67+
68+
func TestTryBind_RebindToDifferentClient(t *testing.T) {
69+
limits := NewLimits(map[Client]int{
70+
ClientAppleMail: 2,
71+
ClientOutlook: 2,
72+
}, 1)
73+
l := NewConnectionLimiter(limits)
74+
l.TryBind(1, imapID("Apple Mail"))
75+
allowed, key, cur, _ := l.TryBind(1, imapID("Microsoft Outlook"))
76+
require.True(t, allowed)
77+
assert.Equal(t, ClientOutlook, key)
78+
assert.Equal(t, 1, cur)
79+
_, _, appleCur, _ := l.TryBind(2, imapID("Apple Mail"))
80+
assert.Equal(t, 1, appleCur)
81+
}
82+
83+
func TestTryBind_UnknownClientUsesUnknownLimit(t *testing.T) {
84+
limits := NewLimits(map[Client]int{ClientAppleMail: 10}, 2)
85+
l := NewConnectionLimiter(limits)
86+
l.TryBind(1, imapID("SomeOtherClient"))
87+
l.TryBind(2, imapID("Unknown"))
88+
allowed, key, current, max := l.TryBind(3, imapID("Custom"))
89+
require.False(t, allowed)
90+
assert.Equal(t, ClientUnknown, key)
91+
assert.Equal(t, 2, current)
92+
assert.Equal(t, 2, max)
93+
}
94+
95+
func TestTryBind_UnlimitedLimit(t *testing.T) {
96+
limits := NewLimits(map[Client]int{ClientAppleMail: 0}, 1)
97+
l := NewConnectionLimiter(limits)
98+
for i := 1; i <= 5; i++ {
99+
allowed, key, current, max := l.TryBind(i, imapID("Apple Mail"))
100+
require.True(t, allowed)
101+
assert.Equal(t, ClientAppleMail, key)
102+
assert.Equal(t, i, current)
103+
assert.Equal(t, 0, max)
104+
}
105+
}
106+
107+
func TestNormalizeClientKey_AppleMail(t *testing.T) {
108+
assert.Equal(t, ClientAppleMail, normalizeClientKey(imapID("Apple Mail")))
109+
assert.Equal(t, ClientAppleMail, normalizeClientKey(imapID("macos mail")))
110+
assert.Equal(t, ClientAppleMail, normalizeClientKey(imapID("OS X Mail")))
111+
}
112+
func TestNormalizeClientKey_OutlookAndThunderbird(t *testing.T) {
113+
assert.Equal(t, ClientOutlook, normalizeClientKey(imapID("Microsoft Outlook")))
114+
assert.Equal(t, ClientThunderbird, normalizeClientKey(imapID("Thunderbird")))
115+
}
116+
func TestNormalizeClientKey_Unknown(t *testing.T) {
117+
assert.Equal(t, ClientUnknown, normalizeClientKey(imapID("")))
118+
assert.Equal(t, ClientUnknown, normalizeClientKey(imapID("Custom Client")))
119+
}

‎imap/connectionlimiter/limits.go‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package connectionlimiter
2+
3+
const (
4+
defaultClientLimit = 25
5+
defaultUnknownLimit = 10
6+
)
7+
8+
type Limits struct {
9+
// normalized client name with max open sessions.
10+
// If we want unlimited connections for a client set the limit to 0.
11+
PerClient map[Client]int
12+
13+
// max open sessions for unknown clients.
14+
// If we want unlimited connections for unknown clients set the limit to 0.
15+
UnknownLimit int
16+
}
17+
18+
func NewDefaultLimits() Limits {
19+
return Limits{
20+
PerClient: map[Client]int{
21+
ClientAppleMail: defaultClientLimit,
22+
ClientOutlook: defaultClientLimit,
23+
ClientThunderbird: defaultClientLimit,
24+
},
25+
UnknownLimit: defaultUnknownLimit,
26+
}
27+
}
28+
29+
func NewLimits(perClient map[Client]int, unknownLimit int) Limits {
30+
return Limits{
31+
PerClient: perClient,
32+
UnknownLimit: unknownLimit,
33+
}
34+
}

0 commit comments

Comments
 (0)