github.com/wangyougui/gf/v2@v2.6.5/net/gudp/gudp_server.go (about)

     1  // Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
     2  //
     3  // This Source Code Form is subject to the terms of the MIT License.
     4  // If a copy of the MIT was not distributed with this file,
     5  // You can obtain one at https://github.com/wangyougui/gf.
     6  
     7  package gudp
     8  
     9  import (
    10  	"fmt"
    11  	"net"
    12  	"sync"
    13  
    14  	"github.com/wangyougui/gf/v2/container/gmap"
    15  	"github.com/wangyougui/gf/v2/errors/gcode"
    16  	"github.com/wangyougui/gf/v2/errors/gerror"
    17  	"github.com/wangyougui/gf/v2/text/gstr"
    18  	"github.com/wangyougui/gf/v2/util/gconv"
    19  )
    20  
    21  const (
    22  	// FreePortAddress marks the server listens using random free port.
    23  	FreePortAddress = ":0"
    24  )
    25  
    26  const (
    27  	defaultServer = "default"
    28  )
    29  
    30  // Server is the UDP server.
    31  type Server struct {
    32  	mu      sync.Mutex  // Used for Server.listen concurrent safety. -- The golang test with data race checks this.
    33  	conn    *Conn       // UDP server connection object.
    34  	address string      // UDP server listening address.
    35  	handler func(*Conn) // Handler for UDP connection.
    36  }
    37  
    38  var (
    39  	// serverMapping is used for instance name to its UDP server mappings.
    40  	serverMapping = gmap.NewStrAnyMap(true)
    41  )
    42  
    43  // GetServer creates and returns an UDP server instance with given name.
    44  func GetServer(name ...interface{}) *Server {
    45  	serverName := defaultServer
    46  	if len(name) > 0 && name[0] != "" {
    47  		serverName = gconv.String(name[0])
    48  	}
    49  	if s := serverMapping.Get(serverName); s != nil {
    50  		return s.(*Server)
    51  	}
    52  	s := NewServer("", nil)
    53  	serverMapping.Set(serverName, s)
    54  	return s
    55  }
    56  
    57  // NewServer creates and returns an UDP server.
    58  // The optional parameter `name` is used to specify its name, which can be used for
    59  // GetServer function to retrieve its instance.
    60  func NewServer(address string, handler func(*Conn), name ...string) *Server {
    61  	s := &Server{
    62  		address: address,
    63  		handler: handler,
    64  	}
    65  	if len(name) > 0 && name[0] != "" {
    66  		serverMapping.Set(name[0], s)
    67  	}
    68  	return s
    69  }
    70  
    71  // SetAddress sets the server address for UDP server.
    72  func (s *Server) SetAddress(address string) {
    73  	s.address = address
    74  }
    75  
    76  // SetHandler sets the connection handler for UDP server.
    77  func (s *Server) SetHandler(handler func(*Conn)) {
    78  	s.handler = handler
    79  }
    80  
    81  // Close closes the connection.
    82  // It will make server shutdowns immediately.
    83  func (s *Server) Close() (err error) {
    84  	s.mu.Lock()
    85  	defer s.mu.Unlock()
    86  	err = s.conn.Close()
    87  	if err != nil {
    88  		err = gerror.Wrap(err, "connection failed")
    89  	}
    90  	return
    91  }
    92  
    93  // Run starts listening UDP connection.
    94  func (s *Server) Run() error {
    95  	if s.handler == nil {
    96  		err := gerror.NewCode(gcode.CodeMissingConfiguration, "start running failed: socket handler not defined")
    97  		return err
    98  	}
    99  	addr, err := net.ResolveUDPAddr("udp", s.address)
   100  	if err != nil {
   101  		err = gerror.Wrapf(err, `net.ResolveUDPAddr failed for address "%s"`, s.address)
   102  		return err
   103  	}
   104  	conn, err := net.ListenUDP("udp", addr)
   105  	if err != nil {
   106  		err = gerror.Wrapf(err, `net.ListenUDP failed for address "%s"`, s.address)
   107  		return err
   108  	}
   109  	s.mu.Lock()
   110  	s.conn = NewConnByNetConn(conn)
   111  	s.mu.Unlock()
   112  	s.handler(s.conn)
   113  	return nil
   114  }
   115  
   116  // GetListenedAddress retrieves and returns the address string which are listened by current server.
   117  func (s *Server) GetListenedAddress() string {
   118  	if !gstr.Contains(s.address, FreePortAddress) {
   119  		return s.address
   120  	}
   121  	var (
   122  		address      = s.address
   123  		listenedPort = s.GetListenedPort()
   124  	)
   125  	address = gstr.Replace(address, FreePortAddress, fmt.Sprintf(`:%d`, listenedPort))
   126  	return address
   127  }
   128  
   129  // GetListenedPort retrieves and returns one port which is listened to by current server.
   130  func (s *Server) GetListenedPort() int {
   131  	s.mu.Lock()
   132  	defer s.mu.Unlock()
   133  	if ln := s.conn; ln != nil {
   134  		return ln.LocalAddr().(*net.UDPAddr).Port
   135  	}
   136  	return -1
   137  }