Golang Conn.NextWriter方法代碼示例

func sendTermSize(control *websocket.Conn) error {
	width, height, err := terminal.GetSize(int(syscall.Stdout))
	if err != nil {
		return err
	}

	shared.Debugf("Window size is now: %dx%d", width, height)

	w, err := control.NextWriter(websocket.TextMessage)
	if err != nil {
		return err
	}

	msg := shared.ContainerExecControl{}
	msg.Command = "window-resize"
	msg.Args = make(map[string]string)
	msg.Args["width"] = strconv.Itoa(width)
	msg.Args["height"] = strconv.Itoa(height)

	buf, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	_, err = w.Write(buf)

	w.Close()
	return err
}

func datawsHandler(w http.ResponseWriter, r *http.Request) {
	var websocket *ws.Conn
	var upgradeerr error
	websocket, upgradeerr = upgrader.Upgrade(w, r, nil)
	if upgradeerr != nil {
		w.Write([]byte(fmt.Sprintf("Could not upgrade HTTP connection to WebSocket: %v\n", upgradeerr)))
		return
	}

	cw := ConnWrapper{
		Writing: &sync.Mutex{},
		Conn:    websocket,
	}

	websocket.SetReadLimit(MAX_REQSIZE)

	for {
		_, payload, err := websocket.ReadMessage()

		if err != nil {
			return // Most likely the connection was closed or the message was too big
		}

		uuidBytes, startTime, endTime, pw, token, echoTag, success := parseDataRequest(string(payload), &cw)

		if success {
			var loginsession *LoginSession
			if token != "" {
				loginsession = validateToken(token)
				if loginsession == nil {
					w.Write([]byte(ERROR_INVALID_TOKEN))
					return
				}
			}
			if hasPermission(loginsession, uuidBytes) {
				dr.MakeDataRequest(uuidBytes, startTime, endTime, uint8(pw), &cw)
			} else {
				cw.GetWriter().Write([]byte("[]"))
			}
		}
		if cw.CurrWriter != nil {
			cw.CurrWriter.Close()
		}

		writer, err := websocket.NextWriter(ws.TextMessage)
		if err != nil {
			fmt.Println("Could not echo tag to client")
		}

		if cw.CurrWriter != nil {
			_, err = writer.Write([]byte(echoTag))
			if err != nil {
				fmt.Println("Could not echo tag to client")
			}
			writer.Close()
		}

		cw.Writing.Unlock()
	}
}

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

	if r == nil {
		close(ch)
		return ch
	}

	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
			}
		}
		conn.WriteMessage(websocket.TextMessage, []byte{})
		ch <- true
	}(conn, r)

	return ch
}

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
}

// 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 writeMessage(c *websocket.Conn, str string, t *testing.T) {
	w, err := c.NextWriter(websocket.TextMessage)

	if err != nil {
		t.Error("cannot create websocket write")
	}

	io.WriteString(w, str)
	w.Close()
}

func primeRequestMainLoop(ws *websocket.Conn) {
	wantDiagOutput := true
	msgType, msg, err := ws.ReadMessage()
MainLoop:
	for ; err == nil && msgType == websocket.TextMessage; msgType, msg, err = ws.ReadMessage() {
		req := strings.ToLower(strings.Trim(string(msg), " \r\n\t\v\f"))
		switch req {
		case "nodiag":
			wantDiagOutput = false
			ws.WriteMessage(websocket.TextMessage, []byte("Diag OFF\n"))
		case "wantdiag":
			wantDiagOutput = true
			ws.WriteMessage(websocket.TextMessage, []byte("Diag ON\n"))
		case "exit":
			fallthrough
		case "quit":
			fallthrough
		case "close":
			break MainLoop
		default:
			if w, e := ws.NextWriter(websocket.TextMessage); e == nil {
				if req == "?" {
					w.Write([]byte(beginHelpMsg))
				}
				_, out, diag := primelib.DoCmd(req)
				if out != nil {
					for output := range out {
						outputStr := fmt.Sprintf("%v", output)
						_, err = w.Write([]byte(outputStr))
						_, err = w.Write([]byte("\n"))
					}
				}
				if wantDiagOutput {
					for diagOutput := range diag {
						diagOutputStr := fmt.Sprintf("%v", diagOutput)
						_, err = w.Write([]byte(diagOutputStr))
						_, err = w.Write([]byte("\n"))
					}
				}
				if req == "?" {
					w.Write([]byte(endHelpMsg))
				}
				w.Close()
			} else {
				log.Printf("ERROR: %v\n", e.Error())
			}
		}
	}
}

// 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 wsWriter(conn *websocket.Conn, m <-chan []byte, e chan<- error) {
	for b := range m {
		writer, err := conn.NextWriter(websocket.TextMessage)
		if err != nil {
			e <- io.EOF
			return
		}

		_, err = writer.Write(b)
		if err != nil {
			e <- err
			return
		}

		err = writer.Close()
		if err != nil {
			e <- err
			return
		}
	}
}

func (c *execCmd) forwardSignal(control *websocket.Conn, sig syscall.Signal) error {
	shared.LogDebugf("Forwarding signal: %s", sig)

	w, err := control.NextWriter(websocket.TextMessage)
	if err != nil {
		return err
	}

	msg := shared.ContainerExecControl{}
	msg.Command = "signal"
	msg.Signal = sig

	buf, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	_, err = w.Write(buf)

	w.Close()
	return err
}

func getThumbnail(requestedPath string, conn *websocket.Conn) {
	fr, err := os.Open(requestedPath)
	if err != nil {
		return
	}
	img, _, err := image.Decode(fr)
	if err != nil {
		return
	}
	thumb := resize.Resize(0, 150, img, resize.Bilinear)
	// w.Header().Add("Content-type", "image/jpeg")
	w, err := conn.NextWriter(websocket.TextMessage)
	if err != nil {
		return
	}
	buf := bytes.Buffer{}
	jpeg.Encode(&buf, thumb, nil)
	thumbBase64 := base64.StdEncoding.EncodeToString(buf.Bytes())
	if _, err := io.WriteString(w, thumbBase64); err != nil {
		return
	}
	w.Close()
	log.Printf("Thumbnail request handled for %s", requestedPath)
}

// WebsocketMirror allows mirroring a reader to a websocket and taking the
// result and writing it to a writer. This function allows for multiple
// mirrorings and correctly negotiates stream endings. However, it means any
// websocket.Conns passed to it are live when it returns, and must be closed
// explicitly.
func WebsocketMirror(conn *websocket.Conn, w io.WriteCloser, r io.ReadCloser) (chan bool, chan bool) {
	readDone := make(chan bool, 1)
	writeDone := make(chan bool, 1)
	go func(conn *websocket.Conn, w io.WriteCloser) {
		for {
			mt, r, err := conn.NextReader()
			if err != nil {
				Debugf("Got error getting next reader %s, %s", err, w)
				break
			}

			if mt == websocket.CloseMessage {
				Debugf("Got close message for reader")
				break
			}

			if mt == websocket.TextMessage {
				Debugf("Got message barrier, resetting stream")
				break
			}

			buf, err := ioutil.ReadAll(r)
			if err != nil {
				Debugf("Got error writing to writer %s", err)
				break
			}
			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
			}
		}
		writeDone <- true
		w.Close()
	}(conn, w)

	go func(conn *websocket.Conn, r io.ReadCloser) {
		in := ReaderToChannel(r)
		for {
			buf, ok := <-in
			if !ok {
				readDone <- true
				r.Close()
				Debugf("sending write barrier")
				conn.WriteMessage(websocket.TextMessage, []byte{})
				return
			}
			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)
		readDone <- true
		r.Close()
	}(conn, r)

	return readDone, writeDone
}

// WebsocketMirror allows mirroring a reader to a websocket and taking the
// result and writing it to a writer.
func WebsocketMirror(conn *websocket.Conn, w io.WriteCloser, r io.Reader) chan bool {
	done := make(chan bool, 2)
	go func(conn *websocket.Conn, w io.WriteCloser) {
		for {
			mt, r, err := conn.NextReader()
			if mt == websocket.CloseMessage {
				Debugf("got close message for reader")
				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
			}
			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
			}
		}
		done <- true
		w.Close()
	}(conn, w)

	go func(conn *websocket.Conn, r io.Reader) {
		in := ReaderToChannel(r)
		for {
			buf, ok := <-in
			if !ok {
				done <- true
				conn.WriteMessage(websocket.CloseMessage, []byte{})
				conn.Close()
				return
			}
			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)
		done <- true
	}(conn, r)

	return done
}

func bracketwsHandler(w http.ResponseWriter, r *http.Request) {
	var websocket *ws.Conn
	var upgradeerr error
	websocket, upgradeerr = upgrader.Upgrade(w, r, nil)
	if upgradeerr != nil {
		w.Write([]byte(fmt.Sprintf("Could not upgrade HTTP connection to WebSocket: %v\n", upgradeerr)))
		return
	}

	cw := ConnWrapper{
		Writing: &sync.Mutex{},
		Conn:    websocket,
	}

	websocket.SetReadLimit(MAX_REQSIZE)

	for {
		_, payload, err := websocket.ReadMessage()

		if err != nil {
			return // Most likely the connection was closed or the message was too big
		}

		uuids, token, echoTag, success := parseBracketRequest(string(payload), &cw, true)

		if success {
			var loginsession *LoginSession
			if token != "" {
				loginsession = validateToken(token)
				if loginsession == nil {
					w.Write([]byte(ERROR_INVALID_TOKEN))
					return
				}
			}
			var viewable []uuid.UUID = uuids[:0]
			for _, uuid := range uuids {
				if hasPermission(loginsession, uuid) {
					viewable = append(viewable, uuid)
				}
			}
			br.MakeBracketRequest(uuids, &cw)
		}
		if cw.CurrWriter != nil {
			cw.CurrWriter.Close()
		}

		writer, err := websocket.NextWriter(ws.TextMessage)
		if err != nil {
			fmt.Println("Could not echo tag to client")
		}

		if cw.CurrWriter != nil {
			_, err = writer.Write([]byte(echoTag))
			if err != nil {
				fmt.Println("Could not echo tag to client")
			}
			writer.Close()
		}

		cw.Writing.Unlock()
	}
}

			requestPayload = `{"path":"ls", "user": "snoopy"}`
		})

		JustBeforeEach(func() {
			wsURL, err := url.Parse(server.URL)
			Expect(err).NotTo(HaveOccurred())

			wsURL.Scheme = "ws"
			wsURL.Path = "/api/v1/containers/" + handle + "/hijack"

			dialer := websocket.Dialer{}
			conn, response, err = dialer.Dial(wsURL.String(), nil)
			if !expectBadHandshake {
				Expect(err).NotTo(HaveOccurred())

				writer, err := conn.NextWriter(websocket.TextMessage)
				Expect(err).NotTo(HaveOccurred())

				_, err = writer.Write([]byte(requestPayload))
				Expect(err).NotTo(HaveOccurred())

				err = writer.Close()
				Expect(err).NotTo(HaveOccurred())
			}
		})

		AfterEach(func() {
			if !expectBadHandshake {
				conn.Close()
			}
		})