Forráskód Böngészése

Use new Upgrader API in examples

Joachim Bauch 11 éve
szülő
commit
9474333b60
3 módosított fájl, 30 hozzáadás és 27 törlés
  1. 8 3
      examples/autobahn/server.go
  2. 11 10
      examples/chat/conn.go
  3. 11 14
      examples/filewatch/main.go

+ 8 - 3
examples/autobahn/server.go

@@ -28,7 +28,10 @@ import (
 
 // 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, nil, 4096, 4096)
+	u := websocket.Upgrader{
+		ReadBufferSize:  4096,
+		WriteBufferSize: 4096}
+	conn, err := u.Upgrade(w, r, nil)
 	if err != nil {
 		log.Println("Upgrade:", err)
 		http.Error(w, "Bad request", 400)
@@ -87,10 +90,12 @@ func echoCopyFull(w http.ResponseWriter, r *http.Request) {
 // 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, nil, 4096, 4096)
+	u := websocket.Upgrader{
+		ReadBufferSize:  4096,
+		WriteBufferSize: 4096}
+	conn, err := u.Upgrade(w, r, nil)
 	if err != nil {
 		log.Println("Upgrade:", err)
-		http.Error(w, "Bad request", 400)
 		return
 	}
 	defer conn.Close()

+ 11 - 10
examples/chat/conn.go

@@ -89,16 +89,17 @@ func serveWs(w http.ResponseWriter, r *http.Request) {
 		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, 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)
+	u := websocket.Upgrader{
+		ReadBufferSize:  1024,
+		WriteBufferSize: 1024,
+		CheckOrigin: func(r *http.Request) bool {
+			return r.Header.Get("Origin") == "http://"+r.Host
+		}}
+	ws, err := u.Upgrade(w, r, nil)
+	if err != nil {
+		if _, ok := err.(websocket.HandshakeError); !ok {
+			log.Println(err)
+		}
 		return
 	}
 	c := &connection{send: make(chan []byte, 256), ws: ws}

+ 11 - 14
examples/filewatch/main.go

@@ -107,20 +107,17 @@ func writer(ws *websocket.Conn, lastMod time.Time) {
 }
 
 func serveWs(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, 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)
+	u := websocket.Upgrader{
+		ReadBufferSize:  1024,
+		WriteBufferSize: 1024,
+		CheckOrigin: func(r *http.Request) bool {
+			return r.Header.Get("Origin") == "http://"+r.Host
+		}}
+	ws, err := u.Upgrade(w, r, nil)
+	if err != nil {
+		if _, ok := err.(websocket.HandshakeError); !ok {
+			log.Println(err)
+		}
 		return
 	}