Advertisement
Guest User

Untitled

a guest
Mar 24th, 2017
67
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 2.99 KB | None | 0 0
  1. // Compile with go build -o idserv main.go
  2. package main
  3.  
  4. import (
  5. "bytes"
  6. "flag"
  7. "fmt"
  8. "log"
  9. "net"
  10. "os"
  11. "strconv"
  12. "strings"
  13. "time"
  14. )
  15.  
  16. type server struct {
  17. name string
  18. ip string
  19. udpPort int
  20. tcpPort int
  21. }
  22.  
  23. type serverSlice []server
  24.  
  25. func usage() {
  26. fmt.Println(`Example Usage: ../../bin/rmb [-i siip] [-p sipt]
  27. Arguments:
  28. -i server ip
  29. -p server port
  30. -h show help
  31. return`)
  32. os.Exit(0)
  33. }
  34.  
  35. func main() {
  36. var ip = flag.String("i", "127.0.0.1", "server ip")
  37. var port = flag.Int("p", 59000, "server port")
  38. var help = flag.Bool("h", false, "show help")
  39.  
  40. if *help {
  41. usage()
  42. }
  43. flag.Parse()
  44.  
  45. addr := net.UDPAddr{
  46. Port: *port,
  47. IP: net.ParseIP(*ip),
  48. }
  49.  
  50. conn, err := net.ListenUDP("udp", &addr)
  51. if err != nil {
  52. log.Fatal(err)
  53. os.Exit(1)
  54. }
  55. defer conn.Close()
  56.  
  57. var servers serverSlice
  58. serverRecv := make(chan server, 10)
  59. askForServers := make(chan serverSlice)
  60.  
  61. timer := time.NewTimer(time.Minute * 1)
  62.  
  63. go handleUDPConn(conn, serverRecv, askForServers)
  64.  
  65. log.Println("Identity Server setup complete")
  66. for {
  67. select {
  68. case <-timer.C:
  69. servers = servers[:0]
  70. log.Println("Wiped Servers")
  71. case serverToAdd := <-serverRecv:
  72. for _, o := range servers {
  73. if (o.name == serverToAdd.name) &&
  74. (o.ip == serverToAdd.ip) &&
  75. (o.tcpPort == serverToAdd.tcpPort) &&
  76. (o.udpPort == serverToAdd.udpPort) {
  77. break
  78. }
  79. }
  80. log.Println("Server Added:", serverToAdd.name, "IP:", serverToAdd.ip, "Ports:", serverToAdd.tcpPort, serverToAdd.udpPort)
  81. servers = append(servers, serverToAdd)
  82. case <-askForServers:
  83. askForServers <- servers
  84. }
  85. }
  86. }
  87.  
  88. func handleRegistrations(data string, serverRecv chan server) {
  89. buffer := strings.Split(data, ";")
  90. var err error
  91.  
  92. if len(buffer) < 4 {
  93. log.Println("Invalid Register")
  94. return
  95. }
  96.  
  97. newServer := server{buffer[0], buffer[1], 0, 0}
  98.  
  99. newServer.udpPort, err = strconv.Atoi(buffer[2])
  100. if err != nil {
  101. fmt.Println(err)
  102. return
  103. }
  104.  
  105. newServer.tcpPort, err = strconv.Atoi(buffer[3])
  106. if err != nil {
  107. fmt.Println(err)
  108. return
  109. }
  110.  
  111. serverRecv <- newServer
  112. }
  113.  
  114. func handleServerRequests(askForServers chan serverSlice, addr *net.UDPAddr, conn *net.UDPConn) {
  115. askForServers <- nil
  116. servers := <-askForServers
  117. var buffer bytes.Buffer
  118.  
  119. buffer.WriteString("SERVERS\n")
  120. for _, srv := range servers {
  121. buffer.WriteString(fmt.Sprintf("%s;%s;%d;%d\n", srv.name, srv.ip, srv.udpPort, srv.tcpPort))
  122. }
  123.  
  124. log.Println("Sending Servers to", addr.String())
  125. conn.WriteToUDP(buffer.Bytes(), addr)
  126. }
  127.  
  128. func handleUDPConn(conn *net.UDPConn, serverRecv chan server, askForServers chan serverSlice) {
  129. buf := make([]byte, 1024)
  130.  
  131. for {
  132. n, addr, err := conn.ReadFromUDP(buf)
  133. if err != nil {
  134. log.Println(err)
  135. }
  136.  
  137. switch {
  138. case strings.Contains(string(buf[0:4]), "REG"):
  139. handleRegistrations(string(buf[4:n-2]), serverRecv)
  140. case strings.Contains(string(buf[0:11]), "GET_SERVERS"):
  141. go handleServerRequests(askForServers, addr, conn)
  142. default:
  143. log.Println("Unknown Command:", string(buf[0:n]))
  144. }
  145. }
  146. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement