From a368fbbe9b4ec7ebe2bd81b0c4903ab444a35ba5 Mon Sep 17 00:00:00 2001
From: Dalton <dalton@cloudflare.com>
Date: Mon, 23 Mar 2020 10:22:58 -0500
Subject: [PATCH] AUTH-2394 fixed header for websockets. Added TCP alias

---
 cmd/cloudflared/access/cmd.go |  2 +-
 cmd/cloudflared/tunnel/cmd.go |  2 ++
 h2mux/header.go               | 17 +++++++++++++++--
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/cmd/cloudflared/access/cmd.go b/cmd/cloudflared/access/cmd.go
index f1c01fe5..1ea949cc 100644
--- a/cmd/cloudflared/access/cmd.go
+++ b/cmd/cloudflared/access/cmd.go
@@ -119,7 +119,7 @@ func Commands() []*cli.Command {
 				{
 					Name:        "ssh",
 					Action:      ssh,
-					Aliases:     []string{"rdp"},
+					Aliases:     []string{"rdp", "tcp"},
 					Usage:       "",
 					ArgsUsage:   "",
 					Description: `The ssh subcommand sends data over a proxy to the Cloudflare edge.`,
diff --git a/cmd/cloudflared/tunnel/cmd.go b/cmd/cloudflared/tunnel/cmd.go
index 9a43fb46..e901c0bd 100644
--- a/cmd/cloudflared/tunnel/cmd.go
+++ b/cmd/cloudflared/tunnel/cmd.go
@@ -602,6 +602,8 @@ func hostnameFromURI(uri string) string {
 		return addPortIfMissing(u, 22)
 	case "rdp":
 		return addPortIfMissing(u, 3389)
+	case "tcp":
+		return addPortIfMissing(u, 7864) // just a random port since there isn't a default in this case
 	}
 	return ""
 }
diff --git a/h2mux/header.go b/h2mux/header.go
index 4a853c09..c79846bf 100644
--- a/h2mux/header.go
+++ b/h2mux/header.go
@@ -3,11 +3,12 @@ package h2mux
 import (
 	"encoding/base64"
 	"fmt"
-	"github.com/pkg/errors"
 	"net/http"
 	"net/url"
 	"strconv"
 	"strings"
+
+	"github.com/pkg/errors"
 )
 
 type Header struct {
@@ -71,6 +72,9 @@ func H2RequestHeadersToH1Request(h2 []Header, h1 *http.Request) error {
 				return fmt.Errorf("unparseable content length")
 			}
 			h1.ContentLength = contentLength
+		case "connection", "upgrade":
+			// for websocket header support
+			h1.Header.Add(http.CanonicalHeaderKey(header.Name), header.Value)
 		default:
 			// Ignore any other header;
 			// User headers will be read from `RequestUserHeadersField`
@@ -109,6 +113,15 @@ func IsControlHeader(headerName string) bool {
 		strings.HasPrefix(headerName, "cf-")
 }
 
+// IsWebsocketClientHeader returns true if the header name is required by the client to upgrade properly
+func IsWebsocketClientHeader(headerName string) bool {
+	headerName = strings.ToLower(headerName)
+
+	return headerName == "sec-websocket-accept" ||
+		headerName == "connection" ||
+		headerName == "upgrade"
+}
+
 func H1ResponseToH2ResponseHeaders(h1 *http.Response) (h2 []Header) {
 	h2 = []Header{
 		{Name: ":status", Value: strconv.Itoa(h1.StatusCode)},
@@ -122,7 +135,7 @@ func H1ResponseToH2ResponseHeaders(h1 *http.Response) (h2 []Header) {
 
 				// Since these are http2 headers, they're required to be lowercase
 				h2 = append(h2, Header{Name: strings.ToLower(header), Value: value})
-			} else if !IsControlHeader(header) {
+			} else if !IsControlHeader(header) || IsWebsocketClientHeader(header) {
 				// User headers, on the other hand, must all be serialized so that
 				// HTTP/2 header validation won't be applied to HTTP/1 header values
 				if _, ok := userHeaders[header]; ok {