gorilla-wsserver_20201129141444.go 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188
  1. package main
  2. import (
  3. "encoding/json"
  4. "flag"
  5. "fmt"
  6. "log"
  7. "net/http"
  8. "github.com/google/uuid"
  9. "github.com/gorilla/websocket"
  10. )
  11. var wsaddr = flag.String("wsaddr", "localhost:12345", "websocket service address")
  12. type WSMsg struct {
  13. Mtype string `json:"mtype"`
  14. ID int `json:"id,omitempty"`
  15. Payload string `json:"payload,string"`
  16. }
  17. type Handler struct {
  18. MType string
  19. Callback func(WSMsg, *WSPeer) error
  20. }
  21. type WSServer struct {
  22. upgrader *websocket.Upgrader
  23. handler map[string]func(WSMsg, *WSPeer) error
  24. peers map[uuid.UUID]*WSPeer
  25. }
  26. func NewWSServer() *WSServer {
  27. upgrader := websocket.Upgrader{} // use default options
  28. upgrader.CheckOrigin = func(r *http.Request) bool { return true } // cors-origin '*'
  29. wss := WSServer{upgrader: &upgrader, peers: map[uuid.UUID]*WSPeer{}}
  30. return &wss
  31. }
  32. func (wss *WSServer) Write(name string, sendMsg WSMsg) {
  33. peers := wss.GetPeers(name)
  34. if len(peers) == 0 {
  35. log.Println("WSServer:Write - no sendChans for", name)
  36. return
  37. }
  38. for _, peer := range peers {
  39. peer.Answer(sendMsg)
  40. }
  41. }
  42. func (wss *WSServer) GetPeers(searchname string) []*WSPeer {
  43. cons := []*WSPeer{}
  44. for key, con := range wss.peers {
  45. username := key.String()
  46. if username == searchname {
  47. cons = append(cons, con)
  48. }
  49. }
  50. return cons
  51. }
  52. func (wss *WSServer) CleanPeer(name uuid.UUID) {
  53. delete(wss.peers, name)
  54. }
  55. func (wss *WSServer) RegisterHandler(handlers []Handler) {
  56. if wss.handler == nil {
  57. wss.handler = map[string]func(WSMsg, *WSPeer) error{}
  58. }
  59. for _, handler := range handlers {
  60. log.Println("Registrering handler for MType: ", handler.MType)
  61. wss.handler[handler.MType] = handler.Callback
  62. }
  63. }
  64. type rawMsg struct {
  65. mtype int
  66. p []byte
  67. }
  68. type WSPeer struct {
  69. uuid uuid.UUID
  70. conn *websocket.Conn
  71. answers *chan WSMsg
  72. }
  73. func (peer *WSPeer) WorkAnswers() {
  74. for {
  75. msg := WSMsg{}
  76. msg = <-*peer.answers
  77. fmt.Println("sending on: ", peer.uuid)
  78. peer._answer(msg)
  79. }
  80. }
  81. func (peer *WSPeer) Answer(msg WSMsg) {
  82. if peer.answers == nil {
  83. fmt.Println("empty chan for :", peer.uuid)
  84. }
  85. *peer.answers <- msg
  86. }
  87. func (peer *WSPeer) _answer(msg WSMsg) error {
  88. log.Println("sending message: ", msg)
  89. jsonMsg, err := json.Marshal(msg)
  90. if err != nil {
  91. log.Printf("write: err marshalling \"%s\"\n%s", msg, err.Error())
  92. return err
  93. }
  94. err = peer.conn.WriteMessage(websocket.BinaryMessage, jsonMsg)
  95. if err != nil {
  96. log.Println("write:", err)
  97. return err
  98. }
  99. return nil
  100. }
  101. func (wss WSServer) WebsocketHandler(w http.ResponseWriter, r *http.Request) {
  102. log.Println("request on /ws from", r.RemoteAddr)
  103. c, err := wss.upgrader.Upgrade(w, r, nil)
  104. if err != nil {
  105. log.Print("upgrade:", err)
  106. return
  107. }
  108. defer c.Close()
  109. peer := WSPeer{uuid: uuid.New(), conn: c}
  110. bla := make(chan WSMsg, 10)
  111. peer.answers = &bla
  112. go peer.WorkAnswers()
  113. wss.peers[peer.uuid] = &peer
  114. for {
  115. _, message, err := c.ReadMessage()
  116. if err != nil {
  117. log.Println("read:", err)
  118. break
  119. }
  120. var rawmsg map[string]interface{}
  121. err = json.Unmarshal(message, &rawmsg)
  122. msg := WSMsg{}
  123. for k, v := range rawmsg {
  124. switch k {
  125. case "id":
  126. id := v.(float64)
  127. msg.ID = int(id)
  128. case "mtype":
  129. msg.Mtype = v.(string)
  130. case "payload":
  131. payload, ok := v.(string)
  132. if !ok {
  133. payloadmap := v.(map[string]interface{})
  134. var payloadbyte []byte
  135. payloadbyte, err = json.Marshal(payloadmap)
  136. payload = string(payloadbyte)
  137. }
  138. msg.Payload = payload
  139. }
  140. }
  141. if err != nil {
  142. log.Println("error unmarshalling message:", err)
  143. log.Printf("recv: %s", message)
  144. return
  145. }
  146. handler, ok := wss.handler[msg.Mtype]
  147. if !ok {
  148. log.Println("unknown MType", msg.Mtype)
  149. keys := ""
  150. for k := range wss.handler {
  151. keys = fmt.Sprintf("%s\n%s", keys, k)
  152. }
  153. log.Println("available MTypes:\n", keys, "----")
  154. remsg := WSMsg{Mtype: "UNKNOWN", Payload: ""}
  155. peer.Answer(remsg)
  156. return
  157. }
  158. err = handler(msg, &peer)
  159. if err != nil {
  160. log.Println("handler failed", err.Error())
  161. //writeChan <- msg
  162. }
  163. log.Println("message handled")
  164. }
  165. }