From bc015995d80f3e0d0a85746e433f88d35531fc5c Mon Sep 17 00:00:00 2001
From: Adam Chalmers <achalmers@cloudflare.com>
Date: Wed, 4 Nov 2020 12:22:21 -0600
Subject: [PATCH] TUN-3484: OriginService that responds with configured HTTP
 status

---
 ingress/ingress.go        |  8 ++++++++
 ingress/ingress_test.go   | 15 +++++++++++++++
 ingress/origin_service.go | 37 +++++++++++++++++++++++++++++++++++++
 ingress/rule_test.go      | 18 ++++++++++++++++++
 4 files changed, 78 insertions(+)

diff --git a/ingress/ingress.go b/ingress/ingress.go
index 614a6b5a..9c910c7d 100644
--- a/ingress/ingress.go
+++ b/ingress/ingress.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net/url"
 	"regexp"
+	"strconv"
 	"strings"
 	"sync"
 
@@ -129,6 +130,13 @@ func validate(ingress []config.UnvalidatedIngressRule, defaults OriginRequestCon
 			// No validation necessary for unix socket filepath services
 			path := strings.TrimPrefix(r.Service, prefix)
 			service = &unixSocketPath{path: path}
+		} else if prefix := "http_status:"; strings.HasPrefix(r.Service, prefix) {
+			status, err := strconv.Atoi(strings.TrimPrefix(r.Service, prefix))
+			if err != nil {
+				return Ingress{}, errors.Wrap(err, "invalid HTTP status")
+			}
+			srv := newStatusCode(status)
+			service = &srv
 		} else if r.Service == "hello_world" || r.Service == "hello-world" || r.Service == "helloworld" {
 			service = new(helloWorld)
 		} else {
diff --git a/ingress/ingress_test.go b/ingress/ingress_test.go
index 7c98bbb3..88b79d0e 100644
--- a/ingress/ingress_test.go
+++ b/ingress/ingress_test.go
@@ -24,6 +24,7 @@ ingress:
 func Test_parseIngress(t *testing.T) {
 	localhost8000 := MustParseURL(t, "https://localhost:8000")
 	localhost8001 := MustParseURL(t, "https://localhost:8001")
+	fourOhFour := newStatusCode(404)
 	defaultConfig := setConfig(originRequestFromYAML(config.OriginRequestConfig{}), config.OriginRequestConfig{})
 	require.Equal(t, defaultKeepAliveConnections, defaultConfig.KeepAliveConnections)
 	type args struct {
@@ -172,6 +173,20 @@ ingress:
 `},
 			wantErr: true,
 		},
+		{
+			name: "Valid HTTP status",
+			args: args{rawYAML: `
+ingress:
+ - service: http_status:404
+`},
+			want: []Rule{
+				{
+					Hostname: "",
+					Service:  &fourOhFour,
+					Config:   defaultConfig,
+				},
+			},
+		},
 		{
 			name: "Valid hello world service",
 			args: args{rawYAML: `
diff --git a/ingress/origin_service.go b/ingress/origin_service.go
index 7ac32bf4..2bca701f 100644
--- a/ingress/origin_service.go
+++ b/ingress/origin_service.go
@@ -223,6 +223,43 @@ func originRequiresProxy(staticHost string, cfg OriginRequestConfig) bool {
 	return staticHost != "" || cfg.BastionMode
 }
 
+// statusCode is an OriginService that just responds with a given HTTP status.
+// Typical use-case is "user wants the catch-all rule to just respond 404".
+type statusCode struct {
+	resp *http.Response
+}
+
+func newStatusCode(status int) statusCode {
+	resp := &http.Response{
+		StatusCode: status,
+		Status:     http.StatusText(status),
+		Body:       new(NopReadCloser),
+	}
+	return statusCode{resp: resp}
+}
+
+func (o *statusCode) String() string {
+	return fmt.Sprintf("HTTP %d", o.resp.StatusCode)
+}
+
+func (o *statusCode) start(wg *sync.WaitGroup, log logger.Service, shutdownC <-chan struct{}, errC chan error, cfg OriginRequestConfig) error {
+	return nil
+}
+
+func (o *statusCode) RoundTrip(_ *http.Request) (*http.Response, error) {
+	return o.resp, nil
+}
+
+type NopReadCloser struct{}
+
+func (nrc *NopReadCloser) Read(buf []byte) (int, error) {
+	return 0, nil
+}
+
+func (nrc *NopReadCloser) Close() error {
+	return nil
+}
+
 func newHTTPTransport(service OriginService, cfg OriginRequestConfig) (*http.Transport, error) {
 	originCertPool, err := tlsconfig.LoadOriginCA(cfg.CAPool, nil)
 	if err != nil {
diff --git a/ingress/rule_test.go b/ingress/rule_test.go
index ef908090..c42b1b65 100644
--- a/ingress/rule_test.go
+++ b/ingress/rule_test.go
@@ -4,6 +4,8 @@ import (
 	"net/url"
 	"regexp"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func Test_rule_matches(t *testing.T) {
@@ -117,3 +119,19 @@ func Test_rule_matches(t *testing.T) {
 		})
 	}
 }
+
+func TestStaticHTTPStatus(t *testing.T) {
+	o := newStatusCode(404)
+	buf := make([]byte, 100)
+
+	sendReq := func() {
+		resp, err := o.RoundTrip(nil)
+		require.NoError(t, err)
+		_, err = resp.Body.Read(buf)
+		require.NoError(t, err)
+		require.NoError(t, resp.Body.Close())
+		require.Equal(t, 404, resp.StatusCode)
+	}
+	sendReq()
+	sendReq()
+}