gitee.com/zhaochuninhefei/gmgo@v0.0.31-0.20240209061119-069254a02979/grpc/credentials/alts/internal/testutil/testutil.go (about)

     1  /*
     2   *
     3   * Copyright 2018 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  // Package testutil include useful test utilities for the handshaker.
    20  package testutil
    21  
    22  import (
    23  	"bytes"
    24  	"encoding/binary"
    25  	"io"
    26  	"net"
    27  	"sync"
    28  
    29  	"gitee.com/zhaochuninhefei/gmgo/grpc/credentials/alts/internal/conn"
    30  )
    31  
    32  // Stats is used to collect statistics about concurrent handshake calls.
    33  type Stats struct {
    34  	mu                 sync.Mutex
    35  	calls              int
    36  	MaxConcurrentCalls int
    37  }
    38  
    39  // Update updates the statistics by adding one call.
    40  func (s *Stats) Update() func() {
    41  	s.mu.Lock()
    42  	s.calls++
    43  	if s.calls > s.MaxConcurrentCalls {
    44  		s.MaxConcurrentCalls = s.calls
    45  	}
    46  	s.mu.Unlock()
    47  
    48  	return func() {
    49  		s.mu.Lock()
    50  		s.calls--
    51  		s.mu.Unlock()
    52  	}
    53  }
    54  
    55  // Reset resets the statistics.
    56  func (s *Stats) Reset() {
    57  	s.mu.Lock()
    58  	defer s.mu.Unlock()
    59  	s.calls = 0
    60  	s.MaxConcurrentCalls = 0
    61  }
    62  
    63  // testConn mimics a net.Conn to the peer.
    64  type testConn struct {
    65  	net.Conn
    66  	in  *bytes.Buffer
    67  	out *bytes.Buffer
    68  }
    69  
    70  // NewTestConn creates a new instance of testConn object.
    71  func NewTestConn(in *bytes.Buffer, out *bytes.Buffer) net.Conn {
    72  	return &testConn{
    73  		in:  in,
    74  		out: out,
    75  	}
    76  }
    77  
    78  // Read reads from the in buffer.
    79  func (c *testConn) Read(b []byte) (n int, err error) {
    80  	return c.in.Read(b)
    81  }
    82  
    83  // Write writes to the out buffer.
    84  func (c *testConn) Write(b []byte) (n int, err error) {
    85  	return c.out.Write(b)
    86  }
    87  
    88  // Close closes the testConn object.
    89  func (c *testConn) Close() error {
    90  	return nil
    91  }
    92  
    93  // unresponsiveTestConn mimics a net.Conn for an unresponsive peer. It is used
    94  // for testing the PeerNotResponding case.
    95  type unresponsiveTestConn struct {
    96  	net.Conn
    97  }
    98  
    99  // NewUnresponsiveTestConn creates a new instance of unresponsiveTestConn object.
   100  func NewUnresponsiveTestConn() net.Conn {
   101  	return &unresponsiveTestConn{}
   102  }
   103  
   104  // Read reads from the in buffer.
   105  func (c *unresponsiveTestConn) Read(b []byte) (n int, err error) {
   106  	return 0, io.EOF
   107  }
   108  
   109  // Write writes to the out buffer.
   110  func (c *unresponsiveTestConn) Write(b []byte) (n int, err error) {
   111  	return 0, nil
   112  }
   113  
   114  // Close closes the TestConn object.
   115  func (c *unresponsiveTestConn) Close() error {
   116  	return nil
   117  }
   118  
   119  // MakeFrame creates a handshake frame.
   120  func MakeFrame(pl string) []byte {
   121  	f := make([]byte, len(pl)+conn.MsgLenFieldSize)
   122  	binary.LittleEndian.PutUint32(f, uint32(len(pl)))
   123  	copy(f[conn.MsgLenFieldSize:], []byte(pl))
   124  	return f
   125  }