Fix nil pointer dereference in TCPTransport.DoRaw

Signed-off-by: robert-cronin <robert.owen.cronin@gmail.com>

Author: robert-cronin

transport/tcp.go

@@ -1,6 +1,7 @@
 package transport
 
 import (
+	"context"
 	"net"
 	"net/http"
 )
@@ -14,19 +15,27 @@ func TCPTransport(host string, opts ...ConnectionOption) (*Transport, error) {
 		}
 	}
 
+	httpTransport := &http.Transport{
+		DialContext:     new(net.Dialer).DialContext,
+		TLSClientConfig: cfg.TLSConfig,
+	}
+
+	scheme := "http"
+	if cfg.TLSConfig != nil {
+		scheme = "https"
+	}
+
+	dial := func(ctx context.Context) (net.Conn, error) {
+		return httpTransport.DialContext(ctx, "tcp", host)
+	}
+
 	t := &Transport{
-		scheme: "http",
+		scheme: scheme,
 		host:   host,
 		c: &http.Client{
-			Transport: &http.Transport{
-				DialContext:     new(net.Dialer).DialContext,
-				TLSClientConfig: cfg.TLSConfig,
-			},
+			Transport: httpTransport,
 		},
-	}
-
-	if cfg.TLSConfig != nil {
-		t.scheme = "https"
+		dial: dial,
 	}
 
 	return t, nil

transport/tcp_test.go

@@ -2,6 +2,7 @@ package transport
 
 import (
 	"context"
+	"crypto/tls"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -35,3 +36,47 @@ func TestTCPTransport(t *testing.T) {
 	assert.NilError(t, err)
 	assert.Equal(t, string(buf), data)
 }
+
+func TestTCPTransportDoRaw(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		// for DoRaw test, we expect an upgrade request
+		if req.Header.Get("Connection") == "Upgrade" {
+			w.WriteHeader(http.StatusSwitchingProtocols)
+		} else {
+			w.WriteHeader(http.StatusNotFound)
+		}
+	}))
+	defer srv.Close()
+
+	ctx := context.Background()
+
+	u, err := url.Parse(srv.URL)
+	assert.NilError(t, err)
+
+	tr, err := TCPTransport(u.Host)
+	assert.NilError(t, err)
+
+	conn, err := tr.DoRaw(ctx, "POST", "/grpc", WithUpgrade("h2c"))
+	assert.Assert(t, conn != nil, "expected a connection but got nil")
+	assert.NilError(t, err)
+}
+
+func TestTCPTransportDoRawWithTLS(t *testing.T) {
+	// test with TLS configuration
+	tlsOpt := func(cfg *ConnectionConfig) error {
+		cfg.TLSConfig = &tls.Config{
+			InsecureSkipVerify: true,
+		}
+		return nil
+	}
+
+	tr, err := TCPTransport("localhost:2376", tlsOpt)
+	assert.NilError(t, err)
+
+	ctx := context.Background()
+
+	_, err = tr.DoRaw(ctx, "POST", "/grpc", WithUpgrade("h2c"))
+	// we expect this to fail with connection error since theres no server
+	// but it should not panic
+	assert.Assert(t, err != nil, "expected an error but got none")
+}