Golang Conn.SetReadLimit方法代碼示例

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

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

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

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

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 pumpStdin(ws *websocket.Conn, w io.Writer) {
	defer ws.Close()
	ws.SetReadLimit(maxMessageSize)
	for {
		_, message, err := ws.ReadMessage()
		if err != nil {
			break
		}
		message = append(message, '\n')
		if _, err := w.Write(message); err != nil {
			break
		}
	}
	log.Println("exit stdin pump")
}

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

func wsReader(ws *websocket.Conn, mc chan string) {
	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 {
		_, msg, err := ws.ReadMessage()
		if err != nil {
			log.Println("Error when reading websocket message")
			log.Println(err)
			break
		}
		log.Println("Recieved message: " + string(msg))
		mc <- string(msg)
	}
}

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

// readPump pumps messages from the websocket connection to the hub.
func responseAndAnswer(ws *websocket.Conn, client string, r *http.Request, w http.ResponseWriter) {
	defer func() {
		ws.Close()
	}()
	ws.SetReadLimit(maxMessageSize)
	ws.SetReadDeadline(time.Now().Add(pongWait))
	ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	//enemyFromCookie:=""
	for {
		_, message, err := ws.ReadMessage()
		if err != nil {
			break
		}

		//enemyNameCookie, _ := r.Cookie("enemyName")
		//if (enemyNameCookie!=nil){
		//	enemyFromCookie =enemyNameCookie.Value
		//}
		//fmt.Println(string(message))
		names := string(message[:bytes.Index(message, []byte(":"))])
		colons := bytes.Index(message, []byte(":"))
		identifier := string(message[colons+1 : colons+9])
		//fmt.Println("IDDDDDDDDDDD "+identifier)

		switch identifier {
		case "clientTo":
			askInvitation(names)
			break
		case "Accepted":
			acceptInvitation(names, colons, client, w)
			break
		case "CheckVal":
			getSideValue(client)
			break
		case "MadeStep":
			madeStep(client, names)
		case "QuitGame":
			quitGame(w, r, names)
		}

	}
}

// NewShellHandler returns a new ShellHandler structure for that can
// serve/expose a shell over a websocket.
func NewShellHandler(ws *websocket.Conn, log *logrus.Entry) *ShellHandler {
	tellIn := make(chan int, 10)
	stdin, stdinWriter := ioext.AsyncPipe(shellconsts.ShellMaxPendingBytes, tellIn)
	stdoutReader, stdout := ioext.BlockedPipe()
	stderrReader, stderr := ioext.BlockedPipe()
	stdoutReader.Unblock(shellconsts.ShellMaxPendingBytes)
	stderrReader.Unblock(shellconsts.ShellMaxPendingBytes)

	s := &ShellHandler{
		log:          log,
		ws:           ws,
		stdin:        stdin,
		stdout:       stdout,
		stderr:       stderr,
		stdinWriter:  stdinWriter,
		stdoutReader: stdoutReader,
		stderrReader: stderrReader,
		tellIn:       tellIn,
	}

	ws.SetReadLimit(shellconsts.ShellMaxMessageSize)

	return s
}

func (user *User) Receiver(location *Location, ws *websocket.Conn) {
	// receives and decodes messages from users
	go func() {
		ws.SetReadLimit(MAX_MESSAGE_SIZE)
		ws.SetReadDeadline(time.Now().Add(PONG_WAIT))
		ws.SetPongHandler(func(string) error {
			ws.SetReadDeadline(time.Now().Add(PONG_WAIT))
			return nil
		})
		for {
			op, r, err := ws.NextReader()
			if err != nil {
				user.Kick(err.Error())
				return
			}
			switch op {
			case websocket.BinaryMessage:
				data, err := ioutil.ReadAll(r)
				if err != nil {
					user.Kick(err.Error())
					return
				}
				var event []interface{}
				err = codec.NewDecoderBytes(data, &msgpackHandle).Decode(&event)
				if err != nil {
					user.Kick(err.Error())
					return
				}
				location.Message <- &UserAndEvent{user, &event}
			default:
				user.Kick(fmt.Sprintf("bad message type: %v", op))
				return
			}
		}
	}()
}