Golang Conn.NextReader方法代码示例

// 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()
}

/**
Listen for Pods. We're only interested in MODIFIED and DELETED events.
*/
func listen(wsConn *websocket.Conn, wsErrors chan string) {
	log.Println("Listening for pods")

	for {
		_, r, err := wsConn.NextReader()

		if err != nil {
			log.Printf("Error getting reader: %v", err)
			wsErrors <- "Error"
			return
		}

		dec := json.NewDecoder(r)
		var objmap map[string]*json.RawMessage
		dec.Decode(&objmap)

		var actionType string
		json.Unmarshal(*objmap["type"], &actionType)

		var pod api.Pod
		err = json.Unmarshal(*objmap["object"], &pod)

		switch actionType {
		case "MODIFIED":
			register(pod)
		case "DELETED":
			deletePod(pod)
		}
	}
}

// Listen for Services.
func svcListener(wsConn *websocket.Conn, wsErrors chan string) {
	log.Println("Listening for services")

	for {
		_, r, err := wsConn.NextReader()

		if err != nil {
			log.Printf("Error getting reader: %v", err)
			wsErrors <- "Error"
			return
		}

		dec := json.NewDecoder(r)

		var objmap map[string]*json.RawMessage
		dec.Decode(&objmap)

		// debug purpose
		// b,_ := json.Marshal(objmap)
		// log.Printf("%s", b)

		var actionType string
		json.Unmarshal(*objmap["type"], &actionType)

		var svc api.Service
		err = json.Unmarshal(*objmap["object"], &svc)

		switch actionType {
		case "ADDED":
			registerSvc(svc)
		case "DELETED":
			unregisterSvc(svc)
		}
	}
}

func readLoop(c *websocket.Conn) {
	for {
		if _, _, err := c.NextReader(); err != nil {
			c.Close()
			break
		}
	}
}

func (ws *Websocket) Reader(c *goWs.Conn, closed chan<- bool) {
	defer c.Close()
	for {
		messageType, _, err := c.NextReader()
		if err != nil || messageType == goWs.CloseMessage {
			break
		}
	}
	closed <- true
}

func keepWsAlive(ws *websocket.Conn) {
	// We repeatedly read from the websocket connections and discard
	// the reader in order to process the underlying ping/pong messages
	// of the websocket connection
	for {
		_, _, err := ws.NextReader()
		if err != nil {
			ws.Close()
			break
		}
	}
}

func ws_null_reader(conn *websocket.Conn, info_ptr *WsInfo) {

	// Apparently reading WebSocket messages from clients is mandatory.
	// This function also closes connections if needed.

	for {
		if _, _, err := conn.NextReader(); err != nil {
			remove_from_ws_client_list(info_ptr)
			conn.Close()
			return
		}
	}
}

func streamWebsocketIOFromContainerizer(ws *websocket.Conn, pidChannel chan<- string, processIO garden.ProcessIO) (int, error) {
	defer close(pidChannel)

	defer func() {
		ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))

		go func() {
			for {
				_, _, err := ws.NextReader()
				if err != nil {
					ws.Close()
					break
				}
			}
		}()

	}()

	receiveStream := ProcessStreamEvent{}
	for {
		err := websocket.ReadJSON(ws, &receiveStream)
		if err != nil {
			return -1, err
		}

		if receiveStream.MessageType == "pid" {
			pidChannel <- receiveStream.Data
		}

		if receiveStream.MessageType == "stdout" && processIO.Stdout != nil {
			io.WriteString(processIO.Stdout, receiveStream.Data)
		}
		if receiveStream.MessageType == "stderr" && processIO.Stderr != nil {
			io.WriteString(processIO.Stderr, receiveStream.Data)
		}

		if receiveStream.MessageType == "error" {
			return -1, errors.New(receiveStream.Data)
		}
		if receiveStream.MessageType == "close" {
			exitCode, err := strconv.Atoi(receiveStream.Data)
			if err != nil {
				return -1, err
			}
			return exitCode, nil
		}
	}
}

func (this *Server) NewConnection(conn *websocket.Conn) (string, error) {
	_, reader, err := conn.NextReader()
	if err != nil {
		glog.Errorf("Error establishing new connection: %s", err.Error())
	}
	user, err := message.GetUser(reader)
	if err != nil {
		return "", err
	}

	glog.Infof("Register conenction with uid: %s", user.Uid)
	this.Users[user.Uid] = NewClient(conn, user)
	this.Users[user.Uid].Send(user, message.MSG_USER_DATA)

	return user.Uid, nil
}

func WebsocketRecvStream(w io.WriteCloser, conn *websocket.Conn) chan bool {
	ch := make(chan bool)

	go func(w io.WriteCloser, conn *websocket.Conn) {
		for {
			mt, r, err := conn.NextReader()
			if mt == websocket.CloseMessage {
				Debugf("Got close message for reader")
				break
			}

			if mt == websocket.TextMessage {
				Debugf("got message barrier")
				break
			}

			if err != nil {
				Debugf("Got error getting next reader %s, %s", err, w)
				break
			}

			buf, err := ioutil.ReadAll(r)
			if err != nil {
				Debugf("Got error writing to writer %s", err)
				break
			}

			if w == nil {
				continue
			}

			i, err := w.Write(buf)
			if i != len(buf) {
				Debugf("Didn't write all of buf")
				break
			}
			if err != nil {
				Debugf("Error writing buf %s", err)
				break
			}
		}
		ch <- true
	}(w, conn)

	return ch
}

func wsReader(conn *websocket.Conn, m chan<- message, e chan<- error) {
	for {
		t, r, err := conn.NextReader()
		if err != nil {
			e <- io.EOF
			return
		}

		msg, err := ioutil.ReadAll(r)
		if err != nil {
			e <- err
			return
		}

		m <- message{ws.MsgType(t), msg}
	}
}

func (sc *ServerConnection) read(ws *websocket.Conn) {

	// log.Println("Response", resp)

	status := new(ServerUpdate)
	status.ConnID = sc.ConnID

	for {

		select {
		case <-sc.quit:
			log.Println("server reader got quit message")
			sc.quit <- true
			return

		default:

			ws.SetReadDeadline(time.Now().Add(time.Second * 3))
			sc.statusMsg("Ok")
			op, r, err := ws.NextReader()
			if err != nil {
				status := fmt.Sprintf("Error reading from server: %s", err)
				sc.statusErrorMsg(status)
				log.Println(status)
				return
			}
			msg, err := ioutil.ReadAll(r)

			// log.Println("op", op, "msg", string(msg), "err", err)

			if op == websocket.TextMessage {
				err = json.Unmarshal(msg, &status)
				if err != nil {
					log.Printf("Unmarshall err from '%s': '%s', data: '%s'\n", sc.IP.String(), err, msg)
				}
				// log.Printf("Got status: %#v\n", status)
				sc.updateChan <- status
			} else {
				log.Println("op", op, "msg", string(msg), "err", err)
			}

			// os.Exit(0)
		}
	}
}

func processRotondePackets(conn *websocket.Conn, inChan, outChan chan interface{}) {
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()

		for {
			select {
			case dispatcherPacket := <-inChan:
				jsonPacket, err := rotonde.ToJSON(dispatcherPacket)
				if err != nil {
					log.Warning(err)
				}
				if err := conn.WriteMessage(websocket.TextMessage, jsonPacket); err != nil {
					log.Fatal(err)
					return
				}
			}
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()

		for {
			messageType, reader, err := conn.NextReader()
			if err != nil {
				log.Fatal(err)
				return
			}
			if messageType == websocket.TextMessage {
				dispatcherPacket, err := rotonde.FromJSON(reader)
				if err != nil {
					log.Warning(err)
				}
				outChan <- dispatcherPacket
			}
		}
	}()

	log.Info("Treating messages")
	wg.Wait()
}

// 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 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)
	}
}