Clean up some lib/utils code (#59790)

* Clean up some lib/utils code

Rename CloserConn to WaitConn. This utility had the ability to tie
multiple io.Closers to the lifecycle of a net.Conn, but we never
used it for that purpose, and instead use it only as a way to wait
for a connection to be closed.

Remove the unnecessary utils.Stater interface. The only usage of
this interface already happened behind a type assertion, so there
was no need for an interface.

Un-export TrackingReader and TrackingWriter since they are
implementation details of TrackingConn.

* Move SyncBuffer to integration test package

integration/buf_test.go

@@ -16,15 +16,15 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-package utils
+package integration
 
 import (
 	"bytes"
 	"io"
 )
 
-// NewSyncBuffer returns new in memory buffer
-func NewSyncBuffer() *SyncBuffer {
+// newSyncBuffer returns new in memory buffer
+func newSyncBuffer() *syncBuffer {
 	reader, writer := io.Pipe()
 	buf := &bytes.Buffer{}
 	copyDone := make(chan struct{})
@@ -32,7 +32,7 @@ func NewSyncBuffer() *SyncBuffer {
 		defer close(copyDone)
 		io.Copy(buf, reader)
 	}()
-	return &SyncBuffer{
+	return &syncBuffer{
 		reader:   reader,
 		writer:   writer,
 		buf:      buf,
@@ -40,9 +40,9 @@ func NewSyncBuffer() *SyncBuffer {
 	}
 }
 
-// SyncBuffer is in memory bytes buffer that is
+// syncBuffer is in memory bytes buffer that is
 // safe for concurrent writes
-type SyncBuffer struct {
+type syncBuffer struct {
 	reader *io.PipeReader
 	writer *io.PipeWriter
 	buf    *bytes.Buffer
@@ -51,26 +51,26 @@ type SyncBuffer struct {
 	copyDone chan struct{}
 }
 
-func (b *SyncBuffer) Write(data []byte) (n int, err error) {
+func (b *syncBuffer) Write(data []byte) (n int, err error) {
 	return b.writer.Write(data)
 }
 
 // String returns contents of the buffer
 // after this call, all writes will fail
-func (b *SyncBuffer) String() string {
+func (b *syncBuffer) String() string {
 	b.Close()
 	return b.buf.String()
 }
 
 // Bytes returns contents of the buffer
 // after this call, all writes will fail
-func (b *SyncBuffer) Bytes() []byte {
+func (b *syncBuffer) Bytes() []byte {
 	b.Close()
 	return b.buf.Bytes()
 }
 
 // Close closes reads and writes on the buffer
-func (b *SyncBuffer) Close() error {
+func (b *syncBuffer) Close() error {
 	err := b.reader.Close()
 	err2 := b.writer.Close()
 

integration/hostuser_test.go

@@ -52,7 +52,6 @@ import (
 	"github.com/gravitational/teleport/lib/service/servicecfg"
 	"github.com/gravitational/teleport/lib/services/local"
 	"github.com/gravitational/teleport/lib/srv"
-	"github.com/gravitational/teleport/lib/utils"
 	"github.com/gravitational/teleport/lib/utils/host"
 	"github.com/gravitational/teleport/lib/utils/log/logtest"
 	"github.com/gravitational/teleport/lib/utils/testutils"
@@ -701,7 +700,7 @@ func TestRootLoginAsHostUser(t *testing.T) {
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
 			stdin := bytes.NewBufferString(tc.stdinText)
-			stdout := utils.NewSyncBuffer()
+			stdout := newSyncBuffer()
 			t.Cleanup(func() {
 				require.NoError(t, stdout.Close())
 			})

integration/integration_test.go

@@ -674,7 +674,7 @@ func testInteroperability(t *testing.T, suite *integrationTestSuite) {
 
 			// hook up stdin and stdout to a buffer for reading and writing
 			inbuf := bytes.NewReader([]byte(tt.inStdin))
-			outbuf := utils.NewSyncBuffer()
+			outbuf := newSyncBuffer()
 			cl.Stdin = inbuf
 			cl.Stdout = outbuf
 			cl.Stderr = outbuf

lib/srv/app/connections_handler.go

@@ -326,12 +326,12 @@ func (c *ConnectionsHandler) expireSessions() {
 func (c *ConnectionsHandler) HandleConnection(conn net.Conn) {
 	ctx := context.Background()
 
-	// Wrap conn in a CloserConn to detect when it is closed.
+	// Wrap conn to detect when it is closed.
 	// Returning early will close conn before it has been serviced.
 	// httpServer will initiate the close call.
-	closerConn := utils.NewCloserConn(conn)
+	waitConn := utils.NewWaitConn(conn)
 
-	cleanup, err := c.handleConnection(closerConn)
+	cleanup, err := c.handleConnection(waitConn)
 	// Make sure that the cleanup function is run
 	if cleanup != nil {
 		defer cleanup()
@@ -348,7 +348,7 @@ func (c *ConnectionsHandler) HandleConnection(conn net.Conn) {
 	}
 
 	// Wait for connection to close.
-	closerConn.Wait()
+	waitConn.Wait()
 }
 
 // serveSession finds the app session and forwards the request.

lib/srv/ctx.go

@@ -868,7 +868,7 @@ func (c *ServerContext) takeClosers() []io.Closer {
 
 // When the ServerContext (connection) is closed, emit "session.data" event
 // containing how much data was transmitted and received over the net.Conn.
-func (c *ServerContext) reportStats(conn utils.Stater) {
+func (c *ServerContext) reportStats(conn *utils.TrackingConn) {
 	// We may not want to record session data for this connection context, e.g. if this is
 	// for a networking subprocess tied to a shell process.
 	if c.SessionRecordingConfig.GetMode() == types.RecordOff {

lib/srv/mcp/http.go

@@ -49,20 +49,15 @@ func (s *Server) serveHTTPConn(ctx context.Context, conn net.Conn, handler http.
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	waitConn := utils.NewCloserConn(conn)
-	listener := listenerutils.NewSingleUseListener(waitConn)
-	go func() {
-		// Make sure connection is closed when ctx is canceled.
-		<-ctx.Done()
-		waitConn.Close()
-	}()
+	waitConn := utils.NewWaitConn(conn)
+	context.AfterFunc(ctx, func() { waitConn.Close() })
 
 	httpServer := &http.Server{
-		Handler: handler,
-		BaseContext: func(net.Listener) context.Context {
-			return ctx
-		},
+		Handler:     handler,
+		BaseContext: func(net.Listener) context.Context { return ctx },
 	}
+
+	listener := listenerutils.NewSingleUseListener(waitConn)
 	if err := httpServer.Serve(listener); err != nil && !utils.IsOKNetworkError(err) {
 		return trace.Wrap(err)
 	}

lib/utils/conn.go

@@ -19,76 +19,55 @@
 package utils
 
 import (
-	"context"
 	"io"
 	"net"
+	"sync"
 	"sync/atomic"
 
 	"github.com/gravitational/trace"
 )
 
-// NewCloserConn returns new connection wrapper that
-// when closed will also close passed closers
-func NewCloserConn(conn net.Conn, closers ...io.Closer) *CloserConn {
-	c := &CloserConn{
-		Conn:    conn,
-		closers: closers,
+// NewWaitConn returns new connection wrapper that
+// provides the ability to wait for the connection
+// to be closed.
+func NewWaitConn(conn net.Conn) *WaitConn {
+	wc := &WaitConn{
+		Conn:   conn,
+		closed: make(chan struct{}),
 	}
-	c.ctx, c.ctxCancel = context.WithCancel(context.Background())
-	return c
+	wc.close = sync.OnceValue(func() error {
+		err := wc.Conn.Close()
+		close(wc.closed)
+		return trace.Wrap(err)
+	})
+	return wc
 }
 
-// CloserConn wraps connection and attaches additional closers to it
-type CloserConn struct {
+// WaitConn wraps a connection and provides the ability to wait for the
+// connection to be closed.
+type WaitConn struct {
 	net.Conn
-	closers   []io.Closer
-	ctx       context.Context
-	ctxCancel context.CancelFunc
+
+	close  func() error
+	closed chan struct{}
 }
 
-// AddCloser adds any closer in ctx that will be called
-// whenever server closes session channel
-func (c *CloserConn) AddCloser(closer io.Closer) {
-	c.closers = append(c.closers, closer)
+// Close closes the connection.
+func (c *WaitConn) Close() error {
+	return c.close()
 }
 
-// Close connection, all closers, and cancel context.
-func (c *CloserConn) Close() error {
-	var errors []error
-	for _, closer := range c.closers {
-		errors = append(errors, closer.Close())
-	}
-	errors = append(errors, c.Conn.Close())
-	c.ctxCancel()
-	return trace.NewAggregate(errors...)
-}
-
-// Context returns a context that is canceled once the connection is closed.
-func (c *CloserConn) Context() context.Context {
-	return c.ctx
-}
-
-// Wait for connection to close.
-func (c *CloserConn) Wait() {
-	<-c.ctx.Done()
-}
-
-// Stater is extension interface of the net.Conn for implementations that
-// track connection statistics.
-type Stater interface {
-	// Stat returns TX, RX data.
-	Stat() (uint64, uint64)
-}
+func (c *WaitConn) Done() <-chan struct{} { return c.closed }
+func (c *WaitConn) Wait()                 { <-c.Done() }
 
 // TrackingConn is a net.Conn that keeps track of how much data was transmitted
 // (TX) and received (RX) over the net.Conn. A maximum of about 18446
 // petabytes can be kept track of for TX and RX before it rolls over.
 // See https://golang.org/ref/spec#Numeric_types for more details.
 type TrackingConn struct {
-	// net.Conn is the underlying net.Conn.
 	net.Conn
-	r *TrackingReader
-	w *TrackingWriter
+	r *trackingReader
+	w *trackingWriter
 }
 
 // NewTrackingConn returns a net.Conn that can keep track of how much data was
@@ -96,13 +75,13 @@ type TrackingConn struct {
 func NewTrackingConn(conn net.Conn) *TrackingConn {
 	return &TrackingConn{
 		Conn: conn,
-		r:    NewTrackingReader(conn),
-		w:    NewTrackingWriter(conn),
+		r:    &trackingReader{r: conn},
+		w:    &trackingWriter{w: conn},
 	}
 }
 
 // Stat returns the transmitted (TX) and received (RX) bytes over the net.Conn.
-func (s *TrackingConn) Stat() (uint64, uint64) {
+func (s *TrackingConn) Stat() (written, read uint64) {
 	return s.w.Count(), s.r.Count()
 }
 
@@ -114,24 +93,19 @@ func (s *TrackingConn) Write(b []byte) (n int, err error) {
 	return s.w.Write(b)
 }
 
-// TrackingReader is an io.Reader that counts the total number of bytes read.
+// trackingReader is an io.Reader that counts the total number of bytes read.
 // It's thread-safe if the underlying io.Reader is thread-safe.
-type TrackingReader struct {
+type trackingReader struct {
 	r     io.Reader
 	count uint64
 }
 
-// NewTrackingReader creates a TrackingReader around r.
-func NewTrackingReader(r io.Reader) *TrackingReader {
-	return &TrackingReader{r: r}
-}
-
 // Count returns the total number of bytes read so far.
-func (r *TrackingReader) Count() uint64 {
+func (r *trackingReader) Count() uint64 {
 	return atomic.LoadUint64(&r.count)
 }
 
-func (r *TrackingReader) Read(b []byte) (int, error) {
+func (r *trackingReader) Read(b []byte) (int, error) {
 	n, err := r.r.Read(b)
 	atomic.AddUint64(&r.count, uint64(n))
 
@@ -141,25 +115,20 @@ func (r *TrackingReader) Read(b []byte) (int, error) {
 	return n, err
 }
 
-// TrackingWriter is an io.Writer that counts the total number of bytes
+// trackingWriter is an io.Writer that counts the total number of bytes
 // written.
 // It's thread-safe if the underlying io.Writer is thread-safe.
-type TrackingWriter struct {
+type trackingWriter struct {
 	count uint64 // intentionally placed first to ensure 64-bit alignment
 	w     io.Writer
 }
 
-// NewTrackingWriter creates a TrackingWriter around w.
-func NewTrackingWriter(w io.Writer) *TrackingWriter {
-	return &TrackingWriter{w: w}
-}
-
 // Count returns the total number of bytes written so far.
-func (w *TrackingWriter) Count() uint64 {
+func (w *trackingWriter) Count() uint64 {
 	return atomic.LoadUint64(&w.count)
 }
 
-func (w *TrackingWriter) Write(b []byte) (int, error) {
+func (w *trackingWriter) Write(b []byte) (int, error) {
 	n, err := w.w.Write(b)
 	atomic.AddUint64(&w.count, uint64(n))
 	return n, trace.Wrap(err)

lib/utils/conn_test.go

@@ -31,7 +31,7 @@ func TestTrackingReaderEOF(t *testing.T) {
 	reader := bytes.NewReader([]byte{})
 
 	// Wrap the reader in a TrackingReader.
-	tr := NewTrackingReader(reader)
+	tr := &trackingReader{r: reader}
 
 	// Make sure it returns an EOF and not a wrapped exception.
 	buf := make([]byte, 64)