websocket: Simple buildEnv test, and fix for addresses without port

middleware/websocket/websocket.go

@@ -149,11 +149,17 @@ func serveWS(w http.ResponseWriter, r *http.Request, config *Config) (int, error
 // 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) {
+	if !strings.Contains(r.RemoteAddr, ":") {
+		r.RemoteAddr += ":"
+	}
 	remoteHost, remotePort, err := net.SplitHostPort(r.RemoteAddr)
 	if err != nil {
 		return
 	}
 
+	if !strings.Contains(r.Host, ":") {
+		r.Host += ":"
+	}
 	serverHost, serverPort, err := net.SplitHostPort(r.Host)
 	if err != nil {
 		return

middleware/websocket/websocket_test.go

@@ -0,0 +1,22 @@
+package websocket
+
+import (
+	"net/http"
+	"testing"
+)
+
+func TestBuildEnv(t *testing.T) {
+	req, err := http.NewRequest("GET", "http://localhost", nil)
+	if err != nil {
+		t.Fatal("Error setting up request:", err)
+	}
+	req.RemoteAddr = "localhost:50302"
+
+	env, err := buildEnv("/bin/command", req)
+	if err != nil {
+		t.Fatal("Didn't expect an error:", err)
+	}
+	if len(env) == 0 {
+		t.Fatal("Expected non-empty environment; got %#v", env)
+	}
+}