Close HTTP connection if DIAL_RSP contains error

Author: Jefftree

pkg/server/server.go

@@ -55,6 +55,9 @@ func (c *ProxyClientConnection) send(pkt *client.Packet) error {
 			_, err := c.HTTP.Write(pkt.GetData().Data)
 			return err
 		} else if pkt.Type == client.PacketType_DIAL_RSP {
+			if pkt.GetDialResponse().Error != "" {
+				return c.HTTP.Close()
+			}
 			return nil
 		} else {
 			return fmt.Errorf("attempt to send via unrecognized connection type %v", pkt.Type)
@@ -463,15 +466,22 @@ func (s *ProxyServer) serveRecvBackend(stream agent.AgentService_ConnectServer,
 			if client, ok := s.PendingDial.Get(resp.Random); !ok {
 				klog.Warning("<<< DialResp not recognized; dropped")
 			} else {
+				dialErr := false
+				if resp.Error != "" {
+					klog.Warningf("<<< DIAL_RSP contains error: %v", resp.Error)
+					dialErr = true
+				}
 				err := client.send(pkt)
 				s.PendingDial.Remove(resp.Random)
-				if resp.Error != "" {
-					klog.Warningf("<<< DIAL_RSP received error: %v", resp.Error)
-					break
-				} else if err != nil {
+				if err != nil {
 					klog.Warningf("<<< DIAL_RSP send to client stream error: %v", err)
+					dialErr = true
+				}
+				// Avoid adding the frontend if there was an error dialing the destination
+				if dialErr == true {
 					break
 				}
+
 				client.connectID = resp.ConnectID
 				client.agentID = agentID
 				s.addFrontend(agentID, resp.ConnectID, client)

tests/proxy_test.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"strings"
 	"testing"
 	"time"
 
@@ -92,8 +93,8 @@ func TestBasicProxy_GRPC(t *testing.T) {
 	}
 }
 
-func TestProxyHandleError_GRPC(t *testing.T) {
-	invalidServer := httptest.NewUnstartedServer(newEchoServer("hello"))
+func TestProxyHandleDialError_GRPC(t *testing.T) {
+	invalidServer := httptest.NewServer(newEchoServer("hello"))
 
 	stopCh := make(chan struct{})
 	defer close(stopCh)
@@ -121,8 +122,11 @@ func TestProxyHandleError_GRPC(t *testing.T) {
 		},
 	}
 
-	_, err = c.Get(invalidServer.URL)
-	if err == nil {
+	url := invalidServer.URL
+	invalidServer.Close()
+
+	_, err = c.Get(url)
+	if err == nil || !strings.Contains(err.Error(), "connection refused") {
 		t.Error("Expected error when destination is unreachable, did not receive error")
 	}
 }