fix: client should return instantly when no backend

pkg/server/server.go

@@ -425,8 +425,12 @@ func (s *ProxyServer) serveRecvFrontend(stream client.ProxyService_ProxyServer,
 				klog.ErrorS(err, "Failed to get a backend", "serverID", s.serverID)
 
 				resp := &client.Packet{
-					Type:    client.PacketType_DIAL_RSP,
-					Payload: &client.Packet_DialResponse{DialResponse: &client.DialResponse{Error: err.Error()}},
+					Type: client.PacketType_DIAL_RSP,
+					Payload: &client.Packet_DialResponse{
+						DialResponse: &client.DialResponse{
+							Random: pkt.GetDialRequest().Random,
+							Error:  err.Error(),
+						}},
 				}
 				if err := stream.Send(resp); err != nil {
 					klog.V(5).Infoln("Failed to send DIAL_RSP for no backend", "error", err, "serverID", s.serverID)
@@ -459,8 +463,10 @@ func (s *ProxyServer) serveRecvFrontend(stream client.ProxyService_ProxyServer,
 			if err := backend.Send(pkt); err != nil {
 				// TODO: retry with other backends connecting to this agent.
 				klog.ErrorS(err, "CLOSE_REQ to Backend failed", "serverID", s.serverID, "connectionID", connID)
+			} else {
+				klog.V(5).Infoln("CLOSE_REQ sent to backend", "serverID", s.serverID, "connectionID", connID)
+				return
 			}
-			klog.V(5).Infoln("CLOSE_REQ sent to backend", "serverID", s.serverID, "connectionID", connID)
 
 		case client.PacketType_DIAL_CLS:
 			random := pkt.GetCloseDial().Random

pkg/server/server_test.go

@@ -20,10 +20,13 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"math/rand"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/golang/mock/gomock"
+	"github.com/google/uuid"
 	"google.golang.org/grpc/metadata"
 
 	authv1 "k8s.io/api/authentication/v1"
@@ -32,6 +35,7 @@ import (
 	fakeauthenticationv1 "k8s.io/client-go/kubernetes/typed/authentication/v1/fake"
 	k8stesting "k8s.io/client-go/testing"
 
+	client "sigs.k8s.io/apiserver-network-proxy/konnectivity-client/proto/client"
 	agentmock "sigs.k8s.io/apiserver-network-proxy/proto/agent/mocks"
 	"sigs.k8s.io/apiserver-network-proxy/proto/header"
 )
@@ -218,3 +222,140 @@ func TestAddRemoveFrontends(t *testing.T) {
 		t.Errorf("expected %v, got %v", e, a)
 	}
 }
+
+func prepareFrontendConn(ctrl *gomock.Controller) *agentmock.MockAgentService_ConnectServer {
+	// prepare the connection to fontend  of proxy-server
+	frontendConn := agentmock.NewMockAgentService_ConnectServer(ctrl)
+	frontendConnMD := metadata.MD{
+		":authority":   []string{"127.0.0.1:8090"},
+		"content-type": []string{"application/grpc"},
+		"user-agent":   []string{"grpc-go/1.42.0"},
+	}
+	frontendConnCtx := metadata.NewIncomingContext(context.Background(), frontendConnMD)
+	frontendConn.EXPECT().Context().Return(frontendConnCtx).AnyTimes()
+	return frontendConn
+}
+
+func prepareAgentConnMD(ctrl *gomock.Controller, proxyServer *ProxyServer) *agentmock.MockAgentService_ConnectServer {
+	// prepare the the connection to agent of proxy-server
+	agentConn := agentmock.NewMockAgentService_ConnectServer(ctrl)
+	agentConnMD := metadata.MD{
+		":authority":       []string{"127.0.0.1:8091"},
+		"agentid":          []string{uuid.New().String()},
+		"agentidentifiers": []string{},
+		"content-type":     []string{"application/grpc"},
+		"user-agent":       []string{"grpc-go/1.42.0"},
+	}
+	agentConnCtx := metadata.NewIncomingContext(context.Background(), agentConnMD)
+	agentConn.EXPECT().Context().Return(agentConnCtx).AnyTimes()
+
+	_ = proxyServer.addBackend(uuid.New().String(), agentConn)
+	return agentConn
+}
+
+func baseServerProxyTestWithoutBackend(t *testing.T, validate func(*agentmock.MockAgentService_ConnectServer)) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	frontendConn := prepareFrontendConn(ctrl)
+	proxyServer := NewProxyServer(uuid.New().String(), []ProxyStrategy{ProxyStrategyDefault}, 1, &AgentTokenAuthenticationOptions{}, true)
+
+	validate(frontendConn)
+
+	proxyServer.Proxy(frontendConn)
+
+	// add a sleep to make sure `serveRecvFrontend` ends after `Proxy` finished.
+	time.Sleep(1 * time.Second)
+}
+
+func baseServerProxyTestWithBackend(t *testing.T, validate func(*agentmock.MockAgentService_ConnectServer, *agentmock.MockAgentService_ConnectServer)) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	frontendConn := prepareFrontendConn(ctrl)
+
+	// prepare proxy server
+	proxyServer := NewProxyServer(uuid.New().String(), []ProxyStrategy{ProxyStrategyDefault}, 1, &AgentTokenAuthenticationOptions{}, true)
+
+	agentConn := prepareAgentConnMD(ctrl, proxyServer)
+
+	validate(frontendConn, agentConn)
+
+	proxyServer.Proxy(frontendConn)
+
+	// add a sleep to make sure `serveRecvFrontend` ends after `Proxy` finished.
+	time.Sleep(1 * time.Second)
+}
+
+func TestServerProxyNoBackend(t *testing.T) {
+	validate := func(frontendConn *agentmock.MockAgentService_ConnectServer) {
+		// receive DIAL_REQ from frontend and proxy to backend
+		randomID := rand.Int63() /* #nosec G404 */
+		dialReq := &client.Packet{
+			Type: client.PacketType_DIAL_REQ,
+			Payload: &client.Packet_DialRequest{
+				DialRequest: &client.DialRequest{
+					Protocol: "tcp",
+					Address:  "127.0.0.1:8080",
+					Random:   randomID,
+				},
+			},
+		}
+
+		dialResp := &client.Packet{
+			Type: client.PacketType_DIAL_RSP,
+			Payload: &client.Packet_DialResponse{
+				DialResponse: &client.DialResponse{
+					Random: randomID,
+					Error:  (&ErrNotFound{}).Error(),
+				}},
+		}
+
+		gomock.InOrder(
+			frontendConn.EXPECT().Recv().Return(dialReq, nil).Times(1),
+			frontendConn.EXPECT().Recv().Return(nil, io.EOF).Times(1),
+			// NOTE(mainred): `Send` should come before `Recv` io.EOF, but we cannot add wait between
+			//                two Recvs, thus `Recv`` comes before `Send`
+			frontendConn.EXPECT().Send(dialResp).Return(nil).Times(1),
+		)
+
+	}
+	baseServerProxyTestWithoutBackend(t, validate)
+}
+
+func TestServerProxyNormalClose(t *testing.T) {
+	validate := func(frontendConn, agentConn *agentmock.MockAgentService_ConnectServer) {
+		// receive DIAL_REQ from frontend and proxy to backend
+		randomID := rand.Int63() /* #nosec G404 */
+		dialReq := &client.Packet{
+			Type: client.PacketType_DIAL_REQ,
+			Payload: &client.Packet_DialRequest{
+				DialRequest: &client.DialRequest{
+					Protocol: "tcp",
+					Address:  "127.0.0.1:8080",
+					Random:   randomID,
+				},
+			},
+		}
+
+		// recevie CLOSE_REQ from frontend and proxy to backend
+		closeReq := &client.Packet{
+			Type: client.PacketType_CLOSE_REQ,
+			Payload: &client.Packet_CloseRequest{
+				CloseRequest: &client.CloseRequest{
+					ConnectID: 1,
+				}},
+		}
+
+		gomock.InOrder(
+			frontendConn.EXPECT().Recv().Return(dialReq, nil).Times(1),
+			frontendConn.EXPECT().Recv().Return(closeReq, nil).Times(1),
+			frontendConn.EXPECT().Recv().Return(nil, io.EOF).Times(1),
+		)
+		gomock.InOrder(
+			agentConn.EXPECT().Send(dialReq).Return(nil).Times(1),
+			agentConn.EXPECT().Send(closeReq).Return(nil).Times(1),
+		)
+	}
+	baseServerProxyTestWithBackend(t, validate)
+}