github.com/artisanhe/tools@v1.0.1-0.20210607022958-19a8fef2eb04/websocket-wrapper/ws_server_handler.go (about)

     1  package ww
     2  
     3  import (
     4  	"net/http"
     5  
     6  	"github.com/gorilla/websocket"
     7  )
     8  
     9  type WSServerInterface interface {
    10  	OnMessage(WSHandle, []byte, error)
    11  	OnClose(WSHandle)
    12  }
    13  
    14  type WSServerHandler struct {
    15  	On WSServerInterface
    16  }
    17  
    18  // OnHandler contains a connection and open a routine for recv
    19  func (ws *WSServerHandler) OnHandler(w http.ResponseWriter, r *http.Request) {
    20  	// check ,if the interface is nil
    21  	if ws.On == nil {
    22  		panic("WSServerHandler ")
    23  	}
    24  	// use the default options
    25  	upgrader := websocket.Upgrader{}
    26  	c, err := upgrader.Upgrade(w, r, nil)
    27  	if err != nil {
    28  		Error("OnHandler,upgrade the general request connection to a websocket connection error:%s", err.Error())
    29  		return
    30  	}
    31  	defer c.Close()
    32  	handle := wss.saveConn(c)
    33  	writeChan := make(chan *WSMessage, 10)
    34  	wss.saveChan(handle, writeChan)
    35  	go func() {
    36  		closeMsg := WSMessage{Type: WSMessageTypeClose}
    37  		for {
    38  			mt, message, err := c.ReadMessage()
    39  			if err != nil {
    40  				remoteAddr := c.RemoteAddr().String()
    41  				Error("OnHandler,ReadMessage error:%s,handle:%d,remote addr:%s,message type:%d",
    42  					err.Error(), handle, remoteAddr, mt)
    43  				ws.On.OnMessage(handle, nil, err)
    44  				// additional send close msg
    45  				writeChan <- &closeMsg
    46  				break
    47  			}
    48  			switch mt {
    49  			case websocket.TextMessage:
    50  				ws.On.OnMessage(handle, message, nil)
    51  			case websocket.BinaryMessage:
    52  				//Debug("recv keep-alive message")
    53  				// do nothing,binray message only for keep-alive message
    54  			default:
    55  				remoteAddr := c.RemoteAddr().String()
    56  				Error("OnHandler,ReadMessage handle:%d,remote addr:%s,unsupport message type:%d",
    57  					handle, remoteAddr, mt)
    58  				// additional send close msg
    59  				writeChan <- &closeMsg
    60  			}
    61  		}
    62  	}()
    63  WriteDone:
    64  	for {
    65  		select {
    66  		case msg, ok := <-writeChan:
    67  			if !ok {
    68  				break WriteDone
    69  			}
    70  			switch {
    71  			case msg.Type == WSMessageTypeBusiness:
    72  				err = c.WriteMessage(websocket.TextMessage, msg.Data)
    73  				if err != nil {
    74  					remoteAddr := c.RemoteAddr().String()
    75  					Error("OnHandler,WriteMessage error:%s,handle:%d,remote addr:%s,msg type:%d",
    76  						err.Error(), handle, remoteAddr, msg.Type)
    77  					break WriteDone
    78  				}
    79  			case msg.Type == WSMessageTypeClose:
    80  				err := c.WriteMessage(websocket.CloseMessage,
    81  					websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
    82  				if err != nil {
    83  					remoteAddr := c.RemoteAddr().String()
    84  					Error("OnHandler,WriteMessage error:%s,handle:%d,remote addr:%s,msg type:%d",
    85  						err.Error(), handle, remoteAddr, msg.Type)
    86  				}
    87  				break WriteDone
    88  			}
    89  		}
    90  	}
    91  	// clean resources
    92  	close(writeChan)
    93  	wss.removeChan(handle)
    94  	wss.removeConn(c)
    95  	ws.On.OnClose(handle)
    96  }