websocket refactored to use gorilla

config/setup/websocket.go

@@ -2,26 +2,26 @@ package setup
 
 import (
 	"github.com/mholt/caddy/middleware"
-	"github.com/mholt/caddy/middleware/websockets"
+	"github.com/mholt/caddy/middleware/websocket"
 )
 
-// WebSocket configures a new WebSockets middleware instance.
+// WebSocket configures a new WebSocket middleware instance.
 func WebSocket(c *Controller) (middleware.Middleware, error) {
 
 	websocks, err := webSocketParse(c)
 	if err != nil {
 		return nil, err
 	}
-	websockets.GatewayInterface = c.AppName + "-CGI/1.1"
-	websockets.ServerSoftware = c.AppName + "/" + c.AppVersion
+	websocket.GatewayInterface = c.AppName + "-CGI/1.1"
+	websocket.ServerSoftware = c.AppName + "/" + c.AppVersion
 
 	return func(next middleware.Handler) middleware.Handler {
-		return websockets.WebSockets{Next: next, Sockets: websocks}
+		return websocket.WebSocket{Next: next, Sockets: websocks}
 	}, nil
 }
 
-func webSocketParse(c *Controller) ([]websockets.Config, error) {
-	var websocks []websockets.Config
+func webSocketParse(c *Controller) ([]websocket.Config, error) {
+	var websocks []websocket.Config
 	var respawn bool
 
 	optionalBlock := func() (hadBlock bool, err error) {
@@ -74,7 +74,7 @@ func webSocketParse(c *Controller) ([]websockets.Config, error) {
 			return nil, err
 		}
 
-		websocks = append(websocks, websockets.Config{
+		websocks = append(websocks, websocket.Config{
 			Path:      path,
 			Command:   cmd,
 			Arguments: args,

config/setup/websocket_test.go

 package setup
 
 import (
-	"github.com/mholt/caddy/middleware/websockets"
 	"testing"
+
+	"github.com/mholt/caddy/middleware/websocket"
 )
 
 func TestWebSocket(t *testing.T) {
@@ -20,10 +21,10 @@ func TestWebSocket(t *testing.T) {
 	}
 
 	handler := mid(EmptyNext)
-	myHandler, ok := handler.(websockets.WebSockets)
+	myHandler, ok := handler.(websocket.WebSocket)
 
 	if !ok {
-		t.Fatalf("Expected handler to be type WebSockets, got: %#v", handler)
+		t.Fatalf("Expected handler to be type WebSocket, got: %#v", handler)
 	}
 
 	if myHandler.Sockets[0].Path != "/" {
@@ -38,15 +39,15 @@ func TestWebSocketParse(t *testing.T) {
 	tests := []struct {
 		inputWebSocketConfig    string
 		shouldErr               bool
-		expectedWebSocketConfig []websockets.Config
+		expectedWebSocketConfig []websocket.Config
 	}{
-		{`websocket /api1 cat`, false, []websockets.Config{{
+		{`websocket /api1 cat`, false, []websocket.Config{{
 			Path:    "/api1",
 			Command: "cat",
 		}}},
 
 		{`websocket /api3 cat  
-		  websocket /api4 cat `, false, []websockets.Config{{
+		  websocket /api4 cat `, false, []websocket.Config{{
 			Path:    "/api3",
 			Command: "cat",
 		}, {

middleware/websocket/websocket.go

+// Package websocket implements a WebSocket server by executing
+// a command and piping its input and output through the WebSocket
+// connection.
+package websocket
+
+import (
+	"io"
+	"net"
+	"net/http"
+	"os/exec"
+	"strings"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/mholt/caddy/middleware"
+)
+
+const (
+	// Time allowed to write a message to the peer.
+	writeWait = 10 * time.Second
+
+	// Time allowed to read the next pong message from the peer.
+	pongWait = 60 * time.Second
+
+	// Send pings to peer with this period. Must be less than pongWait.
+	pingPeriod = (pongWait * 9) / 10
+
+	// Maximum message size allowed from peer.
+	maxMessageSize = 1024 * 1024 * 10 // 10 MB default.
+)
+
+var (
+	// GatewayInterface is the dialect of CGI being used by the server
+	// to communicate with the script.  See CGI spec, 4.1.4
+	GatewayInterface string
+
+	// ServerSoftware is the name and version of the information server
+	// software making the CGI request.  See CGI spec, 4.1.17
+	ServerSoftware string
+)
+
+type (
+	// WebSocket is a type that holds configuration for the
+	// websocket middleware generally, like a list of all the
+	// websocket endpoints.
+	WebSocket struct {
+		// Next is the next HTTP handler in the chain for when the path doesn't match
+		Next middleware.Handler
+
+		// Sockets holds all the web socket endpoint configurations
+		Sockets []Config
+	}
+
+	// Config holds the configuration for a single websocket
+	// endpoint which may serve multiple websocket connections.
+	Config struct {
+		Path      string
+		Command   string
+		Arguments []string
+		Respawn   bool // TODO: Not used, but parser supports it until we decide on it
+	}
+)
+
+// ServeHTTP converts the HTTP request to a WebSocket connection and serves it up.
+func (ws WebSocket) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
+	for _, sockconfig := range ws.Sockets {
+		if middleware.Path(r.URL.Path).Matches(sockconfig.Path) {
+			return serveWS(w, r, &sockconfig)
+		}
+	}
+
+	// Didn't match a websocket path, so pass-thru
+	return ws.Next.ServeHTTP(w, r)
+}
+
+// serveWS is used for setting and upgrading the HTTP connection to a websocket connection.
+// It also spawns the child process that is associated with matched HTTP path/url.
+func serveWS(w http.ResponseWriter, r *http.Request, config *Config) (int, error) {
+	upgrader := websocket.Upgrader{
+		ReadBufferSize:  1024,
+		WriteBufferSize: 1024,
+		CheckOrigin:     func(r *http.Request) bool { return true },
+	}
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		return 0, err
+	}
+	defer conn.Close()
+
+	cmd := exec.Command(config.Command, config.Arguments...)
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		panic(err) // TODO
+	}
+
+	stdin, err := cmd.StdinPipe()
+	if err != nil {
+		panic(err) // TODO
+	}
+
+	metavars, err := buildEnv(cmd.Path, r)
+	if err != nil {
+		panic(err) // TODO
+	}
+
+	cmd.Env = metavars
+
+	if err := cmd.Start(); err != nil {
+		panic(err)
+	}
+
+	reader(conn, stdout, stdin)
+
+	return 0, nil // we shouldn't get here.
+}
+
+// buildEnv creates the meta-variables for the child process according
+// to the CGI 1.1 specification: http://tools.ietf.org/html/rfc3875#section-4.1
+// cmdPath should be the path of the command being run.
+// The returned string slice can be set to the command's Env property.
+func buildEnv(cmdPath string, r *http.Request) (metavars []string, err error) {
+	remoteHost, remotePort, err := net.SplitHostPort(r.RemoteAddr)
+	if err != nil {
+		return
+	}
+
+	serverHost, serverPort, err := net.SplitHostPort(r.Host)
+	if err != nil {
+		return
+	}
+
+	metavars = []string{
+		`AUTH_TYPE=`,      // Not used
+		`CONTENT_LENGTH=`, // Not used
+		`CONTENT_TYPE=`,   // Not used
+		`GATEWAY_INTERFACE=` + GatewayInterface,
+		`PATH_INFO=`,       // TODO
+		`PATH_TRANSLATED=`, // TODO
+		`QUERY_STRING=` + r.URL.RawQuery,
+		`REMOTE_ADDR=` + remoteHost,
+		`REMOTE_HOST=` + remoteHost, // Host lookups are slow - don't do them
+		`REMOTE_IDENT=`,             // Not used
+		`REMOTE_PORT=` + remotePort,
+		`REMOTE_USER=`, // Not used,
+		`REQUEST_METHOD=` + r.Method,
+		`REQUEST_URI=` + r.RequestURI,
+		`SCRIPT_NAME=` + cmdPath, // path of the program being executed
+		`SERVER_NAME=` + serverHost,
+		`SERVER_PORT=` + serverPort,
+		`SERVER_PROTOCOL=` + r.Proto,
+		`SERVER_SOFTWARE=` + ServerSoftware,
+	}
+
+	// Add each HTTP header to the environment as well
+	for header, values := range r.Header {
+		value := strings.Join(values, ", ")
+		header = strings.ToUpper(header)
+		header = strings.Replace(header, "-", "_", -1)
+		value = strings.Replace(value, "\n", " ", -1)
+		metavars = append(metavars, "HTTP_"+header+"="+value)
+	}
+
+	return
+}
+
+// reader is the guts of this package. It takes the stdin and stdout pipes
+// of the cmd we created in ServeWS and pipes them between the client and server
+// over websockets.
+func reader(conn *websocket.Conn, stdout io.ReadCloser, stdin io.WriteCloser) {
+	// Setup our connection's websocket ping/pong handlers from our const values.
+	conn.SetReadLimit(maxMessageSize)
+	conn.SetReadDeadline(time.Now().Add(pongWait))
+	conn.SetPongHandler(func(string) error { conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
+	go ticker(conn)
+
+	for {
+		msgType, r, err := conn.NextReader()
+		if err != nil {
+			if msgType == -1 {
+				return // we are done, as we got a close method.
+			}
+			panic(err) // TODO do something else here.
+		}
+
+		w, err := conn.NextWriter(msgType)
+		if err != nil {
+			panic(err) // TODO do something else here.
+		}
+
+		if _, err := io.Copy(stdin, r); err != nil {
+			panic(err) // TODO do something else here.
+		}
+
+		go func() {
+			if _, err := io.Copy(w, stdout); err != nil {
+				panic(err) // TODO do something else here.
+			}
+			if err := w.Close(); err != nil {
+				panic(err) // TODO do something else here.
+			}
+		}()
+	}
+}
+
+// ticker is start by the reader. Basically it is the method that simulates the websocket
+// between the server and client to keep it alive with ping messages.
+func ticker(conn *websocket.Conn) {
+	ticker := time.NewTicker(pingPeriod)
+	defer func() {
+		ticker.Stop()
+		conn.WriteMessage(websocket.CloseMessage, nil)
+	}()
+
+	for { // blocking loop with select to wait for stimulation.
+		select {
+		case <-ticker.C:
+			conn.WriteMessage(websocket.PingMessage, nil)
+		}
+	}
+}

middleware/websockets/websocket.go

-package websockets
-
-import (
-	"net"
-	"net/http"
-	"os/exec"
-	"strings"
-
-	"golang.org/x/net/websocket"
-)
-
-// WebSocket represents a web socket server instance. A WebSocket
-// is instantiated for each new websocket request/connection.
-type WebSocket struct {
-	Config
-	*http.Request
-}
-
-// 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
-	cmd.Stderr = conn // TODO: Make this configurable from the Caddyfile
-
-	metavars, err := ws.buildEnv(cmd.Path)
-	if err != nil {
-		panic(err) // TODO
-	}
-
-	cmd.Env = metavars
-
-	err = cmd.Run()
-	if err != nil {
-		panic(err)
-	}
-}
-
-// buildEnv creates the meta-variables for the child process according
-// to the CGI 1.1 specification: http://tools.ietf.org/html/rfc3875#section-4.1
-// cmdPath should be the path of the command being run.
-// The returned string slice can be set to the command's Env property.
-func (ws WebSocket) buildEnv(cmdPath string) (metavars []string, err error) {
-	remoteHost, remotePort, err := net.SplitHostPort(ws.RemoteAddr)
-	if err != nil {
-		return
-	}
-
-	serverHost, serverPort, err := net.SplitHostPort(ws.Host)
-	if err != nil {
-		return
-	}
-
-	metavars = []string{
-		`AUTH_TYPE=`,      // Not used
-		`CONTENT_LENGTH=`, // Not used
-		`CONTENT_TYPE=`,   // Not used
-		`GATEWAY_INTERFACE=` + GatewayInterface,
-		`PATH_INFO=`,       // TODO
-		`PATH_TRANSLATED=`, // TODO
-		`QUERY_STRING=` + ws.URL.RawQuery,
-		`REMOTE_ADDR=` + remoteHost,
-		`REMOTE_HOST=` + remoteHost, // Host lookups are slow - don't do them
-		`REMOTE_IDENT=`,             // Not used
-		`REMOTE_PORT=` + remotePort,
-		`REMOTE_USER=`, // Not used,
-		`REQUEST_METHOD=` + ws.Method,
-		`REQUEST_URI=` + ws.RequestURI,
-		`SCRIPT_NAME=` + cmdPath, // path of the program being executed
-		`SERVER_NAME=` + serverHost,
-		`SERVER_PORT=` + serverPort,
-		`SERVER_PROTOCOL=` + ws.Proto,
-		`SERVER_SOFTWARE=` + ServerSoftware,
-	}
-
-	// Add each HTTP header to the environment as well
-	for header, values := range ws.Header {
-		value := strings.Join(values, ", ")
-		header = strings.ToUpper(header)
-		header = strings.Replace(header, "-", "_", -1)
-		value = strings.Replace(value, "\n", " ", -1)
-		metavars = append(metavars, "HTTP_"+header+"="+value)
-	}
-
-	return
-}

middleware/websockets/websockets.go

-// Package websockets implements a WebSocket server by executing
-// a command and piping its input and output through the WebSocket
-// connection.
-package websockets
-
-import (
-	"net/http"
-
-	"github.com/mholt/caddy/middleware"
-	"golang.org/x/net/websocket"
-)
-
-type (
-	// WebSockets is a type that holds configuration for the
-	// websocket middleware generally, like a list of all the
-	// websocket endpoints.
-	WebSockets struct {
-		// Next is the next HTTP handler in the chain for when the path doesn't match
-		Next middleware.Handler
-
-		// Sockets holds all the web socket endpoint configurations
-		Sockets []Config
-	}
-
-	// Config holds the configuration for a single websocket
-	// endpoint which may serve multiple websocket connections.
-	Config struct {
-		Path      string
-		Command   string
-		Arguments []string
-		Respawn   bool // TODO: Not used, but parser supports it until we decide on it
-	}
-)
-
-// ServeHTTP converts the HTTP request to a WebSocket connection and serves it up.
-func (ws WebSockets) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
-	for _, sockconfig := range ws.Sockets {
-		if middleware.Path(r.URL.Path).Matches(sockconfig.Path) {
-			socket := WebSocket{
-				Config:  sockconfig,
-				Request: r,
-			}
-			websocket.Handler(socket.Handle).ServeHTTP(w, r)
-			return 0, nil
-		}
-	}
-
-	// Didn't match a websocket path, so pass-thru
-	return ws.Next.ServeHTTP(w, r)
-}
-
-var (
-	// GatewayInterface is the dialect of CGI being used by the server
-	// to communicate with the script.  See CGI spec, 4.1.4
-	GatewayInterface string
-
-	// ServerSoftware is the name and version of the information server
-	// software making the CGI request.  See CGI spec, 4.1.17
-	ServerSoftware string
-)