From d22f64e6d41117c242fc4ecd00179e68b8f118c5 Mon Sep 17 00:00:00 2001
From: Matthew Holt <mholt@users.noreply.github.com>
Date: Mon, 20 May 2019 15:46:47 -0600
Subject: Implement headers middleware

---
 modules/caddyhttp/headers/headers.go      | 82 +++++++++++++++++++++++++++++++
 modules/caddyhttp/headers/headers_test.go |  7 +++
 2 files changed, 89 insertions(+)
 create mode 100644 modules/caddyhttp/headers/headers.go
 create mode 100644 modules/caddyhttp/headers/headers_test.go

(limited to 'modules/caddyhttp/headers')

diff --git a/modules/caddyhttp/headers/headers.go b/modules/caddyhttp/headers/headers.go
new file mode 100644
index 0000000..1826c7a
--- /dev/null
+++ b/modules/caddyhttp/headers/headers.go
@@ -0,0 +1,82 @@
+package headers
+
+import (
+	"net/http"
+	"strings"
+
+	"bitbucket.org/lightcodelabs/caddy2"
+	"bitbucket.org/lightcodelabs/caddy2/modules/caddyhttp"
+)
+
+func init() {
+	caddy2.RegisterModule(caddy2.Module{
+		Name: "http.middleware.headers",
+		New:  func() (interface{}, error) { return new(Headers), nil },
+	})
+}
+
+// Headers is a middleware which can mutate HTTP headers.
+type Headers struct {
+	Request  HeaderOps
+	Response RespHeaderOps
+}
+
+// HeaderOps defines some operations to
+// perform on HTTP headers.
+type HeaderOps struct {
+	Add    http.Header
+	Set    http.Header
+	Delete []string
+}
+
+// RespHeaderOps is like HeaderOps, but
+// optionally deferred until response time.
+type RespHeaderOps struct {
+	HeaderOps
+	Deferred bool `json:"deferred"`
+}
+
+func (h Headers) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
+	apply(h.Request, r.Header)
+	if h.Response.Deferred {
+		w = &responseWriterWrapper{
+			ResponseWriterWrapper: &caddyhttp.ResponseWriterWrapper{ResponseWriter: w},
+			headerOps:             h.Response.HeaderOps,
+		}
+	} else {
+		apply(h.Response.HeaderOps, w.Header())
+	}
+	return next.ServeHTTP(w, r)
+}
+
+func apply(ops HeaderOps, hdr http.Header) {
+	for fieldName, vals := range ops.Add {
+		for _, v := range vals {
+			hdr.Add(fieldName, v)
+		}
+	}
+	for fieldName, vals := range ops.Set {
+		hdr.Set(fieldName, strings.Join(vals, ","))
+	}
+	for _, fieldName := range ops.Delete {
+		hdr.Del(fieldName)
+	}
+}
+
+// responseWriterWrapper defers response header
+// operations until WriteHeader is called.
+type responseWriterWrapper struct {
+	*caddyhttp.ResponseWriterWrapper
+	headerOps HeaderOps
+}
+
+func (rww *responseWriterWrapper) WriteHeader(status int) {
+	apply(rww.headerOps, rww.ResponseWriterWrapper.Header())
+	rww.ResponseWriterWrapper.WriteHeader(status)
+}
+
+// Interface guards
+var (
+	_ caddyhttp.MiddlewareHandler = (*Headers)(nil)
+	_ caddyhttp.HTTPInterfaces    = (*responseWriterWrapper)(nil)
+)
diff --git a/modules/caddyhttp/headers/headers_test.go b/modules/caddyhttp/headers/headers_test.go
new file mode 100644
index 0000000..cb83d47
--- /dev/null
+++ b/modules/caddyhttp/headers/headers_test.go
@@ -0,0 +1,7 @@
+package headers
+
+import "testing"
+
+func TestReqHeaders(t *testing.T) {
+	// TODO: write tests
+}
-- 
cgit v1.2.3