From 811c6a986fe23fa11cae34295dd050767b0a9916 Mon Sep 17 00:00:00 2001
From: Matthew Holt <Matthew.Holt+git@gmail.com>
Date: Tue, 3 Mar 2015 09:49:45 -0700
Subject: [PATCH] Added WebSocket middleware

---
 config/middleware.go                |  2 +
 middleware/websockets/websocket.go  | 32 +++++++++++
 middleware/websockets/websockets.go | 87 +++++++++++++++++++++++++++++
 3 files changed, 121 insertions(+)
 create mode 100644 middleware/websockets/websocket.go
 create mode 100644 middleware/websockets/websockets.go

diff --git a/config/middleware.go b/config/middleware.go
index d08b718fc..bc526b9bb 100644
--- a/config/middleware.go
+++ b/config/middleware.go
@@ -10,6 +10,7 @@ import (
 	"github.com/mholt/caddy/middleware/proxy"
 	"github.com/mholt/caddy/middleware/redirect"
 	"github.com/mholt/caddy/middleware/rewrite"
+	"github.com/mholt/caddy/middleware/websockets"
 )
 
 // This init function registers middleware. Register middleware
@@ -24,6 +25,7 @@ func init() {
 	register("ext", extensionless.New)
 	register("proxy", proxy.New)
 	register("fastcgi", fastcgi.New)
+	register("websocket", websockets.New)
 }
 
 var (
diff --git a/middleware/websockets/websocket.go b/middleware/websockets/websocket.go
new file mode 100644
index 000000000..201fa40ce
--- /dev/null
+++ b/middleware/websockets/websocket.go
@@ -0,0 +1,32 @@
+package websockets
+
+import (
+	"os/exec"
+
+	"golang.org/x/net/websocket"
+)
+
+// WebSocket represents a web socket server configuration.
+type WebSocket struct {
+	Path      string
+	Command   string
+	Arguments []string
+}
+
+// Handle handles a WebSocket connection. It launches the
+// specified command and streams input and output through
+// the command's stdin and stdout.
+func (ws WebSocket) Handle(conn *websocket.Conn) {
+	cmd := exec.Command(ws.Command, ws.Arguments...)
+	cmd.Stdin = conn
+	cmd.Stdout = conn
+
+	// TODO: Set environment variables according to CGI 1.1
+	// cf. http://tools.ietf.org/html/rfc3875#section-4.1.4
+	cmd.Env = append(cmd.Env, `GATEWAY_INTERFACE="caddy-CGI/1.1"`)
+
+	err := cmd.Run()
+	if err != nil {
+		panic(err)
+	}
+}
diff --git a/middleware/websockets/websockets.go b/middleware/websockets/websockets.go
new file mode 100644
index 000000000..eab35b389
--- /dev/null
+++ b/middleware/websockets/websockets.go
@@ -0,0 +1,87 @@
+// Package websockets implements a WebSocket server by executing
+// a command and piping its input and output through the WebSocket
+// connection.
+package websockets
+
+import (
+	"log"
+	"net/http"
+
+	"github.com/flynn/go-shlex"
+	"github.com/mholt/caddy/middleware"
+	"golang.org/x/net/websocket"
+)
+
+// WebSockets is a type which holds configuration
+// for the websocket middleware collectively.
+type WebSockets struct {
+	Sockets []WebSocket
+}
+
+// ServeHTTP more or less converts the HTTP request to a WebSocket connection.
+func (ws WebSockets) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	for _, socket := range ws.Sockets {
+		if middleware.Path(r.URL.Path).Matches(socket.Path) {
+			websocket.Handler(socket.Handle).ServeHTTP(w, r)
+			return
+		}
+	}
+}
+
+// New constructs and configures a new websockets middleware instance.
+func New(c middleware.Controller) (middleware.Middleware, error) {
+	var websocks []WebSocket
+
+	var path string
+	var command string
+
+	for c.Next() {
+		var val string
+
+		// Path or command; not sure which yet
+		if !c.NextArg() {
+			return nil, c.ArgErr()
+		}
+		val = c.Val()
+
+		// The rest of the arguments are the command
+		if c.NextArg() {
+			path = val
+			command = c.Val()
+			for c.NextArg() {
+				command += " " + c.Val()
+			}
+		} else {
+			path = "/"
+			command = val
+		}
+
+		// Split command into the actual command and its arguments
+		var cmd string
+		var args []string
+
+		parts, err := shlex.Split(command)
+		if err != nil {
+			log.Fatal("Error parsing command for websocket use: " + err.Error())
+		} else if len(parts) == 0 {
+			log.Fatal("No command found for use by websocket.")
+		}
+
+		cmd = parts[0]
+		if len(parts) > 1 {
+			args = parts[1:]
+		}
+
+		websocks = append(websocks, WebSocket{
+			Path:      path,
+			Command:   cmd,
+			Arguments: args,
+		})
+	}
+
+	return func(next http.HandlerFunc) http.HandlerFunc {
+		// We don't use next because websockets aren't HTTP,
+		// so we don't invoke other middleware after this.
+		return WebSockets{Sockets: websocks}.ServeHTTP
+	}, nil
+}