Skip to content

Commit

Permalink
transport: WrapConn, close connection when initialization fails
Browse files Browse the repository at this point in the history
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
  • Loading branch information
Kulezi committed Sep 23, 2022
1 parent 5369a73 commit a443dbf
Show file tree
Hide file tree
Showing 4 changed files with 66 additions and 13 deletions.
4 changes: 1 addition & 3 deletions transport/cluster.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"))
Expand Down
6 changes: 2 additions & 4 deletions transport/conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
63 changes: 63 additions & 0 deletions transport/conn_integration_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"context"
"fmt"
"math/rand"
"net"
"os/signal"
"strconv"
"sync"
Expand Down Expand Up @@ -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")
}
})
}
}
6 changes: 0 additions & 6 deletions transport/pool.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down Expand Up @@ -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 {
Expand Down

0 comments on commit a443dbf

Please sign in to comment.