Add support for graceful shutdown

Fixes #1014

Signed-off-by: Nicola Murino <nicola.murino@gmail.com>
This commit is contained in:
Nicola Murino
2022-10-22 11:56:41 +02:00
parent 87045284cc
commit db0e58ae7e
32 changed files with 371 additions and 124 deletions

View File

@@ -497,10 +497,10 @@ func TestWhitelist(t *testing.T) {
err = Initialize(Config, 0)
assert.NoError(t, err)
assert.True(t, Connections.IsNewConnectionAllowed("172.18.1.1"))
assert.False(t, Connections.IsNewConnectionAllowed("172.18.1.3"))
assert.True(t, Connections.IsNewConnectionAllowed("10.8.7.3"))
assert.False(t, Connections.IsNewConnectionAllowed("10.8.8.2"))
assert.NoError(t, Connections.IsNewConnectionAllowed("172.18.1.1"))
assert.Error(t, Connections.IsNewConnectionAllowed("172.18.1.3"))
assert.NoError(t, Connections.IsNewConnectionAllowed("10.8.7.3"))
assert.Error(t, Connections.IsNewConnectionAllowed("10.8.8.2"))
wl.IPAddresses = append(wl.IPAddresses, "172.18.1.3")
wl.CIDRNetworks = append(wl.CIDRNetworks, "10.8.8.0/24")
@@ -508,14 +508,14 @@ func TestWhitelist(t *testing.T) {
assert.NoError(t, err)
err = os.WriteFile(wlFile, data, 0664)
assert.NoError(t, err)
assert.False(t, Connections.IsNewConnectionAllowed("10.8.8.3"))
assert.Error(t, Connections.IsNewConnectionAllowed("10.8.8.3"))
err = Reload()
assert.NoError(t, err)
assert.True(t, Connections.IsNewConnectionAllowed("10.8.8.3"))
assert.True(t, Connections.IsNewConnectionAllowed("172.18.1.3"))
assert.True(t, Connections.IsNewConnectionAllowed("172.18.1.2"))
assert.False(t, Connections.IsNewConnectionAllowed("172.18.1.12"))
assert.NoError(t, Connections.IsNewConnectionAllowed("10.8.8.3"))
assert.NoError(t, Connections.IsNewConnectionAllowed("172.18.1.3"))
assert.NoError(t, Connections.IsNewConnectionAllowed("172.18.1.2"))
assert.Error(t, Connections.IsNewConnectionAllowed("172.18.1.12"))
Config = configCopy
}
@@ -550,12 +550,12 @@ func TestMaxConnections(t *testing.T) {
Config.MaxPerHostConnections = 0
ipAddr := "192.168.7.8"
assert.True(t, Connections.IsNewConnectionAllowed(ipAddr))
assert.NoError(t, Connections.IsNewConnectionAllowed(ipAddr))
Config.MaxTotalConnections = 1
Config.MaxPerHostConnections = perHost
assert.True(t, Connections.IsNewConnectionAllowed(ipAddr))
assert.NoError(t, Connections.IsNewConnectionAllowed(ipAddr))
c := NewBaseConnection("id", ProtocolSFTP, "", "", dataprovider.User{})
fakeConn := &fakeConnection{
BaseConnection: c,
@@ -563,18 +563,18 @@ func TestMaxConnections(t *testing.T) {
err := Connections.Add(fakeConn)
assert.NoError(t, err)
assert.Len(t, Connections.GetStats(), 1)
assert.False(t, Connections.IsNewConnectionAllowed(ipAddr))
assert.Error(t, Connections.IsNewConnectionAllowed(ipAddr))
res := Connections.Close(fakeConn.GetID())
assert.True(t, res)
assert.Eventually(t, func() bool { return len(Connections.GetStats()) == 0 }, 300*time.Millisecond, 50*time.Millisecond)
assert.True(t, Connections.IsNewConnectionAllowed(ipAddr))
assert.NoError(t, Connections.IsNewConnectionAllowed(ipAddr))
Connections.AddClientConnection(ipAddr)
Connections.AddClientConnection(ipAddr)
assert.False(t, Connections.IsNewConnectionAllowed(ipAddr))
assert.Error(t, Connections.IsNewConnectionAllowed(ipAddr))
Connections.RemoveClientConnection(ipAddr)
assert.True(t, Connections.IsNewConnectionAllowed(ipAddr))
assert.NoError(t, Connections.IsNewConnectionAllowed(ipAddr))
Connections.RemoveClientConnection(ipAddr)
Config.MaxTotalConnections = oldValue
@@ -587,13 +587,13 @@ func TestMaxConnectionPerHost(t *testing.T) {
ipAddr := "192.168.9.9"
Connections.AddClientConnection(ipAddr)
assert.True(t, Connections.IsNewConnectionAllowed(ipAddr))
assert.NoError(t, Connections.IsNewConnectionAllowed(ipAddr))
Connections.AddClientConnection(ipAddr)
assert.True(t, Connections.IsNewConnectionAllowed(ipAddr))
assert.NoError(t, Connections.IsNewConnectionAllowed(ipAddr))
Connections.AddClientConnection(ipAddr)
assert.False(t, Connections.IsNewConnectionAllowed(ipAddr))
assert.Error(t, Connections.IsNewConnectionAllowed(ipAddr))
assert.Equal(t, int32(3), Connections.GetClientConnections())
Connections.RemoveClientConnection(ipAddr)
@@ -697,7 +697,7 @@ func TestCloseConnection(t *testing.T) {
fakeConn := &fakeConnection{
BaseConnection: c,
}
assert.True(t, Connections.IsNewConnectionAllowed("127.0.0.1"))
assert.NoError(t, Connections.IsNewConnectionAllowed("127.0.0.1"))
err := Connections.Add(fakeConn)
assert.NoError(t, err)
assert.Len(t, Connections.GetStats(), 1)