Add static pages to PageWriter interface

oauth2-proxy · Mar 21, 2021 · 9782fc7 · 9782fc7
1 parent c1267bb
commit 9782fc7
Show file tree

Hide file tree

Showing 4 changed files with 261 additions and 0 deletions.
diff --git a/pkg/app/pagewriter/pagewriter.go b/pkg/app/pagewriter/pagewriter.go
@@ -13,12 +13,14 @@ type Writer interface {
 	WriteSignInPage(rw http.ResponseWriter, req *http.Request, redirectURL string)
 	WriteErrorPage(rw http.ResponseWriter, opts ErrorPageOpts)
 	ProxyErrorHandler(rw http.ResponseWriter, req *http.Request, proxyErr error)
+	WriteRobotsTxt(rw http.ResponseWriter, req *http.Request)
 }
 
 // pageWriter implements the Writer interface
 type pageWriter struct {
 	*errorPageWriter
 	*signInPageWriter
+	*staticPageWriter
 }
 
 // Opts contains all options required to configure the template
@@ -88,8 +90,14 @@ func NewWriter(opts Opts) (Writer, error) {
 		logoData:         logoData,
 	}
 
+	staticPages, err := newStaticPageWriter(opts.TemplatesPath, errorPage)
+	if err != nil {
+		return nil, fmt.Errorf("error loading static page writer: %v", err)
+	}
+
 	return &pageWriter{
 		errorPageWriter:  errorPage,
 		signInPageWriter: signInPage,
+		staticPageWriter: staticPages,
 	}, nil
 }
diff --git a/pkg/app/pagewriter/robots.txt b/pkg/app/pagewriter/robots.txt
@@ -0,0 +1,2 @@
+User-agent: *
+Disallow: /
diff --git a/pkg/app/pagewriter/static_pages.go b/pkg/app/pagewriter/static_pages.go
@@ -0,0 +1,98 @@
+package pagewriter
+
+import (
+	// Import embed to allow importing default page templates
+	_ "embed"
+	"fmt"
+	"net/http"
+	"os"
+	"path/filepath"
+
+	middlewareapi "github.com/oauth2-proxy/oauth2-proxy/v7/pkg/apis/middleware"
+	"github.com/oauth2-proxy/oauth2-proxy/v7/pkg/logger"
+)
+
+const (
+	robotsTxtName = "robots.txt"
+)
+
+//go:embed robots.txt
+var defaultRobotsTxt []byte
+
+// staticPageWriter is used to write static pages.
+type staticPageWriter struct {
+	pages           map[string][]byte
+	errorPageWriter *errorPageWriter
+}
+
+// WriteRobotsTxt writes the robots.txt content to the response writer.
+func (s *staticPageWriter) WriteRobotsTxt(rw http.ResponseWriter, req *http.Request) {
+	s.writePage(rw, req, robotsTxtName)
+}
+
+// writePage writes the content of the page to the response writer.
+func (s *staticPageWriter) writePage(rw http.ResponseWriter, req *http.Request, pageName string) {
+	content, ok := s.pages[pageName]
+	if !ok {
+		// If the page isn't regiested, something went wrong and there is a bug.
+		// Tests should make sure this code path is never hit.
+		panic(fmt.Sprintf("Static page %q not found", pageName))
+	}
+	_, err := rw.Write(content)
+	if err != nil {
+		logger.Printf("Error writing %q: %v", pageName, err)
+		scope := middlewareapi.GetRequestScope(req)
+		s.errorPageWriter.WriteErrorPage(rw, ErrorPageOpts{
+			Status:    http.StatusInternalServerError,
+			RequestID: scope.RequestID,
+			AppError:  err.Error(),
+		})
+		return
+	}
+}
+
+func newStaticPageWriter(customDir string, errorWriter *errorPageWriter) (*staticPageWriter, error) {
+	pages, err := loadStaticPages(customDir)
+	if err != nil {
+		return nil, fmt.Errorf("could not load static pages: %v", err)
+	}
+
+	return &staticPageWriter{
+		pages:           pages,
+		errorPageWriter: errorWriter,
+	}, nil
+}
+
+// loadStaticPages loads static page content from the custom directory provided.
+// If any file is not provided in the custom directory, the default will be used
+// instead.
+// Statis files include:
+// - robots.txt
+func loadStaticPages(customDir string) (map[string][]byte, error) {
+	pages := make(map[string][]byte)
+
+	if err := addStaticPage(pages, customDir, robotsTxtName, defaultRobotsTxt); err != nil {
+		return nil, fmt.Errorf("could not add robots.txt: %v", err)
+	}
+
+	return pages, nil
+}
+
+// addStaticPage tries to load the named file from the custom directory.
+// If no custom directory is provided, the default content is used instead.
+func addStaticPage(pages map[string][]byte, customDir, fileName string, defaultContent []byte) error {
+	filePath := filepath.Join(customDir, fileName)
+	if customDir != "" && isFile(filePath) {
+		content, err := os.ReadFile(filePath)
+		if err != nil {
+			return fmt.Errorf("could not read file: %v", err)
+		}
+
+		pages[fileName] = content
+		return nil
+	}
+
+	// No custom content defined, use the default.
+	pages[fileName] = defaultContent
+	return nil
+}
diff --git a/pkg/app/pagewriter/static_pages_test.go b/pkg/app/pagewriter/static_pages_test.go
@@ -0,0 +1,153 @@
+package pagewriter
+
+import (
+	"errors"
+	"html/template"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+
+	middlewareapi "github.com/oauth2-proxy/oauth2-proxy/v7/pkg/apis/middleware"
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("Static Pages", func() {
+	var customDir string
+	const customRobots = "I AM A ROBOT!!!"
+	var errorPage *errorPageWriter
+	var request *http.Request
+
+	BeforeEach(func() {
+		errorTmpl, err := template.New("").Parse("{{.Title}}")
+		Expect(err).ToNot(HaveOccurred())
+		errorPage = &errorPageWriter{
+			template: errorTmpl,
+		}
+
+		customDir, err = ioutil.TempDir("", "oauth2-proxy-static-pages-test")
+		Expect(err).ToNot(HaveOccurred())
+
+		robotsTxtFile := filepath.Join(customDir, robotsTxtName)
+		Expect(ioutil.WriteFile(robotsTxtFile, []byte(customRobots), 0400)).To(Succeed())
+
+		request = httptest.NewRequest("", "http://127.0.0.1/", nil)
+		request = middlewareapi.AddRequestScope(request, &middlewareapi.RequestScope{
+			RequestID: testRequestID,
+		})
+	})
+
+	AfterEach(func() {
+		Expect(os.RemoveAll(customDir)).To(Succeed())
+	})
+
+	Context("Static Page Writer", func() {
+		Context("With custom content", func() {
+			var pageWriter *staticPageWriter
+
+			BeforeEach(func() {
+				var err error
+				pageWriter, err = newStaticPageWriter(customDir, errorPage)
+				Expect(err).ToNot(HaveOccurred())
+			})
+
+			Context("WriterRobotsTxt", func() {
+				It("Should write the custom robots txt", func() {
+					recorder := httptest.NewRecorder()
+					pageWriter.WriteRobotsTxt(recorder, request)
+
+					body, err := ioutil.ReadAll(recorder.Result().Body)
+					Expect(err).ToNot(HaveOccurred())
+					Expect(string(body)).To(Equal(customRobots))
+
+					Expect(recorder.Result().StatusCode).To(Equal(http.StatusOK))
+				})
+			})
+		})
+
+		Context("Without custom content", func() {
+			var pageWriter *staticPageWriter
+
+			BeforeEach(func() {
+				var err error
+				pageWriter, err = newStaticPageWriter("", errorPage)
+				Expect(err).ToNot(HaveOccurred())
+			})
+
+			Context("WriterRobotsTxt", func() {
+				It("Should write the custom robots txt", func() {
+					recorder := httptest.NewRecorder()
+					pageWriter.WriteRobotsTxt(recorder, request)
+
+					body, err := ioutil.ReadAll(recorder.Result().Body)
+					Expect(err).ToNot(HaveOccurred())
+					Expect(string(body)).To(Equal(string(defaultRobotsTxt)))
+
+					Expect(recorder.Result().StatusCode).To(Equal(http.StatusOK))
+				})
+
+				It("Should serve an error if it cannot write the page", func() {
+					recorder := &testBadResponseWriter{
+						ResponseRecorder: httptest.NewRecorder(),
+					}
+					pageWriter.WriteRobotsTxt(recorder, request)
+
+					body, err := ioutil.ReadAll(recorder.Result().Body)
+					Expect(err).ToNot(HaveOccurred())
+					Expect(string(body)).To(Equal(string("Internal Server Error")))
+
+					Expect(recorder.Result().StatusCode).To(Equal(http.StatusInternalServerError))
+				})
+			})
+		})
+	})
+
+	Context("loadStaticPages", func() {
+		Context("With custom content", func() {
+			Context("And a custom robots txt", func() {
+				It("Loads the custom content", func() {
+					pages, err := loadStaticPages(customDir)
+					Expect(err).ToNot(HaveOccurred())
+					Expect(pages).To(HaveLen(1))
+					Expect(pages).To(HaveKeyWithValue(robotsTxtName, []byte(customRobots)))
+				})
+			})
+
+			Context("And no custom robots txt", func() {
+				It("returns the default content", func() {
+					robotsTxtFile := filepath.Join(customDir, robotsTxtName)
+					Expect(os.Remove(robotsTxtFile)).To(Succeed())
+
+					pages, err := loadStaticPages(customDir)
+					Expect(err).ToNot(HaveOccurred())
+					Expect(pages).To(HaveLen(1))
+					Expect(pages).To(HaveKeyWithValue(robotsTxtName, defaultRobotsTxt))
+				})
+			})
+		})
+
+		Context("Without custom content", func() {
+			It("Loads the default content", func() {
+				pages, err := loadStaticPages("")
+				Expect(err).ToNot(HaveOccurred())
+				Expect(pages).To(HaveLen(1))
+				Expect(pages).To(HaveKeyWithValue(robotsTxtName, defaultRobotsTxt))
+			})
+		})
+	})
+})
+
+type testBadResponseWriter struct {
+	*httptest.ResponseRecorder
+	firstWriteCalled bool
+}
+
+func (b *testBadResponseWriter) Write(buf []byte) (int, error) {
+	if !b.firstWriteCalled {
+		b.firstWriteCalled = true
+		return 0, errors.New("write closed")
+	}
+	return b.ResponseRecorder.Write(buf)
+}