webs.go 3.2KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. package chocococon
  2. import (
  3. "github.com/pkg/errors"
  4. "strings"
  5. "fmt"
  6. "github.com/gorilla/websocket"
  7. "sync"
  8. "time"
  9. )
  10. func init() {
  11. cmds["startws"] = startws
  12. cmds["closews"] = closews
  13. }
  14. var (
  15. currentServer *(chan<- struct{})
  16. currentServerLock sync.Mutex
  17. )
  18. func startws(params []string) error {
  19. if len(params) != 1 {
  20. return errors.New("Wrong parameter count")
  21. }
  22. if !strings.HasPrefix(params[0], "url=") {
  23. return errors.New("Wrong parameter on position for URL")
  24. }
  25. currentServerLock.Lock()
  26. if currentServer != nil {
  27. fmt.Print("Shutting down current websocket...")
  28. *currentServer <- struct{}{}
  29. currentServer = nil
  30. fmt.Print("OK!\n")
  31. }
  32. currentServerLock.Unlock()
  33. fmt.Println("Opening connection to websocket...")
  34. return websocket_client(strings.TrimPrefix(params[0], "url="))
  35. }
  36. func closews(params []string) error {
  37. currentServerLock.Lock()
  38. *currentServer <- struct{}{}
  39. currentServer = nil
  40. currentServerLock.Unlock()
  41. return nil
  42. }
  43. func websocket_client(url string) error {
  44. fmt.Println("Connecting to URL...")
  45. uDial := &websocket.Dialer{}
  46. uDial.HandshakeTimeout = 20 * time.Second
  47. c, _, err := uDial.Dial(url, nil)
  48. if err != nil {
  49. return err
  50. }
  51. fmt.Println("Setting up IO Routines...")
  52. readShutdown := make(chan struct{})
  53. reader := readLoop(c, readShutdown)
  54. writeShutdown := make(chan struct{})
  55. writer := writeLoop(c, writeShutdown)
  56. fmt.Println("Main Server started...")
  57. shutdownChan := make(chan struct{})
  58. go func() {
  59. ticker := time.Tick(10 * time.Second)
  60. for {
  61. select {
  62. case msg := <-reader:
  63. if strings.TrimSpace(string(msg)) == "" {
  64. //fmt.Printf("Command empty, skipping")
  65. continue
  66. }
  67. fmt.Printf("Received Command: %sExeccuting...\n", msg)
  68. err := exec(makeTokens(string(msg)))
  69. if err != nil {
  70. fmt.Println("Error")
  71. fmt.Println(": ", err)
  72. writer<-[]byte(fmt.Sprintf("Error: ", err))
  73. continue
  74. }
  75. fmt.Println("OK")
  76. writer<-[]byte("OK")
  77. case <-ticker:
  78. fmt.Println("I am still alive")
  79. case <-shutdownChan:
  80. readShutdown<-struct{}{}
  81. writeShutdown<-struct{}{}
  82. return
  83. }
  84. }
  85. }()
  86. currentServerLock.Lock()
  87. if currentServer != nil {
  88. fmt.Println("Websocket Server already running")
  89. return errors.New("Double Invocation of WS Server")
  90. }
  91. {
  92. s := chan <- struct{}(shutdownChan)
  93. currentServer = &(s)
  94. }
  95. currentServerLock.Unlock()
  96. return nil
  97. }
  98. func readLoop(conn *websocket.Conn, stop <-chan struct{}) (<-chan []byte) {
  99. outChan := make(chan []byte)
  100. go func(stop <-chan struct{}) {
  101. for {
  102. select {
  103. case <-stop:
  104. return
  105. default:
  106. _, m, err := conn.ReadMessage()
  107. if strings.TrimSpace(string(m)) == "" {
  108. continue
  109. }
  110. if err != nil {
  111. fmt.Printf("Error on WS read: ", err)
  112. }
  113. outChan <- m
  114. }
  115. }
  116. }(stop)
  117. return outChan
  118. }
  119. func writeLoop(conn *websocket.Conn, stop <-chan struct{}) (chan<- []byte) {
  120. outChan := make(chan []byte)
  121. go func(stop <-chan struct{}) {
  122. for {
  123. select {
  124. case <-stop:
  125. return
  126. case msg := <-outChan:
  127. if strings.TrimSpace(string(msg)) == "" {
  128. continue
  129. }
  130. err := conn.WriteMessage(websocket.BinaryMessage, msg)
  131. if err != nil {
  132. fmt.Printf("Error on WS write: ", err)
  133. }
  134. }
  135. }
  136. }(stop)
  137. return outChan
  138. }