github.com/netdata/go.d.plugin@v0.58.1/pkg/socket/client.go (about)

     1  // SPDX-License-Identifier: GPL-3.0-or-later
     2  
     3  package socket
     4  
     5  import (
     6  	"bufio"
     7  	"crypto/tls"
     8  	"errors"
     9  	"net"
    10  	"time"
    11  )
    12  
    13  // New returns a new pointer to a socket client given the socket
    14  // type (IP, TCP, UDP, UNIX), a network address (IP/domain:port),
    15  // a timeout and a TLS config. It supports both IPv4 and IPv6 address
    16  // and reuses connection where possible.
    17  func New(config Config) *Socket {
    18  	return &Socket{
    19  		Config: config,
    20  		conn:   nil,
    21  	}
    22  }
    23  
    24  // Socket is the implementation of a socket client.
    25  type Socket struct {
    26  	Config
    27  	conn net.Conn
    28  }
    29  
    30  // Connect connects to the Socket address on the named network.
    31  // If the address is a domain name it will also perform the DNS resolution.
    32  // Address like :80 will attempt to connect to the localhost.
    33  // The config timeout and TLS config will be used.
    34  func (s *Socket) Connect() (err error) {
    35  	network, address := networkType(s.Address)
    36  	if s.TLSConf == nil {
    37  		s.conn, err = net.DialTimeout(network, address, s.ConnectTimeout)
    38  	} else {
    39  		var d net.Dialer
    40  		d.Timeout = s.ConnectTimeout
    41  		s.conn, err = tls.DialWithDialer(&d, network, address, s.TLSConf)
    42  	}
    43  	return err
    44  }
    45  
    46  // Disconnect closes the connection.
    47  // Any in-flight commands will be cancelled and return errors.
    48  func (s *Socket) Disconnect() (err error) {
    49  	if s.conn != nil {
    50  		err = s.conn.Close()
    51  		s.conn = nil
    52  	}
    53  	return err
    54  }
    55  
    56  // Command writes the command string to the connection and passed the
    57  // response bytes line by line to the process function. It uses the
    58  // timeout value from the Socket config and returns read, write and
    59  // timeout errors if any. If a timeout occurs during the processing
    60  // of the responses this function will stop processing and return a
    61  // timeout error.
    62  func (s *Socket) Command(command string, process Processor) error {
    63  	if s.conn == nil {
    64  		return errors.New("cannot send command on nil connection")
    65  	}
    66  	if err := write(command, s.conn, s.WriteTimeout); err != nil {
    67  		return err
    68  	}
    69  	return read(s.conn, process, s.ReadTimeout)
    70  }
    71  
    72  func write(command string, writer net.Conn, timeout time.Duration) error {
    73  	if writer == nil {
    74  		return errors.New("attempt to write on nil connection")
    75  	}
    76  	if err := writer.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
    77  		return err
    78  	}
    79  	_, err := writer.Write([]byte(command))
    80  	return err
    81  }
    82  
    83  func read(reader net.Conn, process Processor, timeout time.Duration) error {
    84  	if process == nil {
    85  		return errors.New("process func is nil")
    86  	}
    87  	if reader == nil {
    88  		return errors.New("attempt to read on nil connection")
    89  	}
    90  	if err := reader.SetReadDeadline(time.Now().Add(timeout)); err != nil {
    91  		return err
    92  	}
    93  	scanner := bufio.NewScanner(reader)
    94  	for scanner.Scan() && process(scanner.Bytes()) {
    95  	}
    96  	return scanner.Err()
    97  }