Golang Conn.WriteControl方法代碼示例

func connectToAgent(timeout int, onReady func(*websocket.Conn)) (*websocket.Conn, error) {
	var ws *websocket.Conn
	var err error

	t := 0
	for {
		if t > timeout {
			return nil, errors.New("Connection to Agent : timeout reached")
		}

		ws, err = newClient()
		if err == nil {
			break
		}
		time.Sleep(1 * time.Second)
		t++
	}

	ready := false
	h := func(message string) error {
		err := ws.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(time.Second))
		if err != nil {
			return err
		}
		if !ready {
			ready = true
			onReady(ws)
		}
		return nil
	}
	ws.SetPingHandler(h)

	return ws, nil
}

func AddWSSink(receivedChan chan []byte, port string, path string) (*websocket.Conn, <-chan struct{}) {
	connectionDroppedChannel := make(chan struct{}, 1)

	var ws *websocket.Conn

	ip, _ := localip.LocalIP()
	fullURL := "ws://" + ip + ":" + port + path

	Eventually(func() error {
		var err error
		ws, _, err = websocket.DefaultDialer.Dial(fullURL, http.Header{})
		return err
	}, 5, 1).ShouldNot(HaveOccurred(), fmt.Sprintf("Unable to connect to server at %s.", fullURL))

	ws.SetPingHandler(func(message string) error {
		ws.WriteControl(websocket.PongMessage, []byte(message), time.Time{})
		return nil
	})

	go func() {
		for {
			_, data, err := ws.ReadMessage()
			if err != nil {
				close(connectionDroppedChannel)
				close(receivedChan)
				return
			}

			receivedChan <- data
		}

	}()

	return ws, connectionDroppedChannel
}

// Pick requests off the RemoteServer queue and send them into the tunnel
func wsWriter(rs *remoteServer, ws *websocket.Conn, ch chan int) {
	var req *remoteRequest
	var err error
	for {
		// fetch a request
		select {
		case req = <-rs.requestQueue:
			// awesome...
		case _ = <-ch:
			// time to close shop
			rs.log.Info("WS closing on signal", "ws", wsp(ws))
			ws.Close()
			return
		}
		//log.Printf("WS->%s#%d start %s", req.token, req.id, req.info)
		// See whether the request has already expired
		if req.deadline.Before(time.Now()) {
			req.replyChan <- responseBuffer{
				err: errors.New("Timeout before forwarding the request"),
			}
			req.log.Info("WS   SND timeout before sending", "ago",
				time.Now().Sub(req.deadline).Seconds())
			continue
		}
		// write the request into the tunnel
		ws.SetWriteDeadline(time.Now().Add(time.Minute))
		var w io.WriteCloser
		w, err = ws.NextWriter(websocket.BinaryMessage)
		// got an error, reply with a "hey, retry" to the request handler
		if err != nil {
			break
		}
		// write the request Id
		_, err = fmt.Fprintf(w, "%04x", req.id)
		if err != nil {
			break
		}
		// write the request itself
		_, err = req.buffer.WriteTo(w)
		if err != nil {
			break
		}
		// done
		err = w.Close()
		if err != nil {
			break
		}
		req.log.Info("WS   SND", "info", req.info)
	}
	// tell the sender to retry the request
	req.replyChan <- responseBuffer{err: RetryError}
	req.log.Info("WS error causes retry")
	// close up shop
	ws.WriteControl(websocket.CloseMessage, nil, time.Now().Add(5*time.Second))
	time.Sleep(2 * time.Second)
	ws.Close()
}

func closeWithErr(log lager.Logger, conn *websocket.Conn, code int, reason string) {
	err := conn.WriteControl(
		websocket.CloseMessage,
		websocket.FormatCloseMessage(code, reason),
		time.Time{},
	)

	if err != nil {
		log.Error("failed-to-close-websocket-connection", err)
	}
}

func (m *multiplexer) routeMessages(ws *websocket.Conn) {
	stopSignal := make(chan bool, 1)

	// Read messages from backend
	go func(stop chan<- bool) {
		for {
			msgType, msg, err := ws.ReadMessage()
			if err != nil {
				m.shutdown(stop)
				return
			}

			if msgType != websocket.TextMessage {
				continue
			}
			message := common.ParseMessage(string(msg))

			m.frontendMu.RLock()
			frontendChan, ok := m.frontendChans[message.Key]
			if ok {
				frontendChan <- message
			}
			m.frontendMu.RUnlock()

			if !ok && message.Type != common.Close {
				m.sendClose(message.Key)
			}
		}
	}(stopSignal)

	// Write messages to backend
	go func(stop <-chan bool) {
		ticker := time.NewTicker(time.Second * 5)
		defer ticker.Stop()
		for {
			select {
			case message, ok := <-m.messagesToBackend:
				if !ok {
					return
				}
				err := ws.WriteMessage(websocket.TextMessage, []byte(message))
				if err != nil {
					log.WithFields(log.Fields{"error": err, "msg": message}).Error("Could not write message.")
				}

			case <-ticker.C:
				ws.WriteControl(websocket.PingMessage, []byte(""), time.Now().Add(time.Second))

			case <-stop:
				return
			}
		}
	}(stopSignal)
}

// ping will loop infinitely and send ping messages over ws. If a write takes
// longer than writeWait, it will remove ws from the connection pool.
func ping(ws *websocket.Conn, id id) {
	b := make([]byte, 2)
	for range time.Tick(pingPeriod) {
		rand.Read(b)
		deadline := time.Now().Add(writeWait)
		if err := ws.WriteControl(websocket.PingMessage, b, deadline); err != nil {
			die(id, err)
			return
		}
		log.Debugf(appengine.BackgroundContext(), "ws: ping %x for %s(%d %q) sent", b, id.kind, id.intID, id.stringID)
	}
}

func (w *WebsocketServer) streamLogs(appId string, ws *websocket.Conn) {
	defer ws.WriteControl(websocket.CloseMessage, []byte{}, time.Now().Add(100*time.Millisecond))
	websocketSink := sinks.NewWebsocketSink(
		appId,
		w.logger,
		ws,
		w.keepAliveInterval,
		w.bufferSize,
	)
	w.logger.Debugf("WebsocketServer: Requesting a wss sink for app %s", websocketSink.AppId())
	w.sinkManager.RegisterSink(websocketSink, false)
}

func stopSend(con *websocket.Conn, done chan int) {
	//send close message
	fmt.Println("wsSend: got stop mssage")
	//fmt.Println("wsSend: closing send")
	//its okay if we continue recieving a bit before the close message is read.
	//channel was closed, send close frame and exit
	fmt.Println("wsSend: sending close frame")
	err := con.WriteControl(websocket.CloseMessage, nil, time.Now().Add(3*time.Second))
	if err != nil { //if theres an error sending, print err, exit
		fmt.Println("wsSend: control frame send err:", err)
	}
	close(done)
}

func ping(ws *websocket.Conn, done chan struct{}) {
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			if err := ws.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait)); err != nil {
				log.Println("ping:", err)
			}
		case <-done:
			return
		}
	}
}

func pingDispatcher(s *websocket.Conn) {
	for {
		if s == nil {
			log.Info("can't ping as connection closed")
			return
		}
		if err := s.WriteControl(websocket.PingMessage, []byte(fmt.Sprintf("server at %s", time.Now().String())), time.Now().Add(WS_TIMEOUT)); err == nil {
			log.Tracef("ping sent")
		} else {
			log.Error("ping failed")
			return
		}
		time.Sleep(WS_TIMEOUT)
	}
}

// pinger simulates the websocket to keep it alive with ping messages.
func pinger(conn *websocket.Conn, done chan struct{}) {
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()

	for { // blocking loop with select to wait for stimulation.
		select {
		case <-ticker.C:
			if err := conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait)); err != nil {
				conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, err.Error()), time.Time{})
				return
			}
		case <-done:
			return // clean up this routine.
		}
	}
}

// pumpStdout handles reading data from stdout of the process and writing
// it to websocket connection.
func pumpStdout(conn *websocket.Conn, stdout io.Reader, done chan struct{}) {
	go pinger(conn, done)
	defer func() {
		conn.Close()
		close(done) // make sure to close the pinger when we are done.
	}()

	s := bufio.NewScanner(stdout)
	for s.Scan() {
		conn.SetWriteDeadline(time.Now().Add(writeWait))
		if err := conn.WriteMessage(websocket.TextMessage, bytes.TrimSpace(s.Bytes())); err != nil {
			break
		}
	}
	if s.Err() != nil {
		conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, s.Err().Error()), time.Time{})
	}
}

// reader is the guts of this package. It takes the stdin and stdout pipes
// of the cmd we created in ServeWS and pipes them between the client and server
// over websockets.
func reader(conn *websocket.Conn, stdout io.ReadCloser, stdin io.WriteCloser) {
	// Setup our connection's websocket ping/pong handlers from our const values.
	conn.SetReadLimit(maxMessageSize)
	conn.SetReadDeadline(time.Now().Add(pongWait))
	conn.SetPongHandler(func(string) error { conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	tickerChan := make(chan bool)
	defer close(tickerChan) // make sure to close the ticker when we are done.
	go ticker(conn, tickerChan)

	for {
		msgType, r, err := conn.NextReader()
		if err != nil {
			if msgType == -1 {
				return // we got a disconnect from the client. We are good to close.
			}
			conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
			return
		}

		w, err := conn.NextWriter(msgType)
		if err != nil {
			conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
			return
		}

		if _, err := io.Copy(stdin, r); err != nil {
			conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
			return
		}

		go func() {
			if _, err := io.Copy(w, stdout); err != nil {
				conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
				return
			}
			if err := w.Close(); err != nil {
				conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
				return
			}
		}()
	}
}

func wsSetPingHandler(t *WSTunnelServer, ws *websocket.Conn, rs *remoteServer) {
	// timeout handler sends a close message, waits a few seconds, then kills the socket
	timeout := func() {
		ws.WriteControl(websocket.CloseMessage, nil, time.Now().Add(1*time.Second))
		time.Sleep(5 * time.Second)
		rs.log.Info("WS closing due to ping timeout", "ws", wsp(ws))
		ws.Close()
	}
	// timeout timer
	timer := time.AfterFunc(t.WSTimeout, timeout)
	// ping handler resets last ping time
	ph := func(message string) error {
		timer.Reset(t.WSTimeout)
		ws.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(t.WSTimeout/3))
		// update lastActivity
		rs.lastActivity = time.Now()
		return nil
	}
	ws.SetPingHandler(ph)
}

func WSHandler(c *websocket.Conn) {
	closed := false
	ret := websocket.CloseError{
		Code: websocket.CloseNormalClosure,
	}
	const writeWait = 10 * time.Second

	defer func() {
		log.Infof("Close websocket [%s]", ret.Error())
		c.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(ret.Code, ret.Text), time.Now().Add(writeWait))
		c.Close()
	}()

	go func() {
		for {
			_, r, err := c.NextReader()
			if err != nil {
				break
			}
			io.Copy(ioutil.Discard, r)
		}
		log.Infof("Pipe closed")

		closed = true
	}()

	for !closed {
		signal := M{
			"pm2.5": atomic.LoadInt32(&currentPM25),
			"pm10":  atomic.LoadInt32(&currentPM100),
		}
		if err := c.WriteJSON(signal); err != nil && err != websocket.ErrCloseSent {
			if _, ok := err.(net.Error); !ok {
				log.Warnf("Error writing to client: %s", err.Error())
			}
			closed = true
		}
		time.Sleep(2 * time.Second)
	}
}