From 322f909edbca5c1385a8f89b833afc07f7790539 Mon Sep 17 00:00:00 2001
From: Areg Harutyunyan <areg@cloudflare.com>
Date: Thu, 9 Apr 2020 21:59:15 +0100
Subject: [PATCH] TUN-2880: Return metadata about source of the response from
 cloudflared

---
 h2mux/header.go                 | 21 +++++++++++++++++++++
 origin/tunnel.go                |  9 +++++++--
 originservice/originservice.go  |  7 +++++--
 streamhandler/stream_handler.go |  5 +++--
 4 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/h2mux/header.go b/h2mux/header.go
index dbec394f..bc28c371 100644
--- a/h2mux/header.go
+++ b/h2mux/header.go
@@ -2,6 +2,7 @@ package h2mux
 
 import (
 	"encoding/base64"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -20,6 +21,10 @@ var headerEncoding = base64.RawStdEncoding
 const (
 	RequestUserHeadersField  = "cf-cloudflared-request-headers"
 	ResponseUserHeadersField = "cf-cloudflared-response-headers"
+
+	ResponseMetaHeaderField   = "cf-cloudflared-response-meta"
+	ResponseSourceCloudflared = "cloudflared"
+	ResponseSourceOrigin      = "origin"
 )
 
 // H2RequestHeadersToH1Request converts the HTTP/2 headers coming from origintunneld
@@ -226,3 +231,19 @@ func CreateSerializedHeaders(headersField string, headers ...http.Header) []Head
 		strings.Join(serializedHeaderChunks, ";"),
 	}}
 }
+
+type responseMetaHeader struct {
+	Source string `json:"src"`
+}
+
+func CreateResponseMetaHeader(source string) Header {
+	jsonResponseMetaHeader, err := json.Marshal(responseMetaHeader{Source: source})
+	if err != nil {
+		panic(err)
+	}
+
+	return Header{
+		Name:  ResponseMetaHeaderField,
+		Value: string(jsonResponseMetaHeader),
+	}
+}
diff --git a/origin/tunnel.go b/origin/tunnel.go
index 3c44e166..a11a514d 100644
--- a/origin/tunnel.go
+++ b/origin/tunnel.go
@@ -688,7 +688,9 @@ func (h *TunnelHandler) serveHTTP(stream *h2mux.MuxedStream, req *http.Request)
 	}
 	defer response.Body.Close()
 
-	err = stream.WriteHeaders(h2mux.H1ResponseToH2ResponseHeaders(response))
+	headers := h2mux.H1ResponseToH2ResponseHeaders(response)
+	headers = append(headers, h2mux.CreateResponseMetaHeader(h2mux.ResponseSourceOrigin))
+	err = stream.WriteHeaders(headers)
 	if err != nil {
 		return nil, errors.Wrap(err, "Error writing response header")
 	}
@@ -725,7 +727,10 @@ func (h *TunnelHandler) isEventStream(response *http.Response) bool {
 
 func (h *TunnelHandler) logError(stream *h2mux.MuxedStream, err error) {
 	h.logger.WithError(err).Error("HTTP request error")
-	stream.WriteHeaders([]h2mux.Header{{Name: ":status", Value: "502"}})
+	stream.WriteHeaders([]h2mux.Header{
+		{Name: ":status", Value: "502"},
+		h2mux.CreateResponseMetaHeader(h2mux.ResponseSourceCloudflared),
+	})
 	stream.Write([]byte("502 Bad Gateway"))
 	h.metrics.incrementResponses(h.connectionID, "502")
 }
diff --git a/originservice/originservice.go b/originservice/originservice.go
index 1c90dfee..86a532f9 100644
--- a/originservice/originservice.go
+++ b/originservice/originservice.go
@@ -17,7 +17,6 @@ import (
 	"github.com/cloudflare/cloudflared/hello"
 	"github.com/cloudflare/cloudflared/log"
 	"github.com/cloudflare/cloudflared/websocket"
-
 	"github.com/pkg/errors"
 )
 
@@ -47,6 +46,8 @@ func NewHTTPService(transport http.RoundTripper, url *url.URL, chunkedEncoding b
 }
 
 func (hc *HTTPService) Proxy(stream *h2mux.MuxedStream, req *http.Request) (*http.Response, error) {
+	const responseSourceOrigin = "origin"
+
 	// Support for WSGI Servers by switching transfer encoding from chunked to gzip/deflate
 	if !hc.chunkedEncoding {
 		req.TransferEncoding = []string{"gzip", "deflate"}
@@ -65,7 +66,9 @@ func (hc *HTTPService) Proxy(stream *h2mux.MuxedStream, req *http.Request) (*htt
 	}
 	defer resp.Body.Close()
 
-	err = stream.WriteHeaders(h1ResponseToH2Response(resp))
+	responseHeaders := h1ResponseToH2Response(resp)
+	responseHeaders = append(responseHeaders, h2mux.CreateResponseMetaHeader(responseSourceOrigin))
+	err = stream.WriteHeaders(responseHeaders)
 	if err != nil {
 		return nil, errors.Wrap(err, "error writing response header to HTTP origin")
 	}
diff --git a/streamhandler/stream_handler.go b/streamhandler/stream_handler.go
index 50ea68d9..aef1609d 100644
--- a/streamhandler/stream_handler.go
+++ b/streamhandler/stream_handler.go
@@ -180,11 +180,12 @@ func (s *StreamHandler) requestLogger(req *http.Request, tunnelHostname h2mux.Tu
 }
 
 func (s *StreamHandler) writeErrorStatus(stream *h2mux.MuxedStream, status *httpErrorStatus) {
-	stream.WriteHeaders([]h2mux.Header{
+	_ = stream.WriteHeaders([]h2mux.Header{
 		{
 			Name:  statusPseudoHeader,
 			Value: status.status,
 		},
+		h2mux.CreateResponseMetaHeader(h2mux.ResponseSourceCloudflared),
 	})
-	stream.Write(status.text)
+	_, _ = stream.Write(status.text)
 }