teleport/integration/terminal_test.go
rosstimothy b17f687f5f
Track active ssh sessions established via the web (#20231)
* Track active ssh sessions established via the web

Adds open session tracking in `web.Handler` and introduces a new
 `web.Server` to prevent closing the web `http.Server` during
a graceful shutdown if there are still open sessions. From the
`http.Server/Shutdown` docs:

> Shutdown does not attempt to close nor wait for hijacked
> connections such as WebSockets. The caller of Shutdown should
> separately notify such long-lived connections of shutdown and wait
> for them to close, if desired.

Because the web sessions use WebSockets, `Shutdown` will not
wait for them to be idle before terminating. This used to work due
to the fact that web sessions dialed the Proxy SSH Port  to establish the
SSH session. Which meant that all accounting of active user connections
by `sshutils/server.go` was accurate. But now that sessions are
established directly via the `proxy.Router` the web sessions were
unaccounted for.

The `web.Server` now mimics the behavior of `sshutils/server.go` for
Shutdown. Only once all the active web sessions have been terminated
will the underlying http.Server be shutdown.

Fixes #20227
2023-01-27 14:38:39 +00:00

114 lines
2.3 KiB
Go

/*
Copyright 2021 Gravitational, Inc.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package integration
import (
"bytes"
"io"
"strings"
"sync"
"time"
)
// Terminal emulates stdin+stdout for integration testing
type Terminal struct {
typed chan byte
close chan struct{}
mu *sync.Mutex
written *bytes.Buffer
}
func NewTerminal(capacity int) *Terminal {
return &Terminal{
typed: make(chan byte, capacity),
mu: &sync.Mutex{},
written: bytes.NewBuffer(nil),
close: make(chan struct{}),
}
}
func (t *Terminal) Type(data string) {
for _, b := range []byte(data) {
t.typed <- b
}
}
// Output returns a number of first 'limit' bytes printed into this fake terminal
func (t *Terminal) Output(limit int) string {
t.mu.Lock()
defer t.mu.Unlock()
buff := t.written.Bytes()
if len(buff) > limit {
buff = buff[:limit]
}
// clean up white space for easier comparison:
return strings.TrimSpace(string(buff))
}
// AllOutput returns the entire recorded output from the fake terminal
func (t *Terminal) AllOutput() string {
t.mu.Lock()
defer t.mu.Unlock()
return strings.TrimSpace(t.written.String())
}
func (t *Terminal) Write(data []byte) (n int, err error) {
t.mu.Lock()
defer t.mu.Unlock()
return t.written.Write(data)
}
func (t *Terminal) Read(p []byte) (n int, err error) {
for n = 0; n < len(p); n++ {
select {
case p[n] = <-t.typed:
case <-t.close:
return n, io.EOF
}
if p[n] == '\r' {
break
}
if p[n] == '\a' { // 'alert' used for debugging, means 'pause for 1 second'
select {
case <-time.After(time.Second):
n--
case <-t.close:
return n, io.EOF
}
}
select {
case <-time.After(time.Millisecond * 10):
case <-t.close:
return n, io.EOF
}
}
return n, nil
}
func (t *Terminal) Close() error {
select {
case <-t.close:
default:
close(t.close)
}
return nil
}