transport: WrapConn, close connection when initialization fails

WrapConn didn't close connReader and connWriter loops, leaving this
responsibility to its callers, potentially leading to goroutine leaks,
this is mitigated by closing those loops and never returning a non-nil
Conn pointer when WrapConn fails.

Fixes #289

Author: Kulezi

transport/cluster.go

@@ -167,14 +167,12 @@ func (c *Cluster) NewControl(ctx context.Context) (*Conn, error) {
 			if err := conn.RegisterEventHandler(ctx, c.handleEvent, c.handledEvents...); err == nil {
 				return conn, nil
 			} else {
+				conn.Close()
 				errs = append(errs, fmt.Sprintf("%s failed to register for events: %s", conn, err))
 			}
 		} else {
 			errs = append(errs, fmt.Sprintf("%s failed to connect: %s", addr, err))
 		}
-		if conn != nil {
-			conn.Close()
-		}
 	}
 
 	return nil, fmt.Errorf("couldn't open control connection to any known host:\n%s", strings.Join(errs, "\n"))

transport/conn.go

@@ -385,9 +385,6 @@ func OpenShardConn(ctx context.Context, addr string, si ShardInfo, cfg ConnConfi
 		conn, err := OpenLocalPortConn(ctx, addr, it(), cfg)
 		if err != nil {
 			cfg.Logger.Infof("%s dial error: %s (try %d/%d)", addr, err, i, maxTries)
-			if conn != nil {
-				conn.Close()
-			}
 			continue
 		}
 		return conn, nil
@@ -506,7 +503,8 @@ func WrapConn(ctx context.Context, conn net.Conn, cfg ConnConfig) (*Conn, error)
 	go c.r.loop(ctx)
 
 	if err := c.init(ctx); err != nil {
-		return c, err
+		c.Close()
+		return nil, err
 	}
 
 	return c, nil

transport/conn_integration_test.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"fmt"
 	"math/rand"
+	"net"
 	"os/signal"
 	"strconv"
 	"sync"
@@ -278,3 +279,65 @@ func testCompression(ctx context.Context, t *testing.T, c frame.Compression, toS
 		}
 	}
 }
+
+func TestConnectedToNonCqlServer(t *testing.T) {
+	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGABRT, syscall.SIGTERM)
+	defer cancel()
+
+	t.Logf("%+v", testingConnConfig)
+	testCases := []struct {
+		name     string
+		response []byte
+	}{
+		{
+			name:     "non-cql response",
+			response: []byte("0"),
+		},
+		{
+			name: "non supported cql response",
+			response: func() []byte {
+				var buf frame.Buffer
+				frame := frame.Header{
+					Version: frame.CQLv4,
+					OpCode:  frame.OpReady,
+				}
+
+				frame.WriteTo(&buf)
+				return buf.Bytes()
+			}(),
+		},
+	}
+
+	for i := 0; i < len(testCases); i++ {
+		tc := testCases[i]
+		t.Run(tc.name, func(t *testing.T) {
+			server, err := net.Listen("tcp", "127.0.0.1:")
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer server.Close()
+			go func() {
+				conn, err := server.Accept()
+				if err != nil {
+					t.Log(err)
+					t.Fail()
+					return
+				}
+				go func(conn net.Conn) {
+					defer conn.Close()
+					conn.Write(tc.response)
+				}(conn)
+			}()
+
+			addr := server.Addr().String()
+			conn, err := OpenConn(ctx, addr, nil, testingConnConfig)
+			if err == nil {
+				t.Fatal("connecting to non-cql server should fail")
+			}
+			t.Log(err)
+			if conn != nil {
+				t.Fatal("connecting to non-cql server should return a nil-conn")
+			}
+		})
+	}
+}

transport/pool.go

@@ -141,9 +141,6 @@ func (r *PoolRefiller) init(ctx context.Context, host string) error {
 	conn, err := OpenConn(ctx, host, nil, r.cfg)
 	span.stop()
 	if err != nil {
-		if conn != nil {
-			conn.Close()
-		}
 		return err
 	}
 
@@ -245,9 +242,6 @@ func (r *PoolRefiller) fill(ctx context.Context) {
 			if r.pool.connObs != nil {
 				r.pool.connObs.OnConnect(ConnectEvent{ConnEvent: ConnEvent{Addr: r.addr, Shard: si.Shard}, span: span, Err: err})
 			}
-			if conn != nil {
-				conn.Close()
-			}
 			continue
 		}
 		if r.pool.connObs != nil {