Golang websocket.Upgrade函数代码示例

func (fl *fakeLoggregator) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	fl.receivedRequests = append(fl.receivedRequests, r)
	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		Fail("Bad Handshake")
		return
	} else if err != nil {
		Fail("Upgrade Close")
		return
	}
	defer ws.Close()

	r.ParseForm()

	for _, msg := range fl.messages {
		ws.WriteMessage(websocket.BinaryMessage, []byte(msg))
	}
	for {
		_, _, err := ws.ReadMessage()

		if err == nil {
			fl.numKeepAlives++
			continue
		}
		return
	}
}

func socket(w http.ResponseWriter, r *http.Request) {
	log.Print("connection established")
	if _case == "hang" {
		hang := time.Minute
		log.Printf("hanging for %s\n", hang.String())
		time.Sleep(hang)
	}
	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if err != nil {
		log.Print(err)
		return
	}
	defer func() {
		ws.Close()
		log.Print("connection closed")
	}()
	for {
		msgt, msg, err := ws.ReadMessage()
		if err != nil {
			log.Print(err)
			return
		}
		log.Print("rcvd: '" + string(msg) + "'")
		ws.WriteMessage(msgt, msg)
	}
}

func wsHandler(w http.ResponseWriter, r *http.Request) {

	// Websocket handshake.
	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(w, "Not a websocket handshake", 400)
		return
	} else if err != nil {
		return
	}

	addr := ws.RemoteAddr()

	log.Printf("Websocket accepted: %s\n", addr)

	// Creating player.
	one := newPlayer(ws)

	// Adding player to sector.
	mainSector.addPlayer(one)

	// Spawing writing goroutine.
	go one.writer()

	// Sending player identification.
	one.ident()

	// Blocking this function on a reader.
	one.reader()

	// Reader has stopped, good bye!
	log.Printf("Websocket finalized: %s", addr)
}

func serveWs(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		return
	}

	// Read request details.
	r.ParseForm()
	iface := r.FormValue("if")
	if iface == "" {
		iface = "eth0"
	}

	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(w, "Not a websocket handshake", 400)
		return
	} else if err != nil {
		log.Println(err)
		return
	}
	c := &connection{send: make(chan []byte, 256), ws: ws, iface: iface}
	h.register <- c
	go c.writePump()
	c.readPump()
}

func websocketHandler(w http.ResponseWriter, r *http.Request) {

	conn, _ := websocket.Upgrade(w, r, nil, 1024, 1024)
	addConnection(conn)
	listen(conn)

}

func wsHandler(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	fmt.Println(r.Method)
	// Taken from gorilla's website
	conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(w, "Not a websocket handshake", 400)
		return
	} else if err != nil {
		log.Println(err)
		return
	}
	log.Println("Succesfully upgraded connection")
	connections[conn] = true

	for {
		// Blocks until a message is read
		_, msg, err := conn.ReadMessage()
		if err != nil {
			delete(connections, conn)
			conn.Close()
			return
		}
		log.Println(string(msg))
		sendAll(msg)
	}
}

// Join method handles WebSocket requests for WebSocketController.
func (this *WebSocketController) Join() {
	uname := this.GetString("uname")
	if len(uname) == 0 {
		this.Redirect("/", 302)
		return
	}

	// Upgrade from http request to WebSocket.
	ws, err := websocket.Upgrade(this.Ctx.ResponseWriter, this.Ctx.Request, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(this.Ctx.ResponseWriter, "Not a websocket handshake", 400)
		return
	} else if err != nil {
		beego.Error("Cannot setup WebSocket connection:", err)
		return
	}

	// Join chat room.
	Join(uname, ws)
	defer Leave(uname)

	// Message receive loop.
	for {
		_, p, err := ws.ReadMessage()
		if err != nil {
			return
		}
		publish <- newEvent(models.EVENT_MESSAGE, uname, string(p))
	}
}

func serveWs(w http.ResponseWriter, r *http.Request) {
	log.Print("serveWs")
	session, _ := store.Get(r, "session-name")
	if session.IsNew {
		log.Printf("rejecting ws without session from %v",
			r.RemoteAddr)
		return
	}
	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		log.Printf("rejecting ws with bad handshake from %v",
			r.RemoteAddr)
		http.Error(w, "Not a websocket handshake", 400)
		return
	} else if err != nil {
		log.Printf("serveWs error: %s", err.Error())
		return
	}
	log.Printf("serveWs ws connection from %v", r.RemoteAddr)
	c := &wsconn{send: make(chan []byte, 256), ws: ws}
	h.register <- c
	defer func() { h.unregister <- c }()
	go c.writer()
	c.reader()
}

func notificationWSHandler(w http.ResponseWriter, r *http.Request) {
	sid := r.URL.Query()["sid"][0]

	cSession := coditorSessions.get(sid)

	if nil == cSession {
		return
	}

	conn, _ := websocket.Upgrade(w, r, nil, 1024, 1024)
	wsChan := util.WSChannel{Sid: sid, Conn: conn, Request: r, Time: time.Now()}

	ret := map[string]interface{}{"notification": "Notification initialized", "cmd": "init-notification"}
	err := wsChan.WriteJSON(&ret)
	if nil != err {
		return
	}

	notificationWS[sid] = &wsChan

	logger.Tracef("Open a new [Notification] with session [%s], %d", sid, len(notificationWS))

	// add user event handler
	cSession.EventQueue.addHandler(eventHandleFunc(event2Notification))

	input := map[string]interface{}{}

	for {
		if err := wsChan.ReadJSON(&input); err != nil {
			return
		}
	}
}

func (h *hub) handleConn(w http.ResponseWriter, r *http.Request) {
	conn, err := websocket.Upgrade(w, r, nil, 0, 0)
	ok(err)
	s := &stream{h: h, conn: conn}

	for {
		_, buf, err := conn.ReadMessage()
		if isReadFromClosedConnError(err) {
			log.Printf("conn closed: %v", err)
			break
		}
		ok(err)
		// TODO: Avoid decoding multiple times.
		var mt MsgType
		ok(json.Unmarshal(buf, &mt))
		switch mt.Type {
		case "SubscribeC2S":
			var msg SubscribeC2S
			ok(json.Unmarshal(buf, &msg))
			ok(s.processSubscribeC2S(&msg))
		case "SubscribeI2R":
			var msg SubscribeI2R
			ok(json.Unmarshal(buf, &msg))
			ok(s.processSubscribeI2R(&msg))
		case "PatchC2S":
			var msg PatchC2S
			ok(json.Unmarshal(buf, &msg))
			ok(s.processPatchC2S(&msg))
		default:
			panic(fmt.Errorf("unknown message type: %s", mt.Type))
		}
	}

	conn.Close()
}

func (b BackendController) WsHandler(writer http.ResponseWriter, request *http.Request) {
	conn, err := websocket.Upgrade(writer, request, nil, 1024, 1024)
	log.Println("getting a connection")
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(writer, "got a websocket handshake", 400)
		return
	} else if err != nil {
		log.Println(err)
		return
	}

	b.ConnectionHandler(conn)
	defer b.CleanUpHandler(conn)      // if this function ever exits, clean up the data
	defer delete(b.connections, conn) // if this function ever exits, clean up the data

	b.connections[conn] = true

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

func (w *WebsocketServer) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	w.logger.Debug("WebsocketServer.ServeHTTP: starting")
	var handler wsHandler
	var err error

	paths := strings.Split(request.URL.Path, "/")
	endpointName := paths[1]

	if endpointName == "firehose" {
		handler, err = w.firehoseHandler(paths, writer, request)
	} else {
		handler, err = w.appHandler(paths, writer, request)
	}

	if err != nil {
		w.logger.Errorf("WebsocketServer.ServeHTTP: %s", err.Error())
		return
	}

	ws, err := gorilla.Upgrade(writer, request, nil, 1024, 1024)
	if err != nil {
		w.logger.Errorf("WebsocketServer.ServeHTTP: Upgrade error (returning 400): %s", err.Error())
		http.Error(writer, err.Error(), 400)
		return
	}

	defer func() {
		ws.WriteControl(gorilla.CloseMessage, gorilla.FormatCloseMessage(gorilla.CloseNormalClosure, ""), time.Time{})
		ws.Close()
	}()

	handler(ws)
}

// Upgrade the connection to a websocket connection
func upgradeRequest(resp http.ResponseWriter, req *http.Request, o *Options) (*websocket.Conn, int, error) {
	if req.Method != "GET" {
		o.log("Method %s is not allowed", LogLevelWarning, req.RemoteAddr, req.Method)
		return nil, http.StatusMethodNotAllowed, errors.New("Method not allowed")
	}

	allowedOrigin := replacementRegexp.ReplaceAllString(o.AllowedOrigin, req.Host)
	if r, err := regexp.MatchString(allowedOrigin, req.Header.Get("Origin")); !r || err != nil {
		o.log("Origin %s is not allowed", LogLevelWarning, req.RemoteAddr, req.Host)
		return nil, http.StatusForbidden, errors.New("Origin not allowed")
	}

	o.log("Request to %s has been allowed for origin %s", LogLevelDebug, req.RemoteAddr, req.Host, req.Header.Get("Origin"))

	ws, err := websocket.Upgrade(resp, req, nil, 1024, 1024)
	if handshakeErr, ok := err.(websocket.HandshakeError); ok {
		o.log("Handshake failed: %s", LogLevelWarning, req.RemoteAddr, handshakeErr)
		return nil, http.StatusBadRequest, handshakeErr
	} else if err != nil {
		o.log("Handshake failed: %s", LogLevelWarning, req.RemoteAddr, err)
		return nil, http.StatusBadRequest, err
	}

	o.log("Connection established", LogLevelInfo, req.RemoteAddr)
	return ws, http.StatusOK, nil
}

func slashws(w http.ResponseWriter, req *http.Request) {
	if req.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	if req.Header.Get("Origin") != "http://"+req.Host {
		http.Error(w, "Forbidden", http.StatusForbidden)
		return
	}
	ws, err := websocket.Upgrade(w, req, nil, 1024, 1024)
	if err != nil {
		if _, ok := err.(websocket.HandshakeError); ok {
			http.Error(w, "websocket.Upgrade errd", http.StatusBadRequest)
			return
		}
		panic(err)
	}

	wc := &wclient{ws: ws, ping: make(chan bool, 1)}
	register <- wc
	defer func() {
		unregister <- wc
	}()
	go wc.waitfor_messages() // read from client
	wc.waitfor_updates()     // write to  client
}

func (h *sockethub) WsHandler(w http.ResponseWriter, r *http.Request) {
	var authenticated bool

	if h.Auth != nil {
		authenticated = h.Auth(r)
	}

	if authenticated {
		ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)

		if _, ok := err.(websocket.HandshakeError); ok {
			http.Error(w, "Not a websocket handshake", 400)
			return
		} else if err != nil {
			log.Printf("WsHandler error: %s", err.Error())
			return
		}

		c := &connection{send: make(chan *Data), ws: ws}
		h.register <- c

		defer func() { h.unregister <- c }()
		go c.writer()
		c.reader()
	} else {
		http.Error(w, "Invalid API key", 401)
	}
}