From cb68659cff217b62a7e8468dcf42205cafb40dca Mon Sep 17 00:00:00 2001
From: miya <miyaweb@163.com>
Date: Thu, 22 Apr 2021 11:03:56 +0800
Subject: [PATCH] ws demo (#853)

* examples/ws: ws demo
---
 examples/go.mod                |  1 +
 examples/go.sum                |  1 +
 examples/ws/client/client.go   | 77 ++++++++++++++++++++++++++++++++++
 examples/ws/handler/handler.go | 32 ++++++++++++++
 examples/ws/main.go            | 28 +++++++++++++
 5 files changed, 139 insertions(+)
 create mode 100644 examples/ws/client/client.go
 create mode 100644 examples/ws/handler/handler.go
 create mode 100644 examples/ws/main.go

diff --git a/examples/go.mod b/examples/go.mod
index 64fac4b5b..c05998e90 100644
--- a/examples/go.mod
+++ b/examples/go.mod
@@ -16,6 +16,7 @@ require (
 	github.com/golang/protobuf v1.4.3
 	github.com/google/wire v0.5.0
 	github.com/gorilla/mux v1.8.0
+	github.com/gorilla/websocket v1.4.2
 	github.com/hashicorp/consul/api v1.8.1
 	github.com/hashicorp/go-sockaddr v1.0.2 // indirect
 	github.com/imdario/mergo v0.3.12 // indirect
diff --git a/examples/go.sum b/examples/go.sum
index 4d4d4e214..17fdbe796 100644
--- a/examples/go.sum
+++ b/examples/go.sum
@@ -231,6 +231,7 @@ github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1 h1:EGx4pi6eqNxGa
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
 github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI=
 github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
+github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/grpc-ecosystem/go-grpc-middleware v1.0.0/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgfV/d3M/q6VIi02HzZEHgUlZvzk=
diff --git a/examples/ws/client/client.go b/examples/ws/client/client.go
new file mode 100644
index 000000000..92bf527a4
--- /dev/null
+++ b/examples/ws/client/client.go
@@ -0,0 +1,77 @@
+package main
+
+import (
+	"flag"
+	"log"
+	"net/url"
+	"os"
+	"os/signal"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+var addr = flag.String("addr", "localhost:8080", "http service address")
+
+func main() {
+	flag.Parse()
+	log.SetFlags(0)
+
+	interrupt := make(chan os.Signal, 1)
+	signal.Notify(interrupt, os.Interrupt)
+
+	//u := url.URL{Scheme: "ws", Host: *addr, Path: "/echo"}
+	u := url.URL{Scheme: "ws", Host: *addr, Path: "/ws"}
+	log.Printf("connecting to %s", u.String())
+
+	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
+	if err != nil {
+		log.Fatal("dial:", err)
+	}
+	defer c.Close()
+
+	done := make(chan struct{})
+
+	go func() {
+		defer close(done)
+		for {
+			_, message, err := c.ReadMessage()
+			if err != nil {
+				log.Println("read:", err)
+				return
+			}
+			log.Printf("recv: %s", message)
+		}
+	}()
+
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-done:
+			return
+		case t := <-ticker.C:
+			err := c.WriteMessage(websocket.TextMessage, []byte(t.String()))
+			if err != nil {
+				log.Println("write:", err)
+				return
+			}
+		case <-interrupt:
+			log.Println("interrupt")
+
+			// Cleanly close the connection by sending a close message and then
+			// waiting (with timeout) for the server to close the connection.
+			err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+			if err != nil {
+				log.Println("write close:", err)
+				return
+			}
+			select {
+			case <-done:
+			case <-time.After(time.Second):
+			}
+			return
+		}
+	}
+}
diff --git a/examples/ws/handler/handler.go b/examples/ws/handler/handler.go
new file mode 100644
index 000000000..139be8e99
--- /dev/null
+++ b/examples/ws/handler/handler.go
@@ -0,0 +1,32 @@
+package handler
+
+import (
+	"log"
+	"net/http"
+
+	"github.com/gorilla/websocket"
+)
+
+var upgrader = websocket.Upgrader{}
+
+func WsHandler(w http.ResponseWriter, r *http.Request) {
+	c, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Print("upgrade:", err)
+		return
+	}
+	defer c.Close()
+	for {
+		mt, message, err := c.ReadMessage()
+		if err != nil {
+			log.Println("read:", err)
+			break
+		}
+		log.Printf("recv: %s", message)
+		err = c.WriteMessage(mt, message)
+		if err != nil {
+			log.Println("write:", err)
+			break
+		}
+	}
+}
diff --git a/examples/ws/main.go b/examples/ws/main.go
new file mode 100644
index 000000000..e48da16bf
--- /dev/null
+++ b/examples/ws/main.go
@@ -0,0 +1,28 @@
+package main
+
+import (
+	"log"
+
+	"github.com/go-kratos/kratos/examples/ws/handler"
+	"github.com/go-kratos/kratos/v2"
+	transhttp "github.com/go-kratos/kratos/v2/transport/http"
+	"github.com/gorilla/mux"
+)
+
+func main() {
+	router := mux.NewRouter()
+	router.HandleFunc("/ws", handler.WsHandler)
+
+	httpSrv := transhttp.NewServer(transhttp.Address(":8080"))
+	httpSrv.HandlePrefix("/", router)
+
+	app := kratos.New(
+		kratos.Name("ws"),
+		kratos.Server(
+			httpSrv,
+		),
+	)
+	if err := app.Run(); err != nil {
+		log.Println(err)
+	}
+}