cloudflared/management/service_test.go
Devin Carr 38cd455e4d TUN-7373: Streaming logs override for same actor
To help accommodate web browser interactions with websockets, when a
streaming logs session is requested for the same actor while already
serving a session for that user in a separate request, the original
request will be closed and the new request start streaming logs
instead. This should help with rogue sessions holding on for too long
with no client on the other side (before idle timeout or connection
close).
2023-04-21 11:54:37 -07:00

127 lines
3.2 KiB
Go

package management
import (
"context"
"io"
"testing"
"time"
"github.com/rs/zerolog"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"nhooyr.io/websocket"
"github.com/cloudflare/cloudflared/internal/test"
)
var (
noopLogger = zerolog.New(io.Discard)
)
func TestReadEventsLoop(t *testing.T) {
sentEvent := EventStartStreaming{
ClientEvent: ClientEvent{Type: StartStreaming},
}
client, server := test.WSPipe(nil, nil)
client.CloseRead(context.Background())
defer func() {
client.Close(websocket.StatusInternalError, "")
}()
go func() {
err := WriteEvent(client, context.Background(), &sentEvent)
require.NoError(t, err)
}()
m := ManagementService{
log: &noopLogger,
}
events := make(chan *ClientEvent)
go m.readEvents(server, context.Background(), events)
event := <-events
require.Equal(t, sentEvent.Type, event.Type)
server.Close(websocket.StatusInternalError, "")
}
func TestReadEventsLoop_ContextCancelled(t *testing.T) {
client, server := test.WSPipe(nil, nil)
ctx, cancel := context.WithCancel(context.Background())
client.CloseRead(ctx)
defer func() {
client.Close(websocket.StatusInternalError, "")
}()
m := ManagementService{
log: &noopLogger,
}
events := make(chan *ClientEvent)
go func() {
time.Sleep(time.Second)
cancel()
}()
// Want to make sure this function returns when context is cancelled
m.readEvents(server, ctx, events)
server.Close(websocket.StatusInternalError, "")
}
func TestCanStartStream_NoSessions(t *testing.T) {
m := ManagementService{
log: &noopLogger,
logger: &Logger{
Log: &noopLogger,
},
}
_, cancel := context.WithCancel(context.Background())
session := newSession(0, actor{}, cancel)
assert.True(t, m.canStartStream(session))
}
func TestCanStartStream_ExistingSessionDifferentActor(t *testing.T) {
m := ManagementService{
log: &noopLogger,
logger: &Logger{
Log: &noopLogger,
},
}
_, cancel := context.WithCancel(context.Background())
session1 := newSession(0, actor{ID: "test"}, cancel)
assert.True(t, m.canStartStream(session1))
m.logger.Listen(session1)
assert.True(t, session1.Active())
// Try another session
session2 := newSession(0, actor{ID: "test2"}, cancel)
assert.Equal(t, 1, m.logger.ActiveSessions())
assert.False(t, m.canStartStream(session2))
// Close session1
m.logger.Remove(session1)
assert.True(t, session1.Active()) // Remove doesn't stop a session
session1.Stop()
assert.False(t, session1.Active())
assert.Equal(t, 0, m.logger.ActiveSessions())
// Try session2 again
assert.True(t, m.canStartStream(session2))
}
func TestCanStartStream_ExistingSessionSameActor(t *testing.T) {
m := ManagementService{
log: &noopLogger,
logger: &Logger{
Log: &noopLogger,
},
}
actor := actor{ID: "test"}
_, cancel := context.WithCancel(context.Background())
session1 := newSession(0, actor, cancel)
assert.True(t, m.canStartStream(session1))
m.logger.Listen(session1)
assert.True(t, session1.Active())
// Try another session
session2 := newSession(0, actor, cancel)
assert.Equal(t, 1, m.logger.ActiveSessions())
assert.True(t, m.canStartStream(session2))
// session1 is removed and stopped
assert.Equal(t, 0, m.logger.ActiveSessions())
assert.False(t, session1.Active())
}