handle data package to agent in the same goroutine

cherry pick unit test from the irozzo-1A/apiserver-hetwork-proxy fork

Author: Mikhail Sakhnov

pkg/server/server.go

@@ -756,18 +756,19 @@ func (s *ProxyServer) serveRecvBackend(backend Backend, stream agent.AgentServic
 			klog.V(5).InfoS("Received data from agent", "bytes", len(resp.Data), "agentID", agentID, "connectionID", resp.ConnectID)
 			kasConnCtx, isFromAgent := s.forwardConnectionManager.Get(resp.ConnectID)
 			if isFromAgent {
-				kasConnCtx.dataCh <- resp.Data
+				s.handleDataPacketFromAgent(resp, kasConnCtx)
 				continue
-			}
-			frontend, err := s.getFrontend(agentID, resp.ConnectID)
-			if err != nil {
-				klog.ErrorS(err, "could not get frontend client", "connectionID", resp.ConnectID)
-				break
-			}
-			if err := frontend.send(pkt); err != nil {
-				klog.ErrorS(err, "send to client stream failure", "serverID", s.serverID, "agentID", agentID, "connectionID", resp.ConnectID)
 			} else {
-				klog.V(5).InfoS("DATA sent to frontend")
+				frontend, err := s.getFrontend(agentID, resp.ConnectID)
+				if err != nil {
+					klog.ErrorS(err, "could not get frontend client", "connectionID", resp.ConnectID)
+					break
+				}
+				if err := frontend.send(pkt); err != nil {
+					klog.ErrorS(err, "send to client stream failure", "serverID", s.serverID, "agentID", agentID, "connectionID", resp.ConnectID)
+				} else {
+					klog.V(5).InfoS("DATA sent to frontend")
+				}
 			}
 
 		case client.PacketType_CLOSE_RSP:
@@ -830,10 +831,8 @@ func (s *ProxyServer) handleDialRequest(pkt *client.Packet, backend Backend) {
 	// Odd identifiers are used for connections from node to master network,
 	// increment by 2 to maintain the invariant.
 	connID := atomic.AddInt64(&s.nextConnID, 2)
-	dataCh := make(chan []byte, 5)
 	kasConnCtx := &connContext{
-		conn:   conn,
-		dataCh: dataCh,
+		conn: conn,
 		cleanFunc: func() {
 			klog.V(4).InfoS("close connection", "connectionID", connID)
 			req := &client.Packet{
@@ -853,7 +852,6 @@ func (s *ProxyServer) handleDialRequest(pkt *client.Packet, backend Backend) {
 				klog.ErrorS(err, "error occurred while closing connection", "connectionID", connID)
 			}
 
-			close(dataCh)
 			s.forwardConnectionManager.Delete(connID)
 		},
 		backend: backend,
@@ -869,13 +867,11 @@ func (s *ProxyServer) handleDialRequest(pkt *client.Packet, backend Backend) {
 
 	// proxy data to and from KAS for the connection
 	go s.agentToProxy(connID, kasConnCtx)
-	go s.proxyToAgent(connID, kasConnCtx)
 }
 
 type connContext struct {
 	conn      net.Conn
 	cleanFunc func()
-	dataCh    chan []byte
 	cleanOnce sync.Once
 	backend   Backend
 }
@@ -915,24 +911,22 @@ func (s *ProxyServer) agentToProxy(connID int64, ctx *connContext) {
 	}
 }
 
-func (s *ProxyServer) proxyToAgent(connID int64, ctx *connContext) {
-	defer ctx.cleanup()
-
-	for d := range ctx.dataCh {
-		pos := 0
-		for {
-			n, err := ctx.conn.Write(d[pos:])
-			if err == nil {
-				klog.V(4).InfoS("write to remote", "connectionID", connID, "lastData", n)
-				break
-			} else if n > 0 {
-				// https://golang.org/pkg/io/#Writer specifies return non nil error if n < len(d)
-				klog.ErrorS(err, "write to remote with failure", "connectionID", connID, "lastData", n)
-				pos += n
-			} else {
-				klog.ErrorS(err, "conn write failure", "connectionID", connID)
-				return
-			}
+func (s *ProxyServer) handleDataPacketFromAgent(resp *client.Data, kasConnCtx *connContext) {
+	pos := 0
+	connID := resp.ConnectID
+	d := resp.Data
+	for {
+		n, err := kasConnCtx.conn.Write(d[pos:])
+		if err == nil {
+			klog.V(4).InfoS("write to remote", "connectionID", connID, "lastData", n)
+			break
+		} else if n > 0 {
+			// https://golang.org/pkg/io/#Writer specifies return non nil error if n < len(d)
+			klog.ErrorS(err, "write to remote with failure", "connectionID", connID, "lastData", n)
+			pos += n
+		} else {
+			klog.ErrorS(err, "conn write failure", "connectionID", connID)
+			return
 		}
 	}
 }

pkg/server/server_test.go

@@ -20,7 +20,10 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"math/rand"
+	"net"
 	"reflect"
+	"sync"
 	"testing"
 
 	"github.com/golang/mock/gomock"
@@ -32,6 +35,8 @@ 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"
+	"sigs.k8s.io/apiserver-network-proxy/proto/agent"
 	agentmock "sigs.k8s.io/apiserver-network-proxy/proto/agent/mocks"
 	"sigs.k8s.io/apiserver-network-proxy/proto/header"
 )
@@ -218,3 +223,139 @@ func TestAddRemoveFrontends(t *testing.T) {
 		t.Errorf("expected %v, got %v", e, a)
 	}
 }
+
+type testStream struct {
+	agent.AgentService_ConnectServer
+	ch chan *client.Packet
+}
+
+func (f testStream) SendHeader(md metadata.MD) error {
+	return nil
+}
+
+func (f testStream) Context() context.Context {
+	ctx := context.Background()
+	ctx = metadata.NewIncomingContext(ctx, metadata.Pairs(header.AgentID, "test-agent-1"))
+	return ctx
+}
+
+func (t testStream) Send(packet *client.Packet) error {
+	return nil
+}
+
+func (t testStream) Recv() (*client.Packet, error) {
+	v, ok := <-t.ch
+	if !ok {
+		return nil, io.EOF
+	}
+	return v, nil
+}
+
+func TestNodeToControlPlane(t *testing.T) {
+	testData := []byte("this is test data")
+	stub := gomock.NewController(t)
+	defer stub.Finish()
+	var wg sync.WaitGroup
+
+	// create proxy server
+	p := NewProxyServer("", []ProxyStrategy{ProxyStrategyDefault}, 1,
+		&AgentTokenAuthenticationOptions{}, false)
+
+	// start a controlplane server
+	lis, err := net.Listen("tcp", ":0")
+	if err != nil {
+		panic(err)
+	}
+
+	var receivedData []byte
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		conn, err := lis.Accept()
+		if err != nil {
+			panic(err)
+		}
+		for {
+			var data [20]byte
+			n, err := conn.Read(data[:])
+			if err == io.EOF {
+				break
+			}
+
+			if err != nil {
+				panic(err)
+			}
+			receivedData = append(receivedData, data[:n]...)
+		}
+	}()
+
+	stream := &testStream{ch: make(chan *client.Packet, 15)}
+	wg.Add(1)
+
+	go func() {
+		connID := int64(3)
+		stream.ch <- packetDialReq(lis.Addr().String())
+		stream.ch <- packetData(connID, testData)
+		stream.ch <- packetDataEOF(connID)
+		stream.ch <- packetCloseReq(connID)
+		close(stream.ch)
+		wg.Done()
+	}()
+	// connect the agent to proxy server
+	err = p.Connect(stream)
+	if err != nil {
+		panic(err)
+	}
+	wg.Wait()
+	if !reflect.DeepEqual(testData, receivedData) {
+		t.Fatalf("Expected: %v Got: %v", testData, receivedData)
+	}
+}
+
+func packetCloseReq(connID int64) *client.Packet {
+	return &client.Packet{
+		Type: client.PacketType_CLOSE_REQ,
+		Payload: &client.Packet_CloseRequest{
+			CloseRequest: &client.CloseRequest{
+				ConnectID: connID,
+			},
+		},
+	}
+}
+
+func packetData(connID int64, testData []byte) *client.Packet {
+	return &client.Packet{
+		Type: client.PacketType_DATA,
+		Payload: &client.Packet_Data{
+			Data: &client.Data{
+				ConnectID: connID,
+				Data:      testData,
+			},
+		},
+	}
+}
+
+func packetDialReq(addr string) *client.Packet {
+	return &client.Packet{
+		Type: client.PacketType_DIAL_REQ,
+		Payload: &client.Packet_DialRequest{
+			DialRequest: &client.DialRequest{
+				Protocol: "tcp",
+				Address:  addr,
+				Random:   rand.Int63(),
+			},
+		},
+	}
+}
+
+func packetDataEOF(connID int64) *client.Packet {
+	return &client.Packet{
+		Type: client.PacketType_DATA,
+		Payload: &client.Packet_Data{
+			Data: &client.Data{
+				ConnectID: connID,
+				Error:     io.EOF.Error(),
+			},
+		},
+	}
+}