From f69a0b7e0e8436fefbe934c9236292ac745a1278 Mon Sep 17 00:00:00 2001
From: Andrew Gaul <andrew@gaul.org>
Date: Sun, 20 Aug 2017 21:55:20 -0700
Subject: [PATCH] Serve diagnostics and CA certificate

Fixes #21.  Reference:

https://mtersch.wordpress.com/2015/03/17/certificate-import-in-firefox-on-android/
---
 compy.go              |  2 +-
 compy_test.go         | 19 ++++++++++-
 proxy/mitmlistener.go | 14 +++++----
 proxy/proxy.go        | 73 +++++++++++++++++++++++++++++++++++++++++--
 4 files changed, 98 insertions(+), 10 deletions(-)

diff --git a/compy.go b/compy.go
index aaf90e8..9d10226 100644
--- a/compy.go
+++ b/compy.go
@@ -31,7 +31,7 @@ var (
 func main() {
 	flag.Parse()
 
-	p := proxy.New()
+	p := proxy.New(*host, cert)
 
 	if (*ca == "") != (*caKey == "") {
 		log.Fatalln("must specify both CA certificate and key")
diff --git a/compy_test.go b/compy_test.go
index f071fc4..887d890 100644
--- a/compy_test.go
+++ b/compy_test.go
@@ -37,7 +37,7 @@ var _ = Suite(&CompyTest{})
 func (s *CompyTest) SetUpSuite(c *C) {
 	s.server = httptest.NewServer(httpbin.GetMux())
 
-	s.proxy = proxy.New()
+	s.proxy = proxy.New("localhost"+*host, nil)
 	s.proxy.AddTranscoder("image/gif", &tc.Gif{})
 	s.proxy.AddTranscoder("image/jpeg", tc.NewJpeg(50))
 	s.proxy.AddTranscoder("image/png", &tc.Png{})
@@ -243,3 +243,20 @@ func (s *CompyTest) TestAuthentication(c *C) {
 	defer resp.Body.Close()
 	c.Assert(resp.StatusCode, Equals, 200)
 }
+
+func (s *CompyTest) TestAdmin(c *C) {
+	resp, err := s.client.Get("http://localhost" + *host)
+	c.Assert(err, IsNil)
+	defer resp.Body.Close()
+	c.Assert(resp.StatusCode, Equals, 200)
+
+	resp, err = s.client.Get("http://localhost" + *host + "/cacert")
+	c.Assert(err, IsNil)
+	defer resp.Body.Close()
+	c.Assert(resp.StatusCode, Equals, 404)
+
+	resp, err = s.client.Get("http://localhost" + *host + "/fake")
+	c.Assert(err, IsNil)
+	defer resp.Body.Close()
+	c.Assert(resp.StatusCode, Equals, 501)
+}
diff --git a/proxy/mitmlistener.go b/proxy/mitmlistener.go
index 519274c..3d394e8 100644
--- a/proxy/mitmlistener.go
+++ b/proxy/mitmlistener.go
@@ -6,14 +6,16 @@ import (
 )
 
 type mitmListener struct {
-	c  chan net.Conn
-	cf *certFaker
+	c      chan net.Conn
+	cf     *certFaker
+	config *tls.Config
 }
 
-func newMitmListener(cf *certFaker) *mitmListener {
+func newMitmListener(cf *certFaker, config *tls.Config) *mitmListener {
 	return &mitmListener{
-		c:  make(chan net.Conn),
-		cf: cf,
+		c:      make(chan net.Conn),
+		cf:     cf,
+		config: config,
 	}
 }
 
@@ -30,7 +32,7 @@ func (l *mitmListener) Addr() net.Addr {
 }
 
 func (l *mitmListener) Serve(conn net.Conn, host string) (net.Conn, error) {
-	sconn, err := tls.Dial("tcp", host, nil)
+	sconn, err := tls.Dial("tcp", host, l.config)
 	if err != nil {
 		return nil, err
 	}
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 52659cf..dda49ab 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -1,11 +1,17 @@
 package proxy
 
 import (
+	"crypto/tls"
+	"crypto/x509"
 	"encoding/base64"
+	"errors"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"log"
 	"net"
 	"net/http"
+	"os"
 	"strings"
 	"sync/atomic"
 )
@@ -17,16 +23,20 @@ type Proxy struct {
 	WriteCount  uint64
 	user        string
 	pass        string
+	host        string
+	cert        *string
 }
 
 type Transcoder interface {
 	Transcode(*ResponseWriter, *ResponseReader, http.Header) error
 }
 
-func New() *Proxy {
+func New(host string, cert *string) *Proxy {
 	p := &Proxy{
 		transcoders: make(map[string]Transcoder),
 		ml:          nil,
+		host:        host,
+		cert:        cert,
 	}
 	return p
 }
@@ -36,7 +46,25 @@ func (p *Proxy) EnableMitm(ca, key string) error {
 	if err != nil {
 		return err
 	}
-	p.ml = newMitmListener(cf)
+
+	var config *tls.Config
+	if p.cert != nil {
+		roots, err := x509.SystemCertPool()
+		if err != nil {
+			return err
+		}
+		pem, err := ioutil.ReadFile(*p.cert)
+		if err != nil {
+			return err
+		}
+		ok := roots.AppendCertsFromPEM([]byte(pem))
+		if !ok {
+			return errors.New("failed to parse root certificate")
+		}
+		config = &tls.Config{RootCAs: roots}
+	}
+
+	p.ml = newMitmListener(cf, config)
 	go http.Serve(p.ml, p)
 	return nil
 }
@@ -95,6 +123,14 @@ func (p *Proxy) handle(w http.ResponseWriter, r *http.Request) error {
 		return p.handleConnect(w, r)
 	}
 
+	host := r.URL.Host
+	if host == "" {
+		host = r.Host
+	}
+	if hostname, err := os.Hostname(); host == p.host || (err == nil && host == hostname+p.host) {
+		return p.handleLocalRequest(w, r)
+	}
+
 	resp, err := forward(r)
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
@@ -112,6 +148,39 @@ func (p *Proxy) handle(w http.ResponseWriter, r *http.Request) error {
 	return err
 }
 
+func (p *Proxy) handleLocalRequest(w http.ResponseWriter, r *http.Request) error {
+	if r.Method == "GET" && (r.URL.Path == "" || r.URL.Path == "/") {
+		w.Header().Set("Content-Type", "text/html")
+		read := atomic.LoadUint64(&p.ReadCount)
+		written := atomic.LoadUint64(&p.WriteCount)
+		io.WriteString(w, fmt.Sprintf(`<html>
+<head>
+<title>compy</title>
+</head>
+<body>
+<h1>compy</h1>
+<ul>
+<li>total transcoded: %d -> %d (%3.1f%%)</li>
+<li><a href="/cacert">CA cert</a></li>
+<li><a href="https://github.com/barnacs/compy">GitHub</a></li>
+</ul>
+</body>
+</html>`, read, written, float64(written)/float64(read)*100))
+		return nil
+	} else if r.Method == "GET" && r.URL.Path == "/cacert" {
+		if p.cert == nil {
+			http.NotFound(w, r)
+			return nil
+		}
+		w.Header().Set("Content-Type", "application/x-x509-ca-cert")
+		http.ServeFile(w, r, *p.cert)
+		return nil
+	} else {
+		w.WriteHeader(http.StatusNotImplemented)
+		return nil
+	}
+}
+
 func forward(r *http.Request) (*http.Response, error) {
 	if r.URL.Scheme == "" {
 		if r.TLS != nil && r.TLS.ServerName == r.Host {