github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/network/tcp_transport.go (about)

     1  package network
     2  
     3  import (
     4  	"errors"
     5  	"net"
     6  	"sync"
     7  	"time"
     8  
     9  	"go.uber.org/zap"
    10  )
    11  
    12  // TCPTransport allows network communication over TCP.
    13  type TCPTransport struct {
    14  	log      *zap.Logger
    15  	server   *Server
    16  	listener net.Listener
    17  	bindAddr string
    18  	hostPort hostPort
    19  	lock     sync.RWMutex
    20  	quit     bool
    21  }
    22  
    23  type hostPort struct {
    24  	Host string
    25  	Port string
    26  }
    27  
    28  // NewTCPTransport returns a new TCPTransport that will listen for
    29  // new incoming peer connections.
    30  func NewTCPTransport(s *Server, bindAddr string, log *zap.Logger) *TCPTransport {
    31  	host, port, err := net.SplitHostPort(bindAddr)
    32  	if err != nil {
    33  		// Only host can be provided, it's OK.
    34  		host = bindAddr
    35  	}
    36  	return &TCPTransport{
    37  		log:      log,
    38  		server:   s,
    39  		bindAddr: bindAddr,
    40  		hostPort: hostPort{
    41  			Host: host,
    42  			Port: port,
    43  		},
    44  	}
    45  }
    46  
    47  // Dial implements the Transporter interface.
    48  func (t *TCPTransport) Dial(addr string, timeout time.Duration) (AddressablePeer, error) {
    49  	conn, err := net.DialTimeout("tcp", addr, timeout)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	p := NewTCPPeer(conn, addr, t.server)
    54  	go p.handleConn()
    55  	return p, nil
    56  }
    57  
    58  // Accept implements the Transporter interface.
    59  func (t *TCPTransport) Accept() {
    60  	l, err := net.Listen("tcp", t.bindAddr)
    61  	if err != nil {
    62  		t.log.Panic("TCP listen error", zap.Error(err))
    63  		return
    64  	}
    65  
    66  	t.lock.Lock()
    67  	if t.quit {
    68  		t.lock.Unlock()
    69  		l.Close()
    70  		return
    71  	}
    72  	t.listener = l
    73  	t.bindAddr = l.Addr().String()
    74  	t.hostPort.Host, t.hostPort.Port, _ = net.SplitHostPort(t.bindAddr) // no error expected as l.Addr() is a valid address.
    75  	t.lock.Unlock()
    76  
    77  	for {
    78  		conn, err := l.Accept()
    79  		if err != nil {
    80  			t.lock.Lock()
    81  			quit := t.quit
    82  			t.lock.Unlock()
    83  			if errors.Is(err, net.ErrClosed) && quit {
    84  				break
    85  			}
    86  			t.log.Warn("TCP accept error", zap.Stringer("address", l.Addr()), zap.Error(err))
    87  			continue
    88  		}
    89  		p := NewTCPPeer(conn, "", t.server)
    90  		go p.handleConn()
    91  	}
    92  }
    93  
    94  // Close implements the Transporter interface.
    95  func (t *TCPTransport) Close() {
    96  	t.lock.Lock()
    97  	if t.listener != nil {
    98  		t.listener.Close()
    99  	}
   100  	t.quit = true
   101  	t.lock.Unlock()
   102  }
   103  
   104  // Proto implements the Transporter interface.
   105  func (t *TCPTransport) Proto() string {
   106  	return "tcp"
   107  }
   108  
   109  // HostPort implements the Transporter interface.
   110  func (t *TCPTransport) HostPort() (string, string) {
   111  	t.lock.RLock()
   112  	defer t.lock.RUnlock()
   113  
   114  	return t.hostPort.Host, t.hostPort.Port
   115  }