mirror of
https://github.com/gravitational/teleport.git
synced 2026-09-01 16:03:55 +08:00
Kenneth
a57a6aa174
* Make rejected accpets retryable * Added better test for server behavior on limiter accept errors * Test for presence rather than absence * Create callback so listener does not delay on failed IP limit Accept() --------- Co-authored-by: Maja Cieślak <maja.cieslak@goteleport.com>
89 lines
2.1 KiB
Go
89 lines
2.1 KiB
Go
/*
|
|
* Teleport
|
|
* Copyright (C) 2024 Gravitational, Inc.
|
|
*
|
|
* This program is free software: you can redistribute it and/or modify
|
|
* it under the terms of the GNU Affero General Public License as published by
|
|
* the Free Software Foundation, either version 3 of the License, or
|
|
* (at your option) any later version.
|
|
*
|
|
* This program 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 Affero General Public License for more details.
|
|
*
|
|
* You should have received a copy of the GNU Affero General Public License
|
|
* along with this program. If not, see <http://www.gnu.org/licenses/>.
|
|
*/
|
|
|
|
package limiter_test
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/gravitational/teleport/lib/limiter"
|
|
)
|
|
|
|
func TestConnectionsLimiter(t *testing.T) {
|
|
l := limiter.NewConnectionsLimiter(0)
|
|
|
|
for range 10 {
|
|
require.NoError(t, l.AcquireConnection("token1"))
|
|
}
|
|
for range 5 {
|
|
require.NoError(t, l.AcquireConnection("token2"))
|
|
}
|
|
|
|
for range 10 {
|
|
l.ReleaseConnection("token1")
|
|
}
|
|
for range 5 {
|
|
l.ReleaseConnection("token2")
|
|
}
|
|
|
|
l = limiter.NewConnectionsLimiter(5)
|
|
|
|
for range 5 {
|
|
require.NoError(t, l.AcquireConnection("token1"))
|
|
}
|
|
|
|
for range 5 {
|
|
require.NoError(t, l.AcquireConnection("token2"))
|
|
}
|
|
for range 5 {
|
|
require.Error(t, l.AcquireConnection("token2"))
|
|
}
|
|
|
|
for range 10 {
|
|
l.ReleaseConnection("token1")
|
|
require.NoError(t, l.AcquireConnection("token1"))
|
|
}
|
|
|
|
for range 5 {
|
|
l.ReleaseConnection("token2")
|
|
}
|
|
for range 5 {
|
|
require.NoError(t, l.AcquireConnection("token2"))
|
|
}
|
|
}
|
|
|
|
func TestConnectionsLimiter_GetNumConnection(t *testing.T) {
|
|
l := limiter.NewConnectionsLimiter(2)
|
|
|
|
numConnections, err := l.GetNumConnection("conn1")
|
|
require.NoError(t, err)
|
|
require.Zero(t, numConnections)
|
|
|
|
require.NoError(t, l.AcquireConnection("conn1"))
|
|
numConnections, err = l.GetNumConnection("conn1")
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(1), numConnections)
|
|
|
|
l.ReleaseConnection("conn1")
|
|
numConnections, err = l.GetNumConnection("conn1")
|
|
require.NoError(t, err)
|
|
require.Zero(t, numConnections)
|
|
}
|