github.com/mailgun/holster/v4@v4.20.0/udp/server.go (about)

     1  package udp
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"strings"
     7  
     8  	"github.com/sirupsen/logrus"
     9  )
    10  
    11  type Server struct {
    12  	conn     net.PacketConn
    13  	shutdown chan struct{}
    14  }
    15  
    16  type Handler func(net.PacketConn, []byte, net.Addr)
    17  
    18  type ServerConfig struct {
    19  	BindAddress string
    20  	Handler     Handler
    21  }
    22  
    23  func NewServer(conf ServerConfig) (*Server, error) {
    24  	conn, err := net.ListenPacket("udp", conf.BindAddress)
    25  	if err != nil {
    26  		return nil, fmt.Errorf("while listening on '%s' - %w", conf.BindAddress, err)
    27  	}
    28  
    29  	shutdown := make(chan struct{})
    30  	logrus.Debugf("Listening [%s]...\n", conf.BindAddress)
    31  
    32  	go func() {
    33  		for {
    34  			select {
    35  			case <-shutdown:
    36  				return
    37  			default:
    38  			}
    39  
    40  			b := make([]byte, 10_000)
    41  			n, addr, err := conn.ReadFrom(b)
    42  			if err != nil {
    43  				if strings.HasSuffix(err.Error(), "use of closed network connection") {
    44  					return
    45  				}
    46  				logrus.WithError(err).Error("ReadFrom() failed")
    47  				return
    48  			}
    49  			// fmt.Printf("packet-received: bytes=%d from=%s\n", n, addr.String())
    50  			conf.Handler(conn, b[:n], addr)
    51  		}
    52  	}()
    53  	return &Server{
    54  		conn:     conn,
    55  		shutdown: shutdown,
    56  	}, nil
    57  }
    58  
    59  func (u *Server) Close() {
    60  	close(u.shutdown)
    61  	u.conn.Close()
    62  }