You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

server.go 2.9KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  1. package main
  2. import (
  3. "encoding/binary"
  4. "flag"
  5. "fmt"
  6. "log"
  7. "net"
  8. "net/http"
  9. "strings"
  10. "time"
  11. "github.com/gorilla/websocket"
  12. )
  13. var (
  14. bind = flag.String("bind", "[::1]:13254", "Bind address")
  15. )
  16. var upgrader = websocket.Upgrader{
  17. HandshakeTimeout: time.Second * 30,
  18. ReadBufferSize: 2048,
  19. WriteBufferSize: 2048,
  20. CheckOrigin: func(r *http.Request) bool {
  21. return true
  22. },
  23. }
  24. func (p *RelayServer) wserror(w http.ResponseWriter, err error) {
  25. h := w.Header()
  26. h.Set("Content-Type", "text/plain; charset=utf8")
  27. w.WriteHeader(500)
  28. fmt.Fprintf(w, "Error: %s", err)
  29. }
  30. func (p *RelayServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  31. c, err := upgrader.Upgrade(w, r, nil)
  32. if err != nil {
  33. log.Println("ws-upgrade:", err)
  34. p.wserror(w, err)
  35. return
  36. }
  37. defer c.Close()
  38. sender := make(chan []byte, 16)
  39. ident := p.add(sender)
  40. //log.Println("wsrelay", ident, "connected")
  41. pfx := make([]byte, 4)
  42. binary.LittleEndian.PutUint32(pfx, ident)
  43. defer p.remove(ident)
  44. exit := append(pfx, 0)
  45. defer p.broadcast(ident, exit)
  46. //defer log.Println("wsrelay", ident, "disconnected")
  47. ping := 10 * time.Second
  48. timeout := 15 * time.Second
  49. sender <- append([]byte{0, 0, 0, 0}, pfx...)
  50. //c.SetReadLimit(maxMessageSize)
  51. end := make(chan struct{}, 2)
  52. go func() {
  53. var err error
  54. ticker := time.NewTicker(ping)
  55. defer ticker.Stop()
  56. for err == nil {
  57. select {
  58. case <-ticker.C:
  59. c.SetWriteDeadline(time.Now().Add(timeout))
  60. err = c.WriteMessage(websocket.PingMessage, nil)
  61. case msg, ok := <-sender:
  62. if !ok {
  63. break
  64. }
  65. c.SetWriteDeadline(time.Now().Add(timeout))
  66. err = c.WriteMessage(websocket.BinaryMessage, msg)
  67. }
  68. }
  69. if err != nil {
  70. log.Println("wsrelay send:", err)
  71. }
  72. end <- struct{}{}
  73. }()
  74. c.SetPongHandler(func(string) error {
  75. c.SetReadDeadline(time.Now().Add(timeout))
  76. return nil
  77. })
  78. go func() {
  79. var err error
  80. for {
  81. c.SetReadDeadline(time.Now().Add(timeout))
  82. mt, buf, err_ := c.ReadMessage()
  83. if err_ != nil {
  84. err = err_
  85. break
  86. } else if mt == websocket.CloseMessage {
  87. break
  88. } else if mt != websocket.BinaryMessage {
  89. log.Printf("wsrelay: Received wrong message: %v", mt)
  90. break
  91. }
  92. s := buf
  93. if len(s) > 250 {
  94. s = s[:250]
  95. }
  96. p.broadcast(ident, append(pfx, s...))
  97. }
  98. if err != nil {
  99. if _, ok := err.(*websocket.CloseError); !ok {
  100. log.Println("wsrelay receive:", err)
  101. }
  102. }
  103. end <- struct{}{}
  104. }()
  105. <-end
  106. }
  107. func main() {
  108. flag.Parse()
  109. log.Println("wsrelay")
  110. n := "tcp"
  111. if strings.HasPrefix(*bind, "/") {
  112. n = "unix"
  113. }
  114. l, err := net.Listen(n, *bind)
  115. if err != nil {
  116. log.Fatalln("Listen error:", err)
  117. }
  118. ws := &http.Server{
  119. ReadTimeout: 60 * time.Second,
  120. WriteTimeout: 60 * time.Second,
  121. MaxHeaderBytes: 1 << 16,
  122. Handler: &RelayServer{
  123. clients: make(map[uint32]RelayClient),
  124. },
  125. }
  126. if err := ws.Serve(l); err != nil {
  127. log.Fatalln("Server error:", err)
  128. }
  129. }