Golang Conn.SetReadDeadline方法代碼示例

// Read responses from the tunnel and fulfill pending requests
func wsReader(rs *remoteServer, ws *websocket.Conn, wsTimeout time.Duration, ch chan int) {
	var err error
	log_token := cutToken(rs.token)
	// continue reading until we get an error
	for {
		ws.SetReadDeadline(time.Time{}) // no timeout, there's the ping-pong for that
		// read a message from the tunnel
		var t int
		var r io.Reader
		t, r, err = ws.NextReader()
		if err != nil {
			break
		}
		if t != websocket.BinaryMessage {
			err = fmt.Errorf("non-binary message received, type=%d", t)
			break
		}
		// give the sender a fixed time to get us the data
		ws.SetReadDeadline(time.Now().Add(wsTimeout))
		// get request id
		var id int16
		_, err = fmt.Fscanf(io.LimitReader(r, 4), "%04x", &id)
		if err != nil {
			break
		}
		// read request itself, the size is limited by the SetReadLimit on the websocket
		var buf []byte
		buf, err = ioutil.ReadAll(r)
		if err != nil {
			break
		}
		rs.log.Info("WS   RCV", "id", id, "ws", wsp(ws), "len", len(buf))
		// try to match request
		rs.requestSetMutex.Lock()
		req := rs.requestSet[id]
		rs.lastActivity = time.Now()
		rs.requestSetMutex.Unlock()
		// let's see...
		if req != nil {
			rb := responseBuffer{response: bytes.NewBuffer(buf)}
			// try to enqueue response
			select {
			case req.replyChan <- rb:
				// great!
			default:
				rs.log.Info("WS   RCV can't enqueue response", "id", id, "ws", wsp(ws))
			}
		} else {
			rs.log.Info("%s #%d: WS   RCV orphan response", "id", id, "ws", wsp(ws))
		}
	}
	// print error message
	if err != nil {
		rs.log.Info("WS   closing", "token", log_token, "err", err.Error(), "ws", wsp(ws))
	}
	// close up shop
	ch <- 0 // notify sender
	time.Sleep(2 * time.Second)
	ws.Close()
}

func (this *subServer) wsReadPump(clientGone chan struct{}, ws *websocket.Conn) {
	ws.SetReadLimit(this.wsReadLimit)
	ws.SetReadDeadline(time.Now().Add(this.wsPongWait))
	ws.SetPongHandler(func(string) error {
		ws.SetReadDeadline(time.Now().Add(this.wsPongWait))
		return nil
	})

	// if kateway shutdown while there are open ws conns, the shutdown will
	// wait 1m: this.subServer.wsPongWait
	for {
		_, message, err := ws.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				log.Warn("%s: %v", ws.RemoteAddr(), err)
			} else {
				log.Debug("%s: %v", ws.RemoteAddr(), err)
			}

			close(clientGone)
			break
		}

		log.Debug("ws[%s] read: %s", ws.RemoteAddr(), string(message))
	}
}

// New takes a websocket and creates a ShellClient object implementing the
// engines.Shell interface.
func New(ws *websocket.Conn) *ShellClient {
	stdinReader, stdin := ioext.BlockedPipe()
	tellOut := make(chan int, 10)
	tellErr := make(chan int, 10)
	stdout, stdoutWriter := ioext.AsyncPipe(shellconsts.ShellMaxPendingBytes, tellOut)
	stderr, stderrWriter := ioext.AsyncPipe(shellconsts.ShellMaxPendingBytes, tellErr)
	stdinReader.Unblock(shellconsts.ShellMaxPendingBytes)

	s := &ShellClient{
		ws:           ws,
		stdin:        stdin,
		stdout:       stdout,
		stderr:       stderr,
		stdinReader:  stdinReader,
		stdoutWriter: stdoutWriter,
		stderrWriter: stderrWriter,
		done:         make(chan struct{}),
	}

	ws.SetReadLimit(shellconsts.ShellMaxMessageSize)
	ws.SetReadDeadline(time.Now().Add(shellconsts.ShellPongTimeout))
	ws.SetPongHandler(s.pongHandler)

	go s.writeMessages()
	go s.readMessages()
	go s.sendPings()
	go s.sendAck(shellconsts.StreamStdout, tellOut)
	go s.sendAck(shellconsts.StreamStderr, tellErr)

	return s
}

// Run is the block's main loop. Here we listen on the different channels we set up.
func (b *FromWebsocket) Run() {
	var ws *websocket.Conn
	var url string
	to, _ := time.ParseDuration("10s")
	var handshakeDialer = &websocket.Dialer{
		Subprotocols:     []string{"p1", "p2"},
		HandshakeTimeout: to,
	}
	listenWS := make(blocks.MsgChan)
	wsHeader := http.Header{"Origin": {"http://localhost/"}}

	toOut := make(blocks.MsgChan)
	toError := make(chan error)

	for {
		select {

		case msg := <-toOut:
			b.out <- msg

		case ruleI := <-b.inrule:
			var err error
			// set a parameter of the block
			url, err = util.ParseString(ruleI, "url")
			if err != nil {
				b.Error(err)
				continue
			}
			if ws != nil {
				ws.Close()
			}

			ws, _, err = handshakeDialer.Dial(url, wsHeader)
			if err != nil {
				b.Error("could not connect to url")
				break
			}
			ws.SetReadDeadline(time.Time{})
			h := recvHandler{toOut, toError}
			go h.recv(ws, listenWS)

		case err := <-toError:
			b.Error(err)

		case <-b.quit:
			// quit the block
			return
		case o := <-b.queryrule:
			o <- map[string]interface{}{
				"url": url,
			}
		case in := <-listenWS:
			b.out <- in
		}
	}
}

// Run is the block's main loop. Here we listen on the different channels we set up.
func (b *FromWebsocket) Run() {
	var ws *websocket.Conn
	var URL string
	var handshakeDialer = &websocket.Dialer{
		Subprotocols: []string{"p1", "p2"},
	}
	listenWS := make(chan interface{})
	wsHeader := http.Header{"Origin": {"http://localhost/"}}

	for {
		select {
		case ruleI := <-b.inrule:
			var err error
			// set a parameter of the block
			r, ok := ruleI.(map[string]interface{})
			if !ok {
				b.Error("bad rule")
				break
			}

			url, ok := r["url"]
			if !ok {
				b.Error("no url specified")
				break
			}
			surl, ok := url.(string)
			if !ok {
				b.Error("error reading url")
				break
			}
			if ws != nil {
				ws.Close()
			}

			ws, _, err = handshakeDialer.Dial(surl, wsHeader)
			if err != nil {
				b.Error("could not connect to url")
				break
			}
			ws.SetReadDeadline(time.Time{})
			go recv(ws, listenWS)

			URL = surl
		case <-b.quit:
			// quit the block
			return
		case o := <-b.queryrule:
			o <- map[string]interface{}{
				"url": URL,
			}
		case in := <-listenWS:
			b.out <- in
		}
	}
}

func reader(ws *websocket.Conn) {
	defer ws.Close()
	ws.SetReadLimit(512)
	ws.SetReadDeadline(time.Now().Add(pongWait))
	ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, _, err := ws.ReadMessage()
		if err != nil {
			break
		}
	}
}

// New returns a new display client implementing the ioext.ReadWriteCloser
// interface using a websocket.
//
// The DisplayClient essentially takes care of sending and receiving ping/pongs
// to keep the websocket alive. However, the DisplayClient does read/write
// directly on websocket without any buffering, hence, you must keep calling
// Read() with a non-empty buffer to keep the connection alive.
func New(ws *websocket.Conn) *DisplayClient {
	c := &DisplayClient{
		ws: ws,
	}

	ws.SetReadLimit(displayconsts.DisplayMaxMessageSize)
	ws.SetReadDeadline(time.Now().Add(displayconsts.DisplayPongTimeout))
	ws.SetPongHandler(c.pongHandler)
	go c.sendPings()

	return c
}

func NewCcusConn(ws *websocket.Conn, cc_in chan updateInfo) *ccusConn {

	log.Printf("ws[%s]:new connect %s", ws.RemoteAddr(), ws.LocalAddr())
	ws.SetReadDeadline(time.Now().Add(pongWait))
	ws.SetWriteDeadline(time.Now().Add(writeWait))
	pc := &ccusConn{}
	pc.ws = ws
	pc.cc_in = cc_in
	go pc.receiver()

	return pc
}

func (ws *WebSocker) reader(conn *websocket.Conn) {
	defer conn.Close()
	conn.SetReadLimit(512)
	conn.SetReadDeadline(time.Now().Add(pongWait))
	conn.SetPongHandler(func(string) error { conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			break
		}
	}
}

// wsConnHandler starts goroutines for forever reading from, and writing to, a
// websocket connection. Reads can be read from the ch_r channel, writes be
// sent to the ch_w channel.
//
// When a read from the websocket returns an error (which for example happens
// when the connection is closed), the read goroutine will terminate, but not
// right away. The faulty read is available on the read channel for a while
// before a timeout kicks in and the channel is closed. This is a little weird,
// but it ensures that faulty reads can be read by some listening goroutine,
// while it is at the same time guaranteed that the goroutine terminates even
// if there is no listener.
//
// When the caller closes the write channel (don't forget to do that!),
// wsConnHandler will close the websocket connection (with Close(), not with a
// control message) and terminate both goroutines. If you want to close the
// websocket with a control message, just do it by sending a control message
// directly over the conn object (this is legal).  After that, close the write
// channel.
func wsConnHandler(c *websocket.Conn) (<-chan *wsReadResult,
	chan<- *wsWriteCmd) {

	// channels we expose
	ch_r := make(chan *wsReadResult)
	ch_w := make(chan *wsWriteCmd)

	// reader
	go func() {
		for cont := true; cont; {
			// read from websocket forever
			res := new(wsReadResult)
			c.SetReadDeadline(time.Now().Add(
				time.Duration(appVars.config.WebsocketConnectionTimeoutS) * time.Second))
			res.messageType, res.data, res.err = c.ReadMessage() // err on socket close

			if res.err == nil {
				// got a message? send to read channel and read from websocket again
				ch_r <- res
			} else {
				log.Printf("ws conn handler reader got error (normal at close)")
				// got an error from the read? offer on read channel until timeout.
				// Eventually, break out of loop
				select {
				case ch_r <- res:
					cont = false
				case <-time.After(30 * time.Second):
					cont = false
				}
			}
		}
		close(ch_r)
		log.Printf("ws conn handler reader terminates")
		return
	}()

	// writer
	go func() {
		// recv from ch_w and send what is received over WriteMessage until channel
		// is closed
		for cmd := range ch_w {
			c.SetWriteDeadline(time.Now().Add(
				time.Duration(appVars.config.WebsocketConnectionTimeoutS) * time.Second))
			err := c.WriteMessage(cmd.messageType, cmd.data)
			cmd.ch_ret <- err
		}
		// when channel is closed, close the websocket
		log.Printf("ws conn handler writer closes websocket connection and terminates")
		c.Close() // reader goroutine will get an error from ReadMessage()
		return
	}()
	return ch_r, ch_w
}

func imwork(wsConn *websocket.Conn) {
	wsConn.SetReadLimit(maxMessageSize)
	wsConn.SetReadDeadline(time.Now().Add(pongWait))
	wsConn.SetPongHandler(func(string) error { wsConn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, msgBytes, err := wsConn.ReadMessage()
		if err != nil {
			fmt.Println(err)
			break
		}
		//		fmt.Println("receive:", msgBytes)
		go procMsg(msgBytes)
	}
}

func countAllIncomingMessages(ws *websocket.Conn, timeout time.Duration) int {
	var count int
	for {
		if timeout != time.Duration(0) {
			ws.SetReadDeadline(time.Now().Add(timeout))
		}
		_, err := getTextMessage(ws)
		if err != nil {
			break
		}
		count += 1
	}
	return count
}

//readMessagesAndKeepTheLastOne takes in a Gorilla websocket connection and how
//many messages to read and reads in that number of messages, returning the last
//message read if no error was encountered or an error if one was encountered.
func readMessagesAndKeepTheLastOne(conn *ws.Conn, howMany int) ([]byte, error) {
	var msg []byte
	var err error
	for i := 0; i < howMany; i++ {
		//Make sure the connection doesn't block if there's no message to read.
		conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))

		_, msg, err = conn.ReadMessage()
		if err != nil {
			return nil, err
		}
	}
	return msg, nil
}

func pumpStdin(ws *websocket.Conn, w io.Writer) {
	defer ws.Close()
	ws.SetReadLimit(maxMessageSize)
	ws.SetReadDeadline(time.Now().Add(pongWait))
	ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, message, err := ws.ReadMessage()
		if err != nil {
			break
		}
		message = append(message, '\n')
		if _, err := w.Write(message); err != nil {
			break
		}
	}
}

// pumpStdin handles reading data from the websocket connection and writing
// it to stdin of the process.
func pumpStdin(conn *websocket.Conn, stdin io.WriteCloser) {
	// Setup our connection's websocket ping/pong handlers from our const values.
	defer conn.Close()
	conn.SetReadLimit(maxMessageSize)
	conn.SetReadDeadline(time.Now().Add(pongWait))
	conn.SetPongHandler(func(string) error { conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			break
		}
		message = append(message, '\n')
		if _, err := stdin.Write(message); err != nil {
			break
		}
	}
}