Golang websocket.FormatCloseMessage函数代码示例

func main() {
	flag.Parse()
	log.SetFlags(0)

	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	u := url.URL{Scheme: "ws", Host: *addr, Path: "/echo"}
	log.Printf("connecting to %s", u.String())

	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		log.Fatal("dial:", err)
	}
	defer c.Close()

	done := make(chan struct{})

	go func() {
		defer c.Close()
		defer close(done)
		for {
			_, message, err := c.ReadMessage()
			if err != nil {
				log.Println("read:", err)
				return
			}
			log.Printf("recv: %s", message)
		}
	}()

	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	for {
		select {
		case t := <-ticker.C:
			err := c.WriteMessage(websocket.TextMessage, []byte(t.String()))
			if err != nil {
				log.Println("write:", err)
				return
			}
		case <-interrupt:
			log.Println("interrupt")
			// To cleanly close a connection, a client should send a close
			// frame and wait for the server to close the connection.
			err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
			if err != nil {
				log.Println("write close:", err)
				return
			}
			select {
			case <-done:
			case <-time.After(time.Second):
			}
			c.Close()
			return
		}
	}
}

// Sends dispose signals to all peers and deletes the room record.
func (room *Room) dispose(status int) {
	if status == 0 {
		status = websocket.CloseNormalClosure
	}

	// Close all Websocket connections.
	for peer := range room.peers {
		peer.ws.WriteControl(
			websocket.CloseMessage,
			websocket.FormatCloseMessage(status, "Room disposed"),
			time.Time{})

		delete(room.peers, peer)
	}

	// Close all the room channels.
	close(room.broadcastQueue)
	close(room.register)
	close(room.unregister)

	delete(rooms, room.Id)

	// Delete the room record from the DB.
	db := dbPool.Get()
	defer db.Close()

	db.Delete(config.CachePrefixRoom + room.Id)
}

// Start connects to chat and starts bot.
func (c *Client) Start() {
	conn, _, err := websocket.DefaultDialer.Dial(c.WebSocketURL, nil)
	if err != nil {
		log.Fatal(err)
	}

	defer conn.Close()
	c.conn = conn

	done := make(chan struct{})
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	go c.processMessages(done)
	for {
		select {
		case <-interrupt:
			err = conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
			if err != nil {
				log.Fatal(err)
			}
			select {
			case <-done:
			case <-time.After(time.Second):
			}
			return
		}
	}
}

func (c *conn) close(closeCode int, closeErr error) error {
	var err error
	var errMsg string

	if closeErr != nil {
		errMsg = closeErr.Error()
		c.server.logError(closeErr)

		// Attempt to set close code from error message
		errMsgLen := len(errMsg)
		if errMsgLen >= 21 && errMsg[:17] == "websocket: close " {
			closeCode, err = strconv.Atoi(errMsg[17:21])
			if errMsgLen > 21 {
				errMsg = errMsg[22:]
			}
		}
	}

	// Default close code
	if closeCode == 0 {
		closeCode = websocket.CloseNoStatusReceived
	}

	// Send close message
	closeMessage := websocket.FormatCloseMessage(closeCode, errMsg)
	deadline := time.Now().Add(time.Second)
	err = c.conn.WriteControl(websocket.CloseMessage, closeMessage, deadline)

	// Kill and remove connection
	c.closeChan <- closeSignal{}
	c.server.connSet.remove(c)
	return err
}

// StartMockServer starts a mock websocket server.
func StartMockServer(t *testing.T, closeWS <-chan bool) (*httptest.Server, chan<- string, <-chan string, <-chan error, error) {
	serverChan := make(chan string)
	requestsChan := make(chan string)
	errChan := make(chan error)

	upgrader := websocket.Upgrader{ReadBufferSize: 1024, WriteBufferSize: 1024}
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ws, err := upgrader.Upgrade(w, r, nil)
		go func() {
			<-closeWS
			ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
		}()
		if err != nil {
			errChan <- err
		}
		go func() {
			_, msg, err := ws.ReadMessage()
			if err != nil {
				errChan <- err
			} else {
				requestsChan <- string(msg)
			}
		}()
		for str := range serverChan {
			err := ws.WriteMessage(websocket.TextMessage, []byte(str))
			if err != nil {
				errChan <- err
			}
		}
	})

	server := httptest.NewTLSServer(handler)
	return server, serverChan, requestsChan, errChan, nil
}

func (w *WebsocketServer) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	w.logger.Debug("WebsocketServer.ServeHTTP: starting")
	var handler wsHandler
	var err error

	paths := strings.Split(request.URL.Path, "/")
	endpointName := paths[1]

	if endpointName == "firehose" {
		handler, err = w.firehoseHandler(paths, writer, request)
	} else {
		handler, err = w.appHandler(paths, writer, request)
	}

	if err != nil {
		w.logger.Errorf("WebsocketServer.ServeHTTP: %s", err.Error())
		return
	}

	ws, err := gorilla.Upgrade(writer, request, nil, 1024, 1024)
	if err != nil {
		w.logger.Errorf("WebsocketServer.ServeHTTP: Upgrade error (returning 400): %s", err.Error())
		http.Error(writer, err.Error(), 400)
		return
	}

	defer func() {
		ws.WriteControl(gorilla.CloseMessage, gorilla.FormatCloseMessage(gorilla.CloseNormalClosure, ""), time.Time{})
		ws.Close()
	}()

	handler(ws)
}

func WebsocketSendStream(conn *websocket.Conn, r io.Reader) chan bool {
	ch := make(chan bool)

	go func(conn *websocket.Conn, r io.Reader) {
		in := ReaderToChannel(r)
		for {
			buf, ok := <-in
			if !ok {
				break
			}

			w, err := conn.NextWriter(websocket.BinaryMessage)
			if err != nil {
				Debugf("got error getting next writer %s", err)
				break
			}

			_, err = w.Write(buf)
			w.Close()
			if err != nil {
				Debugf("got err writing %s", err)
				break
			}
		}
		closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
		conn.WriteMessage(websocket.CloseMessage, closeMsg)
		ch <- true
	}(conn, r)

	return ch
}

func fakeFirehoseHandler(rw http.ResponseWriter, r *http.Request) {
	defer GinkgoRecover()
	authorization := r.Header.Get("Authorization")

	if authorization != "bearer good-token" {
		log.Printf("Bad token passed to firehose: %s", authorization)
		rw.WriteHeader(403)
		r.Body.Close()
		return
	}

	upgrader := websocket.Upgrader{
		CheckOrigin: func(*http.Request) bool { return true },
	}

	ws, _ := upgrader.Upgrade(rw, r, nil)

	defer ws.Close()
	defer ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Time{})

	for envelope := range fakeFirehoseInputChan {
		buffer, err := proto.Marshal(envelope)
		Expect(err).NotTo(HaveOccurred())
		err = ws.WriteMessage(websocket.BinaryMessage, buffer)
		Expect(err).NotTo(HaveOccurred())
	}
}

// Close the Base connection. Closes the send Handler and all channels used
// Since all channels are either internal or channels this middleware is sending on.
func (c *Connection) Close(closeCode int) error {
	c.disconnectSend <- true
	//TODO look for a better way to unblock the reader
	c.ws.SetReadDeadline(time.Now())

	// Send close message to the client
	c.log("Sending close message to client", LEVEL_DEBUG)
	c.ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(closeCode, ""), time.Now().Add(c.WriteWait))

	// If the connection can not be closed, return the error
	c.log("Closing websocket connection", LEVEL_DEBUG)
	if err := c.ws.Close(); err != nil {
		c.log("Connection could not be closed: %s", LEVEL_ERROR, err.Error())
		return err
	}

	// Send disconnect message to the next handler
	c.log("Sending disconnect to handler", LEVEL_DEBUG)
	c.Done <- true

	// Close disconnect and error channels this connection was sending on
	close(c.Done)
	close(c.Error)

	return nil
}

func (c *migrationFields) disconnect() {
	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")

	c.controlLock.Lock()
	if c.controlConn != nil {
		c.controlConn.WriteMessage(websocket.CloseMessage, closeMsg)
		c.controlConn = nil /* don't close twice */
	}
	c.controlLock.Unlock()

	/* Below we just Close(), which doesn't actually write to the
	 * websocket, it just closes the underlying connection. If e.g. there
	 * is still a filesystem transfer going on, but the other side has run
	 * out of disk space, writing an actual CloseMessage here will cause
	 * gorilla websocket to panic. Instead, we just force close this
	 * connection, since we report the error over the control channel
	 * anyway.
	 */
	if c.fsConn != nil {
		c.fsConn.Close()
	}

	if c.criuConn != nil {
		c.criuConn.Close()
	}
}

func (pipe *Pipe) Close(msg string) {
	pipe.logger.Printf("pipe: closing (msg=%q)", msg)
	pipe.conn.WriteMessage(
		websocket.CloseMessage,
		websocket.FormatCloseMessage(websocket.CloseNormalClosure, msg),
	)
	pipe.conn.Close()
}

func NewFakeFirehoseInAppMode(validToken, appName string) *FakeFirehose {
	return &FakeFirehose{
		AppMode:      true,
		AppName:      appName,
		validToken:   validToken,
		closeMessage: websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
	}
}

func (c *Connection) SendClose(closeCode int, text string) {
	// XXX: we're allowed to send control frames from any thread, so it
	// might be easier to write the message directly to the web socket.
	c.send <- message{
		websocket.CloseMessage,
		websocket.FormatCloseMessage(closeCode, text),
	}
}

func (ws *wsconn) Close() error {
	ws.connClosing = true
	ws.wlock.Lock()
	ws.wsc.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add((time.Second * 5)))
	ws.wlock.Unlock()
	err := ws.wsc.Close()
	return err
}

func shouldCloseConnection(actual interface{}, _ ...interface{}) string {
	conn := actual.(*websocket.Conn)
	err := conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
	if err != nil {
		return "Websocket connection didn't close properly: " + err.Error()
	}
	return ""
}