Issue #9: Add websocket support to demo clients

* Add websocket support to demo/server
* Add simple websocket client in demo/wsclient

demo/server/server.go

@@ -1,5 +1,5 @@
-// Package server provides a sample HTTP web server which registers
-// itself in consul using one or more URL prefixes to demonstrate and
+// Package server provides a sample HTTP/Websocket server which registers
+// itself in consul using one or more url prefixes to demonstrate and
 // test the automatic fabio routing table update.
 //
 // During startup the server performs the following steps:
@@ -12,17 +12,24 @@
 //   tag per prefix
 // * Install a signal handler to deregister the service on exit
 //
+// If the protocol is set to "ws" the registered endpoints function
+// as websocket echo servers.
+//
 // Example:
 //
+//   # http server
 //   ./server -addr 127.0.0.1:5000 -name svc-a -prefix /foo,/bar
-//   ./server -addr 127.0.0.1:6000 -name svc-b -prefix /baz,/bar
+//   ./server -addr 127.0.0.1:5001 -name svc-b -prefix /baz,/bar
+//
+//   # websocket server
+//   ./server -addr 127.0.0.1:6000 -name ws-a -prefix /echo1,/echo2 -proto ws
 //
-// This used to be hosted under https://github.com/magiconair/fabio-example
 package main
 
 import (
 	"flag"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
@@ -33,14 +40,16 @@ import (
 	"strings"
 
 	"github.com/eBay/fabio/_third_party/github.com/hashicorp/consul/api"
+	"github.com/eBay/fabio/_third_party/golang.org/x/net/websocket"
 )
 
 func main() {
-	var addr, consul, name, prefix string
+	var addr, consul, name, prefix, proto string
 	flag.StringVar(&addr, "addr", "127.0.0.1:5000", "host:port of the service")
 	flag.StringVar(&consul, "consul", "127.0.0.1:8500", "host:port of the consul agent")
 	flag.StringVar(&name, "name", filepath.Base(os.Args[0]), "name of the service")
 	flag.StringVar(&prefix, "prefix", "", "comma-sep list of host/path prefixes to register")
+	flag.StringVar(&proto, "proto", "http", "protocol for endpoints: http or ws")
 	flag.Parse()
 
 	if prefix == "" {
@@ -51,9 +60,16 @@ func main() {
 	// register prefixes
 	prefixes := strings.Split(prefix, ",")
 	for _, p := range prefixes {
-		http.HandleFunc(p, func(w http.ResponseWriter, r *http.Request) {
-			fmt.Fprintf(w, "Serving %s from %s on %s\n", r.RequestURI, name, addr)
-		})
+		switch proto {
+		case "http":
+			http.HandleFunc(p, func(w http.ResponseWriter, r *http.Request) {
+				fmt.Fprintf(w, "Serving %s from %s on %s\n", r.RequestURI, name, addr)
+			})
+		case "ws":
+			http.Handle(p, websocket.Handler(EchoServer))
+		default:
+			log.Fatal("Invalid protocol ", proto)
+		}
 	}
 
 	// start http server
@@ -123,3 +139,27 @@ func main() {
 	}
 	log.Printf("Deregistered service %q in consul", name)
 }
+
+func EchoServer(ws *websocket.Conn) {
+	addr := ws.LocalAddr().String()
+	pfx := []byte("[" + addr + "] ")
+
+	log.Printf("ws connect on %s", addr)
+
+	// the following could be done with io.Copy(ws, ws)
+	// but I want to add some meta data
+	var msg = make([]byte, 1024)
+	for {
+		n, err := ws.Read(msg)
+		if err != nil && err != io.EOF {
+			log.Printf("ws error on %s. %s", addr, err)
+			break
+		}
+		_, err = ws.Write(append(pfx, msg[:n]...))
+		if err != nil && err != io.EOF {
+			log.Printf("ws error on %s%s. %s", addr, err)
+			break
+		}
+	}
+	log.Printf("ws disconnect on %s", addr)
+}

demo/wsclient/.gitignore

@@ -0,0 +1 @@
+wsclient

demo/wsclient/wsclient.go

@@ -0,0 +1,50 @@
+// Package wsclient implements a simple web socket client
+// which reads lines from stdin and sends them to the
+// websocket url.
+package main
+
+import (
+	"bufio"
+	"flag"
+	"fmt"
+	"log"
+	"os"
+
+	"github.com/eBay/fabio/_third_party/golang.org/x/net/websocket"
+)
+
+func main() {
+	var url, origin string
+	flag.StringVar(&url, "url", "ws://127.0.0.1:9999/echo", "websocket URL")
+	flag.StringVar(&origin, "origin", "http://localhost/", "origin header")
+	flag.Parse()
+
+	if url == "" {
+		flag.Usage()
+		os.Exit(1)
+	}
+
+	ws, err := websocket.Dial(url, "", origin)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	go func() {
+		var msg = make([]byte, 512)
+		for {
+			n, err := ws.Read(msg)
+			if err != nil {
+				log.Fatal(err)
+			}
+			fmt.Printf("R: %s\nS: ", msg[:n])
+		}
+	}()
+
+	fmt.Print("S: ")
+	sc := bufio.NewScanner(os.Stdin)
+	for sc.Scan() {
+		if _, err := ws.Write(sc.Bytes()); err != nil {
+			log.Fatal(err)
+		}
+	}
+}