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

     1  // SPDX-License-Identifier: GPL-3.0-or-later
     2  
     3  package socket
     4  
     5  import (
     6  	"bufio"
     7  	"errors"
     8  	"fmt"
     9  	"net"
    10  	"os"
    11  	"strings"
    12  	"time"
    13  )
    14  
    15  type tcpServer struct {
    16  	addr        string
    17  	server      net.Listener
    18  	rowsNumResp int
    19  }
    20  
    21  func (t *tcpServer) Run() (err error) {
    22  	t.server, err = net.Listen("tcp", t.addr)
    23  	if err != nil {
    24  		return
    25  	}
    26  	return t.handleConnections()
    27  }
    28  
    29  func (t *tcpServer) Close() (err error) {
    30  	return t.server.Close()
    31  }
    32  
    33  func (t *tcpServer) handleConnections() (err error) {
    34  	for {
    35  		conn, err := t.server.Accept()
    36  		if err != nil || conn == nil {
    37  			return errors.New("could not accept connection")
    38  		}
    39  		t.handleConnection(conn)
    40  	}
    41  }
    42  
    43  func (t *tcpServer) handleConnection(conn net.Conn) {
    44  	defer func() { _ = conn.Close() }()
    45  	_ = conn.SetDeadline(time.Now().Add(time.Millisecond * 100))
    46  
    47  	rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
    48  	_, err := rw.ReadString('\n')
    49  	if err != nil {
    50  		_, _ = rw.WriteString("failed to read input")
    51  		_ = rw.Flush()
    52  	} else {
    53  		resp := strings.Repeat("pong\n", t.rowsNumResp)
    54  		_, _ = rw.WriteString(resp)
    55  		_ = rw.Flush()
    56  	}
    57  }
    58  
    59  type udpServer struct {
    60  	addr        string
    61  	conn        *net.UDPConn
    62  	rowsNumResp int
    63  }
    64  
    65  func (u *udpServer) Run() (err error) {
    66  	addr, err := net.ResolveUDPAddr("udp", u.addr)
    67  	if err != nil {
    68  		return err
    69  	}
    70  	u.conn, err = net.ListenUDP("udp", addr)
    71  	if err != nil {
    72  		return
    73  	}
    74  	u.handleConnections()
    75  	return nil
    76  }
    77  
    78  func (u *udpServer) Close() (err error) {
    79  	return u.conn.Close()
    80  }
    81  
    82  func (u *udpServer) handleConnections() {
    83  	for {
    84  		var buf [2048]byte
    85  		_, addr, _ := u.conn.ReadFromUDP(buf[0:])
    86  		resp := strings.Repeat("pong\n", u.rowsNumResp)
    87  		_, _ = u.conn.WriteToUDP([]byte(resp), addr)
    88  	}
    89  }
    90  
    91  type unixServer struct {
    92  	addr        string
    93  	conn        *net.UnixListener
    94  	rowsNumResp int
    95  }
    96  
    97  func (u *unixServer) Run() (err error) {
    98  	_, _ = os.CreateTemp("/tmp", "testSocketFD")
    99  	addr, err := net.ResolveUnixAddr("unix", u.addr)
   100  	if err != nil {
   101  		return err
   102  	}
   103  	u.conn, err = net.ListenUnix("unix", addr)
   104  	if err != nil {
   105  		return
   106  	}
   107  	go u.handleConnections()
   108  	return nil
   109  }
   110  
   111  func (u *unixServer) Close() (err error) {
   112  	_ = os.Remove(testUnixServerAddress)
   113  	return u.conn.Close()
   114  }
   115  
   116  func (u *unixServer) handleConnections() {
   117  	var conn net.Conn
   118  	var err error
   119  	conn, err = u.conn.AcceptUnix()
   120  	if err != nil {
   121  		panic(fmt.Errorf("could not accept connection: %v", err))
   122  	}
   123  	u.handleConnection(conn)
   124  }
   125  
   126  func (u *unixServer) handleConnection(conn net.Conn) {
   127  	_ = conn.SetDeadline(time.Now().Add(time.Second))
   128  
   129  	rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
   130  	_, err := rw.ReadString('\n')
   131  	if err != nil {
   132  		_, _ = rw.WriteString("failed to read input")
   133  		_ = rw.Flush()
   134  	} else {
   135  		resp := strings.Repeat("pong\n", u.rowsNumResp)
   136  		_, _ = rw.WriteString(resp)
   137  		_ = rw.Flush()
   138  	}
   139  }