github.com/glide-im/glide@v1.6.0/pkg/conn/ws_server.go (about)

     1  package conn
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/gorilla/websocket"
     6  	"net/http"
     7  	"time"
     8  )
     9  
    10  type WsServerOptions struct {
    11  	ReadTimeout  time.Duration
    12  	WriteTimeout time.Duration
    13  }
    14  
    15  type WsServer struct {
    16  	options  *WsServerOptions
    17  	upgrader websocket.Upgrader
    18  	handler  ConnectionHandler
    19  }
    20  
    21  // NewWsServer options can be nil, use default value when nil.
    22  func NewWsServer(options *WsServerOptions) Server {
    23  
    24  	if options == nil {
    25  		options = &WsServerOptions{
    26  			ReadTimeout:  8 * time.Minute,
    27  			WriteTimeout: 8 * time.Minute,
    28  		}
    29  	}
    30  	ws := new(WsServer)
    31  	ws.options = options
    32  	ws.upgrader = websocket.Upgrader{
    33  		ReadBufferSize:  1024,
    34  		WriteBufferSize: 65536,
    35  		CheckOrigin: func(r *http.Request) bool {
    36  			return true
    37  		},
    38  	}
    39  	return ws
    40  }
    41  
    42  func (ws *WsServer) handleWebSocketRequest(writer http.ResponseWriter, request *http.Request) {
    43  
    44  	conn, err := ws.upgrader.Upgrade(writer, request, nil)
    45  	if err != nil {
    46  		// logger.E("upgrade http to ws error", err)
    47  		return
    48  	}
    49  
    50  	proxy := ConnectionProxy{
    51  		conn: NewWsConnection(conn, ws.options),
    52  	}
    53  	ws.handler(proxy)
    54  }
    55  
    56  func (ws *WsServer) SetConnHandler(handler ConnectionHandler) {
    57  	ws.handler = handler
    58  }
    59  
    60  func (ws *WsServer) Run(host string, port int) error {
    61  
    62  	http.HandleFunc("/ws", ws.handleWebSocketRequest)
    63  
    64  	addr := fmt.Sprintf("%s:%d", host, port)
    65  	if err := http.ListenAndServe(addr, nil); err != nil {
    66  		return err
    67  	}
    68  	return nil
    69  }