trpc.group/trpc-go/trpc-go@v1.0.3/internal/packetbuffer/packetbuffer_test.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package packetbuffer_test
    15  
    16  import (
    17  	"context"
    18  	"io"
    19  	"log"
    20  	"net"
    21  	"testing"
    22  
    23  	"github.com/stretchr/testify/require"
    24  	"trpc.group/trpc-go/trpc-go/internal/packetbuffer"
    25  )
    26  
    27  type udpServer struct {
    28  	cancel context.CancelFunc
    29  	conn   net.PacketConn
    30  }
    31  
    32  func (s *udpServer) start(ctx context.Context) error {
    33  	var err error
    34  	s.conn, err = net.ListenPacket("udp", "127.0.0.1:0")
    35  	if err != nil {
    36  		return err
    37  	}
    38  	ctx, s.cancel = context.WithCancel(ctx)
    39  	go func() {
    40  		buf := make([]byte, 65535)
    41  		for {
    42  			select {
    43  			case <-ctx.Done():
    44  				return
    45  			default:
    46  			}
    47  			n, addr, err := s.conn.ReadFrom(buf)
    48  			if err != nil {
    49  				log.Println("l.ReadFrom err: ", err)
    50  				return
    51  			}
    52  			s.conn.WriteTo(buf[:n], addr)
    53  		}
    54  	}()
    55  	return nil
    56  }
    57  
    58  func (s *udpServer) stop() {
    59  	s.cancel()
    60  	s.conn.Close()
    61  }
    62  
    63  func TestPacketReaderSucceed(t *testing.T) {
    64  	s := &udpServer{}
    65  	s.start(context.Background())
    66  	t.Cleanup(s.stop)
    67  
    68  	p, err := net.ListenPacket("udp", "127.0.0.1:0")
    69  	require.Nil(t, err)
    70  	_, err = p.WriteTo([]byte("helloworldA"), s.conn.LocalAddr())
    71  	require.Nil(t, err)
    72  	buf := packetbuffer.New(p, 65535)
    73  	defer buf.Close()
    74  	result := make([]byte, 20)
    75  	n, err := buf.Read(result)
    76  	require.Nil(t, err)
    77  	require.Equal(t, []byte("helloworldA"), result[:n])
    78  	require.Equal(t, s.conn.LocalAddr(), buf.CurrentPacketAddr())
    79  	_, err = buf.Read(result)
    80  	require.Equal(t, io.EOF, err)
    81  	require.Nil(t, buf.Next())
    82  
    83  	_, err = p.WriteTo([]byte("helloworldB"), s.conn.LocalAddr())
    84  	require.Nil(t, err)
    85  	n, err = buf.Read(result)
    86  	require.Nil(t, err)
    87  	require.Equal(t, []byte("helloworldB"), result[:n])
    88  }
    89  
    90  func TestPacketReaderFailed(t *testing.T) {
    91  	s := &udpServer{}
    92  	s.start(context.Background())
    93  	t.Cleanup(s.stop)
    94  
    95  	p, err := net.ListenPacket("udp", "127.0.0.1:0")
    96  	require.Nil(t, err)
    97  	_, err = p.WriteTo([]byte("helloworld"), s.conn.LocalAddr())
    98  	require.Nil(t, err)
    99  	buf := packetbuffer.New(p, 65535)
   100  	defer buf.Close()
   101  	n, err := buf.Read(nil)
   102  	require.Nil(t, err)
   103  	require.Equal(t, 0, n)
   104  	result := make([]byte, 5)
   105  	_, err = buf.Read(result)
   106  	require.Nil(t, err)
   107  	// There are some remaining data in the buf that have not been read.
   108  	require.NotNil(t, buf.Next())
   109  }