fix to pass test

Signed-off-by: Patrik Cyvoct <[email protected]>

Author: Sh4d1

pkg/agent/agentclient/client.go

@@ -87,7 +87,10 @@ func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 		pkt, err := a.stream.Recv()
 		if err != nil {
 			if err2, ok := err.(*ReconnectError); ok {
-				err = err2.Wait()
+				err3 := err2.Wait()
+				if err3 != nil {
+					klog.Warningf("reconnect error: %v", err3)
+				}
 				continue
 			} else if err == io.EOF {
 				klog.Info("received EOF, exit")
@@ -102,6 +105,11 @@ func (a *AgentClient) Serve(stopCh <-chan struct{}) {
 
 		klog.Infof("[tracing] recv packet %+v", pkt)
 
+		if pkt == nil {
+			klog.Warningf("empty packet received")
+			continue
+		}
+
 		switch pkt.Type {
 		case agent.PacketType_DIAL_REQ:
 			klog.Info("received DIAL_REQ")

pkg/agent/agentclient/client_test.go

@@ -38,13 +38,16 @@ func TestServeData_HTTP(t *testing.T) {
 	dialPacket := newDialPacket("tcp", ts.URL[len("http://"):], 111)
 	err = stream.Send(dialPacket)
 	if err != nil {
-		t.Error(err.Error())
+		t.Fatal(err.Error())
 	}
 
 	// Expect receiving DIAL_RSP packet from AgentClient
-	pkg, _ := stream.Recv()
+	pkg, err := stream.Recv()
+	if err != nil {
+		t.Fatal(err.Error())
+	}
 	if pkg == nil {
-		t.Error("unexpected nil packet")
+		t.Fatal("unexpected nil packet")
 	}
 	if pkg.Type != agent.PacketType_DIAL_RSP {
 		t.Errorf("expect PacketType_DIAL_RSP; got %v", pkg.Type)
@@ -135,7 +138,7 @@ func TestClose_Client(t *testing.T) {
 	// Expect receiving DIAL_RSP packet from AgentClient
 	pkg, _ := stream.Recv()
 	if pkg == nil {
-		t.Error("unexpected nil packet")
+		t.Fatal("unexpected nil packet")
 	}
 	if pkg.Type != agent.PacketType_DIAL_RSP {
 		t.Errorf("expect PacketType_DIAL_RSP; got %v", pkg.Type)

pkg/agent/agentclient/stream.go

@@ -152,17 +152,20 @@ func (c *RedialableAgentClient) Recv() (*agent.Packet, error) {
 	c.recvLock.Lock()
 	defer c.recvLock.Unlock()
 
-	if pkt, err := c.stream.Recv(); err != nil {
+	var pkt *agent.Packet
+	var err error
+
+	if pkt, err = c.stream.Recv(); err != nil {
 		if err == io.EOF {
 			return pkt, err
 		}
 		return pkt, &ReconnectError{
 			internalErr: err,
 			errChan:     c.triggerReconnect(),
 		}
-	} else {
-		return pkt, nil
 	}
+
+	return pkt, nil
 }
 
 func (c *RedialableAgentClient) Connect() error {

pkg/agent/agentserver/server.go

@@ -54,7 +54,7 @@ func (c *ProxyClientConnection) send(pkt *agent.Packet) error {
 	}
 }
 
-// ProxyServer 
+// ProxyServer
 type ProxyServer struct {
 	Backend agent.AgentService_ConnectServer
 

tests/proxy_test.go

@@ -35,6 +35,9 @@ func TestBasicProxy_GRPC(t *testing.T) {
 	defer close(stopCh)
 
 	proxy, cleanup, err := runGRPCProxyServer()
+	if err != nil {
+		t.Fatal(err)
+	}
 	defer cleanup()
 
 	if err := runAgent(proxy.agent, stopCh); err != nil {
@@ -79,6 +82,9 @@ func TestBasicProxy_HTTPCONN(t *testing.T) {
 	defer close(stopCh)
 
 	proxy, cleanup, err := runHTTPConnProxyServer()
+	if err != nil {
+		t.Fatal(err)
+	}
 	defer cleanup()
 
 	if err := runAgent(proxy.agent, stopCh); err != nil {