Golang websocket.Upgrader类代码示例

func slashws(w http.ResponseWriter, req *http.Request) {
	// Upgrader.Upgrade() has Origin check if .CheckOrigin is nil
	upgrader := gorillawebsocket.Upgrader{}
	wsconn, err := upgrader.Upgrade(w, req, nil)
	if err != nil { // Upgrade() does http.Error() to the client
		return
	}

	// req.Method == "GET" asserted by the mux
	req.Form = nil // reset reused later .Form
	c := &conn{
		Conn:          wsconn,
		requestOrigin: req,
		receive:       make(chan *received, 2),
		push:          make(chan *pageUpdate, 2),
		full:          defaultClient(),
	}
	register <- c
	defer func() {
		unregister <- c
		c.Conn.Close()
	}()
	stop := make(chan struct{}, 1)
	go c.receiveLoop(stop) // read from the client
	c.updateLoop(stop)     // write to the client
}

func webSocket(res http.ResponseWriter, req *http.Request) {
	var upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin:     func(r *http.Request) bool { return true },
	}

	res.Header().Set("Access-Control-Allow-Origin", "*")

	conn, err := upgrader.Upgrade(res, req, nil)
	if err != nil {
		log.Printf("Not a web socket connection: %s \n", err)
		return
	}

	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			return
		}
		if err = conn.WriteMessage(messageType, p); err != nil {
			return
		}
	}
}

func IndexWS(access *Access, errlog *log.Logger, minperiod flags.Period, w http.ResponseWriter, req *http.Request) {
	// Upgrader.Upgrade() has Origin check if .CheckOrigin is nil
	upgrader := gorillawebsocket.Upgrader{
		HandshakeTimeout: 5 * time.Second,
	}
	wsconn, err := upgrader.Upgrade(w, req, nil)
	if err != nil { // Upgrade() does http.Error() to the client
		return
	}

	// req.Method == "GET" asserted by the mux
	req.Form = nil // reset reused later .Form
	c := &conn{
		Conn: wsconn,

		requestOrigin: req,

		receive: make(chan *received, 2),
		pushch:  make(chan *IndexUpdate, 2),
		para:    params.NewParams(minperiod),
		access:  access,
	}
	Register <- c
	defer func() {
		unregister <- c
		c.Conn.Close()
	}()
	stop := make(chan struct{}, 1)
	go c.receiveLoop(stop)     // read from the client
	c.updateLoop(errlog, stop) // write to the client
}

func Upgrade(res http.ResponseWriter, req *http.Request) {
	if req.Header["Origin"][0] != "http://localhost" {
		return
	}
	up := websocket.Upgrader{
		CheckOrigin: func(req *http.Request) bool {
			if req.Header["Origin"][0] != "http://localhost" {
				//	return false
				// gorilla stops the app if wrong origin...
			}
			return true
		},
	}
	conn, err := up.Upgrade(res, req, nil)
	if err != nil {
		log.Fatal(err)
	}
	go hub.Handle(conn)
	hub.H.Register <- &hub.Register{conn, time.Now().UnixNano()}
	r, err := json.Marshal(&hub.BroadcastMessage{
		0,
		"User connected",
	})
	if err != nil {
		log.Fatal(err)
	}
	hub.H.Broadcast <- &hub.Broadcast{
		conn,
		r,
		"default",
	}
}

// Start the websocket server, each peer connecting to this websocket will be added as a connection to the dispatcher
func Start(d *Dispatcher, port int) {
	var upgrader = websocket.Upgrader{
		ReadBufferSize:  2048,
		WriteBufferSize: 2048,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		log.Debug("Connection received")
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Fatal(err)
		}

		defer conn.Close()

		startConnection(conn, d)
	})

	go http.ListenAndServe(fmt.Sprintf(":%d", port), nil)
	log.Println(fmt.Sprintf("Websocket server started on port %d", port))
	select {}
}

func (w *webserver) live(resp http.ResponseWriter, req *http.Request) {
	//get our feeder registered
	wsf := &liveWSFeeder{
		ch: make(chan namedBwSample, chanBufferSize),
	}
	id, err := w.lf.RegisterLiveFeeder(wsf)
	if err != nil {
		resp.WriteHeader(http.StatusInternalServerError)
		return
	}
	defer w.lf.DeregisterLiveFeeder(id)

	//upgrade to a websocket
	var upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	conn, err := upgrader.Upgrade(resp, req, nil)
	if err != nil {
		return
	}
	defer conn.Close()

	//start feeding and relaying
	for s := range wsf.ch {
		if err := websocket.WriteJSON(conn, s); err != nil {
			break
		}
	}
}

func main() {
	var upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}

	http.Handle("/", handlers.LoggingHandler(os.Stderr, http.FileServer(http.Dir("."))))
	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Println(err)
			return
		}
		defer conn.Close()
		for {
			t, msg, err := conn.ReadMessage()
			if err != nil {
				log.Println(err)
				return
			}
			log.Println("Received:", string(msg))
			err = conn.WriteMessage(t, msg)
			if err != nil {
				log.Println(err)
				return
			}
		}
	})
	host := "localhost:3030"
	log.Printf("http://%v", host)
	log.Fatal(http.ListenAndServe(host, nil))
}

func consoleHandler(w http.ResponseWriter, r *http.Request) {
	upgrader := websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	go wsReadLoop(conn)
	defer conn.Close()

	lines, done, err := tailFile(fOptions.LogPath)
	if err != nil {
		log.Println(err)
		return
	}
	defer close(done)

	for line := range lines {
		err = conn.WriteMessage(websocket.TextMessage, append(line, '\n'))
		if err != nil {
			return
		}
	}
}

// ListenAndServe will run a simple WS (HTTP) server.
// return listening error
func (s *Server) ListenAndServe(address string, path string) error {
	l, err := net.Listen("tcp", address)
	if err != nil {
		return err
	}
	s.listener = l

	mux := http.NewServeMux()
	mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		var upgrader = websocket.Upgrader{
			ReadBufferSize:  s.readBufferSize,
			WriteBufferSize: s.writeBufferSize,
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		}

		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Print("upgrade:", err)
			return
		}

		s.wsHandler.ServeWS(
			conn.UnderlyingConn(),
			newAbstractStream(conn),
			r.URL.Query(),
		)
	})
	s.httpHandler = mux

	go s.serve()

	return nil
}

func (ws *webSocketRenderer) handler(w http.ResponseWriter, r *http.Request) {
	var upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}

	h := http.Header{}

	h["Sec-WebSocket-Protocol"] = []string{"nesrender"}
	conn, err := upgrader.Upgrade(w, r, h)

	if err != nil {
		panic("Unable to upgrade connection")
	}

	ws.conns = append(ws.conns, conn)

	for {
		t, _, err := conn.ReadMessage()
		if err != nil {
			break
		}

		if t == websocket.CloseMessage {
			break
		}
	}
}

func startMockAcsServer(t *testing.T, closeWS <-chan bool) (*httptest.Server, chan<- string, <-chan string, <-chan error, error) {
	serverChan := make(chan string)
	requestsChan := make(chan string)
	errChan := make(chan error)

	upgrader := websocket.Upgrader{ReadBufferSize: 1024, WriteBufferSize: 1024}
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ws, err := upgrader.Upgrade(w, r, nil)
		go func() {
			<-closeWS
			ws.Close()
		}()
		if err != nil {
			errChan <- err
		}
		go func() {
			_, msg, err := ws.ReadMessage()
			if err != nil {
				errChan <- err
			} else {
				requestsChan <- string(msg)
			}
		}()
		for str := range serverChan {
			err := ws.WriteMessage(websocket.TextMessage, []byte(str))
			if err != nil {
				errChan <- err
			}
		}
	})

	server := httptest.NewTLSServer(handler)
	return server, serverChan, requestsChan, errChan, nil
}

func Handle(w http.ResponseWriter, r *http.Request) {
	upgrader := websocket.Upgrader{
		CheckOrigin: func(req *http.Request) bool {
			if req.Host == "127.0.0.1:9090" {
				return true
			} else {
				return false
			}
		},
	}
	var err error
	WebSocketConn, err = upgrader.Upgrade(w, r, nil)
	if err != nil {
		utils.Errorf("Error upgrade http: %s", err.Error())
		return
	}
	for {
		mt, bytes, err := WebSocketConn.ReadMessage()
		if err != nil {
			utils.Errorf("Error read message: %v", err)
			WebSocketConn.Close()
			WebSocketConn = nil
			return
		}
		utils.Debugf("Client msg received: %d, %s", mt, string(bytes))
	}
}

func (f *FakeFirehose) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
	f.lock.Lock()
	defer f.lock.Unlock()

	f.lastAuthorization = r.Header.Get("Authorization")
	f.requested = true

	if f.lastAuthorization != f.validToken {
		log.Printf("Bad token passed to firehose: %s", f.lastAuthorization)
		rw.WriteHeader(403)
		r.Body.Close()
		return
	}

	upgrader := websocket.Upgrader{
		CheckOrigin: func(*http.Request) bool { return true },
	}

	ws, _ := upgrader.Upgrade(rw, r, nil)

	defer ws.Close()
	defer ws.WriteControl(websocket.CloseMessage, f.closeMessage, time.Time{})

	for _, envelope := range f.events {
		buffer, _ := proto.Marshal(&envelope)
		err := ws.WriteMessage(websocket.BinaryMessage, buffer)
		if err != nil {
			panic(err)
		}
	}
	f.wg.Wait()
}

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

	// Upgrade from http to websocket protocol
	upgrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		panic(err)
	}
	defer c.Close()
	for {
		mt, message, err := c.ReadMessage()

		if err != nil {
			panic(err)
		}
		fmt.Printf("Received message: %s\n", message)
		err = c.WriteMessage(mt, message)
		if err != nil {
			panic(err)
		}
	}
}

func main() {
	upgrader := websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		id, err := authentication.AuthenticateRequest(r.FormValue("AccessToken"))
		if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}

		if id == -1 {
			http.Error(w, "", 401)
			return
		}

		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Fatal(err)
		}

		c := client.New(id, conn)
		router.AddClientChan <- c
		go c.Run()
	})

	fmt.Println("listening on 8080")
	http.ListenAndServe(":8080", nil)
}