oneandone-billing-mailer/cmd/test/server.go

package main

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"

	"github.com/go-logr/logr"
	"golang.org/x/net/publicsuffix"
)

type OneAndOneLoginProxy struct {
	httpClient    *http.Client
	listener      net.Listener
	cookieJar     http.CookieJar
	responseCodeC chan url.Values
	logger        logr.Logger
}

func NewOneAndOneLoginProxy(
	ctx context.Context,
	listenAddr string,
) (*OneAndOneLoginProxy, error) {
	// create temporary cookie jar for login flow
	cookieJar, err := cookiejar.New(&cookiejar.Options{
		PublicSuffixList: publicsuffix.List,
	})
	if err != nil {
		return nil, err
	}

	// channel to return response code on
	responseCodeC := make(chan url.Values, 1)

	p := &OneAndOneLoginProxy{
		cookieJar: cookieJar,
		httpClient: &http.Client{
			CheckRedirect: func(req *http.Request, via []*http.Request) error {
				return http.ErrUseLastResponse
			},
		},
		logger:        logr.FromContextOrDiscard(ctx).WithName("OneAndOneLoginProxy"),
		responseCodeC: responseCodeC,
	}
	if err := p.setupListener(listenAddr); err != nil {
		return nil, err
	}
	return p, nil
}

func (p *OneAndOneLoginProxy) ResponseCodeC() <-chan url.Values {
	return p.responseCodeC
}

func (p *OneAndOneLoginProxy) Host() string {
	addr := p.listener.Addr().(*net.TCPAddr)
	return net.JoinHostPort(
		addr.IP.String(),
		fmt.Sprintf("%d", addr.Port))
}

func (p *OneAndOneLoginProxy) setupListener(listenAddr string) error {
	if len(listenAddr) == 0 {
		// any random local port
		listenAddr = "localhost:0"
	}
	listener, err := net.Listen("tcp", listenAddr)
	if err != nil {
		return err
	}
	p.listener = listener

	p.logger.V(1).Info("Proxy listener set up",
		"addr", p.listener.Addr())

	return nil
}

func (p *OneAndOneLoginProxy) Teardown() error {
	var err error
	p.logger.V(2).Info("Teardown() called on proxy")
	if p.listener != nil {
		p.logger.V(1).Info("Tearing down proxy")
		err = p.listener.Close()
		p.listener = nil
	}
	return err
}

func (p *OneAndOneLoginProxy) Listen() error {
	p.logger.V(2).Info("Listen() called on proxy")
	p.logger.V(1).Info("Proxy now going to listen")
	return http.Serve(p.listener, p)
}

func convertToProxyPath(u *url.URL) string {
	// Enforce path at least containing /
	if !strings.HasPrefix(u.Path, "/") {
		u.Path = "/" + u.Path
	}

	return fmt.Sprintf("%s/%s%s", u.Scheme, u.Host, u.Path)
}

func (p *OneAndOneLoginProxy) ConvertURL(u *url.URL) *url.URL {
	return &url.URL{
		Scheme:      "http",
		Host:        p.Host(),
		Path:        convertToProxyPath(u),
		RawQuery:    u.RawQuery,
		RawFragment: u.RawFragment,
	}
}

func (p *OneAndOneLoginProxy) ConvertURLString(uStr string) (*url.URL, error) {
	u, err := url.Parse(uStr)
	if err != nil {
		return nil, err
	}
	return p.ConvertURL(u), nil
}

func (p *OneAndOneLoginProxy) decodeURL(u *url.URL) (proxyURL *url.URL) {
	reqURI := strings.TrimPrefix(u.Path, "/")
	reqURIParts := strings.SplitN(reqURI, "/", 3)
	proxyScheme := strings.ToLower(reqURIParts[0])
	if reqURIParts[0] != "http" && reqURIParts[0] != "https" {
		// assume this is a URL relative to https://account.1und1.de
		proxyURL = &url.URL{
			Scheme:   "https",
			Host:     "account.1und1.de",
			Path:     u.Path,
			RawQuery: u.RawQuery,
		}
	} else {
		// process as absolute URL
		proxyURL = &url.URL{
			Scheme:   proxyScheme,
			RawQuery: u.RawQuery,
		}
		if len(reqURIParts) >= 2 {
			proxyURL.Host = reqURIParts[1]
		}
		if len(reqURIParts) > 2 {
			proxyURL.Path = reqURIParts[2]
		}
	}
	return proxyURL
}

func (p *OneAndOneLoginProxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// create a new url from the raw RequestURI sent by the client
	// /{scheme}/{hostport}/{path/with/slashes} => {scheme}://{hostport}/{path/with/slashes}
	proxyURL := p.decodeURL(req.URL)

	// buffer body before forwarding or processing it further
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		p.logger.V(1).Error(err, "Failed to read request body",
			"method", req.Method,
			"url", req.URL)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// you can reassign the body if you need to parse it as multipart
	req.Body = ioutil.NopCloser(bytes.NewReader(body))

	// assemble real request
	proxyReq, err := http.NewRequest(
		req.Method,
		proxyURL.String(),
		bytes.NewReader(body))
	if err != nil {
		p.logger.V(1).Error(err, "Failed to create proxy request",
			"method", req.Method,
			"url", req.URL,
			"proxyURL", proxyURL.String())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// set headers removing everything that we set ourselves
	proxyReq.Header = make(http.Header)
	for key, values := range req.Header {
		for i, value := range values {
			switch http.CanonicalHeaderKey(key) {
			case "Cookie": // we add this from our own jar
				continue
			case "Origin", "Referer":
				u, err := url.Parse(value)
				if err == nil {
					realURL := p.decodeURL(u)
					values[i] = realURL.String()
				}
			}
			proxyReq.Header.Add(key, value)
		}
	}

	// load cookies from jar to request
	for _, c := range p.cookieJar.Cookies(proxyReq.URL) {
		// p.logger.V(1).Info("Proxy sends cookie", "cookie", c.Name)
		proxyReq.AddCookie(c)
	}

	resp, err := p.httpClient.Do(proxyReq)
	if err != nil {
		p.logger.V(1).Error(err, "Failed to do real request",
			"method", req.Method,
			"url", req.URL,
			"proxyURL", proxyURL.String())
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	// store cookies from response to jar
	p.cookieJar.SetCookies(proxyReq.URL, resp.Cookies())

	// copy response headers
	respondNoContent := false
	for key, values := range resp.Header {
		for _, value := range values {
			switch http.CanonicalHeaderKey(key) {
			case "Location":
				u, err := url.Parse(value)
				if err != nil {
					break
				}
				if u.IsAbs() {
					switch strings.ToLower(u.Scheme) {
					case "com.oneandone.controlcenter.android":
						// there's our oauth code, let's catch it before
						// anything bad happens
						p.logger.V(1).Info("Proxy found oauth response", "url", u.String())
						q := u.Query()

						// Special case for empty state where server
						// inexplicably responds with "state=null" which causes
						// 3LO state check in go oauth library to fail
						if q.Get("state") == "null" {
							q.Set("state", "")
						}

						p.responseCodeC <- q
						respondNoContent = true
					default:
						u = p.ConvertURL(u)
						value = u.String()
					}
				}
			case "Set-Cookie":
				// skip these headers as we already handled that ourselves
				continue
			}
			w.Header().Add(key, value)
		}
	}

	if respondNoContent {
		p.logger.V(1).Info("Serving incoming request without content",
			"method", req.Method,
			"url", req.URL,
			"proxyURL", proxyURL.String())
		w.WriteHeader(http.StatusNoContent)
		return
	}

	// copy response body
	p.logger.V(1).Info("Serving incoming request",
		"method", req.Method,
		"url", req.URL,
		"proxyURL", proxyURL.String(),
		"statusCode", resp.StatusCode,
		"status", resp.Status)
	w.WriteHeader(resp.StatusCode)
	contentType := resp.Header.Get("Content-Type")
	switch {
	case resp.Uncompressed &&
		(strings.Contains(contentType, "application/json") ||
			strings.Contains(contentType, "text/plain") ||
			strings.Contains(contentType, "text/html")):
		buf := new(bytes.Buffer)
		io.Copy(buf, resp.Body)
		bufStr := buf.String()
		io.Copy(w, strings.NewReader(bufStr))
	default:
		io.Copy(w, resp.Body)
	}
}