Release resources in client Transport.Dial on error.

Make a stack of cleanup functions to run (as with defer), but clear the
stack before returning if no error occurs.

Uselessly pushing the stream.Close() cleanup just before clearing the
stack is an intentional safeguard, for in case additional operations are
added before the return in the future.

Fixes #40042.

1 files changed, 18 insertions, 1 deletions

diff --git a/client/lib/snowflake.go b/client/lib/snowflake.go
index 6e87b81..f643c0a 100644
--- a/client/lib/snowflake.go
+++ b/client/lib/snowflake.go
@@ -74,11 +74,21 @@ func NewSnowflakeClient(brokerURL, frontDomain string, iceAddresses []string, ke
 // Create a new Snowflake connection. Starts the collection of snowflakes and returns a
 // smux Stream.
 func (t *Transport) Dial() (net.Conn, error) {
+	// Cleanup functions to run before returning, in case of an error.
+	var cleanup []func()
+	defer func() {
+		// Run cleanup in reverse order, as defer does.
+		for i := len(cleanup) - 1; i >= 0; i-- {
+			cleanup[i]()
+		}
+	}()
+
 	// Prepare to collect remote WebRTC peers.
 	snowflakes, err := NewPeers(t.dialer)
 	if err != nil {
 		return nil, err
 	}
+	cleanup = append(cleanup, func() { snowflakes.End() })
 
 	// Use a real logger to periodically output how much traffic is happening.
 	snowflakes.BytesLogger = NewBytesSyncLogger()
@@ -92,15 +102,22 @@ func (t *Transport) Dial() (net.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
+	cleanup = append(cleanup, func() {
+		pconn.Close()
+		sess.Close()
+	})
 
 	// On the smux session we overlay a stream.
 	stream, err := sess.OpenStream()
 	if err != nil {
 		return nil, err
 	}
-
 	// Begin exchanging data.
 	log.Printf("---- SnowflakeConn: begin stream %v ---", stream.ID())
+	cleanup = append(cleanup, func() { stream.Close() })
+
+	// All good, clear the cleanup list.
+	cleanup = nil
 	return &SnowflakeConn{Stream: stream, sess: sess, pconn: pconn, snowflakes: snowflakes}, nil
 }