document_transport.go

package proxy

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
)

type bodyTransformation func(body io.ReadCloser, req *http.Request) (io.ReadCloser, error)

func newDocumentTransport(inner http.RoundTripper) *documentTransport {
	return &documentTransport{inner, replaceImagesURLProtocol()}
}

type documentTransport struct {
	transport     http.RoundTripper
	transformBody bodyTransformation
}

func (t *documentTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	res, err := t.transport.RoundTrip(req)
	if err != nil {
		return nil, err
	}

	if res.StatusCode >= 200 && res.StatusCode < 400 {
		res.Header.Set("Cache-Control", "max-age="+strconv.Itoa(cacheDuration))
	}

	body, err := t.transformBody(res.Body, req)
	if err != nil {
		return nil, err
	}

	res.Body = body

	return res, nil
}

func replaceImagesURLProtocol() bodyTransformation {
	return func(body io.ReadCloser, req *http.Request) (io.ReadCloser, error) {
		content, err := ioutil.ReadAll(body)
		if err != nil {
			return nil, err
		}
		regex := regexp.MustCompile(`(?m)"url":"http(:[^"]+)"`)
		fixed := regex.ReplaceAllString(string(content), `"url":"https$1"`)
		log.Info(fixed)

		return ioutil.NopCloser(bytes.NewReader([]byte(fixed))), nil
	}
}