Golang websocket.Upgrade函數代碼示例

func ServeRequest(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		return
	}
	if r.Header.Get("Origin") != "http://"+r.Host {
		http.Error(w, "Origin not allowed", 403)
		return
	}
	ws, err := websocket.Upgrade(w, r.Header, 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
	}

	tid, err := uuid.NewV4()
	if err != nil {
		log.Error("could not create sessionid")
		return
	}
	id := tid.String()

	c := &connection{SessionId: id, send: make(chan []byte, 256), ws: ws}
	h.register <- c
	go c.writePump()
	c.readPump()
}

// 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.Header, 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 newconn(w http.ResponseWriter, r *http.Request) (*conn, error) {
	wconn, err := websocket.Upgrade(w, r.Header, nil, 1024, 1024)
	if err != nil {
		return nil, err
	}
	hash := fnv.New32()
	hash.Write([]byte(r.RemoteAddr))
	return &conn{Id(hash.Sum32()), make(chan Msg, 64), wconn, time.NewTicker(pingPeriod)}, nil
}

// echoReadAll echoes messages from the client by reading the entire message
// with ioutil.ReadAll.
func echoReadAll(w http.ResponseWriter, r *http.Request, writeMessage bool) {
	conn, err := websocket.Upgrade(w, r.Header, nil, 4096, 4096)
	if err != nil {
		log.Println("Upgrade:", err)
		http.Error(w, "Bad request", 400)
		return
	}
	defer conn.Close()
	for {
		op, r, err := conn.NextReader()
		if err != nil {
			if err != io.EOF {
				log.Println("NextReader:", err)
			}
			return
		}
		if op == websocket.OpPong {
			continue
		}
		if op == websocket.OpText {
			r = &validator{r: r}
		}
		b, err := ioutil.ReadAll(r)
		if err != nil {
			if err == errInvalidUTF8 {
				conn.WriteControl(websocket.OpClose,
					websocket.FormatCloseMessage(websocket.CloseInvalidFramePayloadData, ""),
					time.Time{})
			}
			log.Println("ReadAll:", err)
			return
		}
		if writeMessage {
			err = conn.WriteMessage(op, b)
			if err != nil {
				log.Println("WriteMessage:", err)
			}
		} else {
			w, err := conn.NextWriter(op)
			if err != nil {
				log.Println("NextWriter:", err)
				return
			}
			if _, err := w.Write(b); err != nil {
				log.Println("Writer:", err)
				return
			}
			if err := w.Close(); err != nil {
				log.Println("Close:", err)
				return
			}
		}
	}
}

func wsHandler(w http.ResponseWriter, req *http.Request) {
	ws, err := websocket.Upgrade(w, req.Header, 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
	}
	//When we try to handle this, see if the hub exists.
	id := mux.Vars(req)["id"]
	wshub.Launch(ws, id)
}

func (this *WSController) Get() {
	ws, err := websocket.Upgrade(this.Ctx.ResponseWriter, this.Ctx.Request.Header, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(this.Ctx.ResponseWriter, "Not a websocket handshake", 400)
		return
	} else if err != nil {
		return
	}
	c := &connection{send: make(chan []byte, 256), ws: ws, username: randomString(10)}
	h.register <- c
	go c.writePump()
	c.readPump()
}

// documentStream handles websocket requests from the client for a particular document
func documentStream(ctx *web.Context, documentId string) {
	w := ctx.ResponseWriter
	r := ctx.Request
	if r.Header.Get("Origin") != "http://"+r.Host {
		ctx.Abort(403, "Origin not allowed")
		// http.Error(w, "Origin not allowed", 403)
		return
	}
	ws, err := websocket.Upgrade(w, r.Header, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		ctx.Abort(400, "Not a websocket handshake")
		return
	} else if err != nil {
		log.Println(err)
		return
	}
	redis_conn, _ := hub.GetRedis()
	defer redis_conn.Close()
	s, err := redis.String(redis_conn.Do("GET", documentId))
	//make new document at that stringId
	if err != nil {
		doc := document.NewDoc(documentId)
		document_hubs[doc.Name] = hub.DocumentHub{
			Document:    doc,
			Broadcast:   make(chan hub.Message),
			Register:    make(chan *hub.DocumentConnection),
			Unregister:  make(chan *hub.DocumentConnection),
			Connections: make(map[*hub.DocumentConnection]bool),
		}
		h := document_hubs[doc.Name]
		go h.Run()
		json_bytes, _ := json.Marshal(doc)
		redis_conn.Do("SET", documentId, string(json_bytes))
		c := &hub.DocumentConnection{Send: make(chan hub.Message, 256), Ws: ws, H: document_hubs[doc.Name]}
		document_hubs[doc.Name].Register <- c
		mainHub.Broadcast <- json_bytes
		go c.WritePump()
		c.ReadPump()
	} else {
		var doc document.Document
		fmt.Println(s)
		err := json.Unmarshal([]byte(s), &doc)
		if err != nil {
			fmt.Println("Error:", err, s)
		}
		c := &hub.DocumentConnection{Send: make(chan hub.Message, 256), Ws: ws, H: document_hubs[doc.Name]}
		document_hubs[doc.Name].Register <- c
		go c.WritePump()
		c.ReadPump()
	}
}

// echoCopy echoes messages from the client using io.Copy.
func echoCopy(w http.ResponseWriter, r *http.Request, writerOnly bool) {
	conn, err := websocket.Upgrade(w, r.Header, nil, 4096, 4096)
	if err != nil {
		log.Println("Upgrade:", err)
		http.Error(w, "Bad request", 400)
		return
	}
	defer conn.Close()
	for {
		op, r, err := conn.NextReader()
		if err != nil {
			if err != io.EOF {
				log.Println("NextReader:", err)
			}
			return
		}
		if op == websocket.OpPong {
			continue
		}
		if op == websocket.OpText {
			r = &validator{r: r}
		}
		w, err := conn.NextWriter(op)
		if err != nil {
			log.Println("NextWriter:", err)
			return
		}
		if op == websocket.OpText {
			r = &validator{r: r}
		}
		if writerOnly {
			_, err = io.Copy(struct{ io.Writer }{w}, r)
		} else {
			_, err = io.Copy(w, r)
		}
		if err != nil {
			if err == errInvalidUTF8 {
				conn.WriteControl(websocket.OpClose,
					websocket.FormatCloseMessage(websocket.CloseInvalidFramePayloadData, ""),
					time.Time{})
			}
			log.Println("Copy:", err)
			return
		}
		err = w.Close()
		if err != nil {
			log.Println("Close:", err)
			return
		}
	}
}

func NewWebView(ctl *ui.Controller, state ui.State, port int) *WebView {
	v := &WebView{
		wsMessages: util.NewBroadcast(),
	}

	switch p := state.GetProtocol().(type) {
	case *proto.Http:
		NewWebHttpView(v, ctl, p)
	}

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/http/in", 302)
	})

	http.HandleFunc("/_ws", func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Upgrade(w, r.Header, nil, 1024, 1024)

		if err != nil {
			http.Error(w, "Failed websocket upgrade", 400)
			log.Warn("Failed websocket upgrade: %v", err)
			return
		}

		msgs := v.wsMessages.Reg()
		defer v.wsMessages.UnReg(msgs)
		for m := range msgs {
			err := conn.WriteMessage(websocket.OpText, m.([]byte))
			if err != nil {
				// connection is closed
				break
			}
		}
	})

	http.HandleFunc("/static/", func(w http.ResponseWriter, r *http.Request) {
		parts := strings.Split(r.URL.Path, "/")
		name := parts[len(parts)-1]
		fn, ok := static.AssetMap[name]
		if !ok {
			http.NotFound(w, r)
			return
		}
		w.Write(fn())
	})

	log.Info("Serving web interface on localhost:%d", port)
	go http.ListenAndServe(fmt.Sprintf(":%d", port), nil)
	return v
}

func handleWs(out http.ResponseWriter, request *http.Request) {
	if request.Method != "GET" {
		http.Error(out, "Method not allowed", 405)
		return
	}

	ws, err := websocket.Upgrade(out, request.Header, nil, 4096, 4096)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(out, "Not a websocket handshake", 400)
	} else if err != nil {
		log.Println(err)
	} else {
		handleChat(ws)
	}
}

func NewWebView(ctl mvc.Controller, addr string) *WebView {
	wv := &WebView{
		Logger:     log.NewPrefixLogger("view", "web"),
		wsMessages: util.NewBroadcast(),
		ctl:        ctl,
	}

	// for now, always redirect to the http view
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/http/in", 302)
	})

	// handle web socket connections
	http.HandleFunc("/_ws", func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Upgrade(w, r.Header, nil, 1024, 1024)

		if err != nil {
			http.Error(w, "Failed websocket upgrade", 400)
			wv.Warn("Failed websocket upgrade: %v", err)
			return
		}

		msgs := wv.wsMessages.Reg()
		defer wv.wsMessages.UnReg(msgs)
		for m := range msgs {
			err := conn.WriteMessage(websocket.OpText, m.([]byte))
			if err != nil {
				// connection is closed
				break
			}
		}
	})

	// serve static assets
	http.HandleFunc("/static/", func(w http.ResponseWriter, r *http.Request) {
		buf, err := assets.ReadAsset(path.Join("assets", "client", r.URL.Path[1:]))
		if err != nil {
			wv.Warn("Error serving static file: %s", err.Error())
			http.NotFound(w, r)
			return
		}
		w.Write(buf)
	})

	wv.Info("Serving web interface on %s", addr)
	wv.ctl.Go(func() { http.ListenAndServe(addr, nil) })
	return wv
}

func (t wsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	t.Log("HELLOOO")
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		t.Logf("bad method: %s", r.Method)
		return
	}
	if r.Header.Get("Origin") != "http://"+r.Host {
		http.Error(w, "Origin not allowed", 403)
		t.Logf("bad origin: %s", r.Header.Get("Origin"))
		return
	}
	ws, err := websocket.Upgrade(w, r.Header, "", 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		t.Logf("bad handshake: %v", err)
		http.Error(w, "Not a websocket handshake", 400)
		return
	} else if err != nil {
		t.Logf("upgrade error: %v", err)
		return
	}
	defer ws.Close()
	for {
		op, r, err := ws.NextReader()
		if err != nil {
			if err != io.EOF {
				t.Logf("NextReader: %v", err)
			}
			return
		}
		if op == websocket.OpPong {
			continue
		}
		w, err := ws.NextWriter(op)
		if err != nil {
			t.Logf("NextWriter: %v", err)
			return
		}
		if _, err = io.Copy(w, r); err != nil {
			t.Logf("Copy: %v", err)
			return
		}
		if err := w.Close(); err != nil {
			t.Logf("Close: %v", err)
			return
		}
	}
}

// serverWs handles webocket requests from the client.
func serveWs(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		return
	}
	ws, err := websocket.Upgrade(w, r.Header, "", 1024, 1024)
	if err != nil {
		log.Println(err)
		http.Error(w, "Bad request", 400)
		return
	}
	c := &connection{send: make(chan []byte, 256), ws: ws}
	h.register <- c
	go c.writePump()
	c.readPump()
}

// creats the webocket http upgrade handler requests from the client.
func (h *HttpHandler) initServeGbWsHndlr(path string, world *World) {
	http.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "GET" {
			ErrHttpMethodNotAllowed.Report(w)
			return
		}
		ws, err := gbws.Upgrade(w, r.Header, "", 1024, 1024)
		if err != nil {
			log.Println("Unable to upgrade connection,", err)
			ErrHttpBadRequeset.Report(w)
			return
		}

		h.kickOffPlayer(NewGbWsConn(h.nextConnId, ws), world)
		h.nextConnId++
	})
}

func serveWs(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		return
	}
	ws, err := websocket.Upgrade(w, r.Header, 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}
	wsHub.register <- c
	go c.writePump()
	c.readPump()
}