neo-go/pkg/network/tcp_transport.go
Roman Khimov e79dec15f9 *: use zap.Stringer instead of zap.String where it can be used
It's a bit more efficient in case we're not logging the message (mostly for
debug), makes the code somewhat simpler as well.
2022-12-13 12:44:54 +03:00

115 lines
2.3 KiB
Go

package network
import (
"errors"
"net"
"sync"
"time"
"go.uber.org/zap"
)
// TCPTransport allows network communication over TCP.
type TCPTransport struct {
log *zap.Logger
server *Server
listener net.Listener
bindAddr string
hostPort hostPort
lock sync.RWMutex
quit bool
}
type hostPort struct {
Host string
Port string
}
// NewTCPTransport returns a new TCPTransport that will listen for
// new incoming peer connections.
func NewTCPTransport(s *Server, bindAddr string, log *zap.Logger) *TCPTransport {
host, port, err := net.SplitHostPort(bindAddr)
if err != nil {
// Only host can be provided, it's OK.
host = bindAddr
}
return &TCPTransport{
log: log,
server: s,
bindAddr: bindAddr,
hostPort: hostPort{
Host: host,
Port: port,
},
}
}
// Dial implements the Transporter interface.
func (t *TCPTransport) Dial(addr string, timeout time.Duration) (AddressablePeer, error) {
conn, err := net.DialTimeout("tcp", addr, timeout)
if err != nil {
return nil, err
}
p := NewTCPPeer(conn, addr, t.server)
go p.handleConn()
return p, nil
}
// Accept implements the Transporter interface.
func (t *TCPTransport) Accept() {
l, err := net.Listen("tcp", t.bindAddr)
if err != nil {
t.log.Panic("TCP listen error", zap.Error(err))
return
}
t.lock.Lock()
if t.quit {
t.lock.Unlock()
l.Close()
return
}
t.listener = l
t.bindAddr = l.Addr().String()
t.hostPort.Host, t.hostPort.Port, _ = net.SplitHostPort(t.bindAddr) // no error expected as l.Addr() is a valid address.
t.lock.Unlock()
for {
conn, err := l.Accept()
if err != nil {
t.lock.Lock()
quit := t.quit
t.lock.Unlock()
if errors.Is(err, net.ErrClosed) && quit {
break
}
t.log.Warn("TCP accept error", zap.Stringer("address", l.Addr()), zap.Error(err))
continue
}
p := NewTCPPeer(conn, "", t.server)
go p.handleConn()
}
}
// Close implements the Transporter interface.
func (t *TCPTransport) Close() {
t.lock.Lock()
if t.listener != nil {
t.listener.Close()
}
t.quit = true
t.lock.Unlock()
}
// Proto implements the Transporter interface.
func (t *TCPTransport) Proto() string {
return "tcp"
}
// HostPort implements the Transporter interface.
func (t *TCPTransport) HostPort() (string, string) {
t.lock.RLock()
defer t.lock.RUnlock()
return t.hostPort.Host, t.hostPort.Port
}