Merge pull request #94 from Jefftree/dial_rsp

Don't add frontend on proxy server if DIAL_RSP contains error

Author: k8s-ci-robot

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)
@@ -468,16 +471,26 @@ 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 err != nil {
 					klog.Warningf("<<< DIAL_RSP send to client stream error: %v", err)
-				} else {
-					client.connectID = resp.ConnectID
-					client.agentID = agentID
-					s.addFrontend(agentID, resp.ConnectID, client)
-					close(client.connected)
+					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)
+				close(client.connected)
 			}
 
 		case client.PacketType_DATA:

tests/proxy_test.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"strings"
 	"testing"
 	"time"
 
@@ -92,6 +93,44 @@ func TestBasicProxy_GRPC(t *testing.T) {
 	}
 }
 
+func TestProxyHandleDialError_GRPC(t *testing.T) {
+	invalidServer := httptest.NewServer(newEchoServer("hello"))
+
+	stopCh := make(chan struct{})
+	defer close(stopCh)
+
+	proxy, cleanup, err := runGRPCProxyServer()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cleanup()
+
+	runAgent(proxy.agent, stopCh)
+
+	// Wait for agent to register on proxy server
+	time.Sleep(time.Second)
+
+	// run test client
+	tunnel, err := client.CreateGrpcTunnel(proxy.front, grpc.WithInsecure())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	c := &http.Client{
+		Transport: &http.Transport{
+			Dial: tunnel.Dial,
+		},
+	}
+
+	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")
+	}
+}
+
 func TestProxy_LargeResponse(t *testing.T) {
 	length := 1 << 20 // 1M
 	chunks := 10