From 69b89413bddfb4f4fbbc257871ed4067303fb8c8 Mon Sep 17 00:00:00 2001
From: Andrew Gaul <andrew@gaul.org>
Date: Wed, 5 Jul 2017 18:20:26 -0700
Subject: [PATCH] Add optional HTTP BASIC authorization

Fixes #18.
---
 README.md      |  6 ++++++
 compy.go       |  9 +++++++++
 compy_test.go  | 32 ++++++++++++++++++++++++++++++++
 proxy/proxy.go | 35 +++++++++++++++++++++++++++++++++++
 4 files changed, 82 insertions(+)

diff --git a/README.md b/README.md
index 87374a6..5c5b025 100644
--- a/README.md
+++ b/README.md
@@ -76,6 +76,12 @@ Probably the best option is to run it with both TLS and MitM support, combining
 compy -cert cert.crt -key cert.key -ca ca.crt -cakey ca.key
 ```
 
+You can limit access to your proxy via HTTP BASIC authentication:
+
+```
+compy -cert cert.crt -key cert.key -user myuser -pass mypass
+```
+
 You can also specify the listen port (defaults to 9999):  
 ```
 compy -host :9999
diff --git a/compy.go b/compy.go
index 1068285..aaf90e8 100644
--- a/compy.go
+++ b/compy.go
@@ -17,6 +17,8 @@ var (
 	key   = flag.String("key", "", "proxy cert key path")
 	ca    = flag.String("ca", "", "CA path")
 	caKey = flag.String("cakey", "", "CA key path")
+	user  = flag.String("user", "", "proxy user name")
+	pass  = flag.String("pass", "", "proxy password")
 
 	brotli = flag.Int("brotli", -1, "Brotli compression level (0-11, default 6)")
 	jpeg   = flag.Int("jpeg", 50, "jpeg quality (1-100, 0 to disable)")
@@ -45,6 +47,13 @@ func main() {
 		}
 	}
 
+	// TODO: require cert and key?
+	if (*user == "") != (*pass == "") {
+		log.Fatalln("must specify both user and pass")
+	} else {
+		p.SetAuthentication(*user, *pass)
+	}
+
 	if *jpeg != 0 {
 		p.AddTranscoder("image/jpeg", tc.NewJpeg(*jpeg))
 	}
diff --git a/compy_test.go b/compy_test.go
index 5165287..2179286 100644
--- a/compy_test.go
+++ b/compy_test.go
@@ -5,6 +5,7 @@ import (
 
 	"bytes"
 	gzipp "compress/gzip"
+	"encoding/base64"
 	gifp "image/gif"
 	jpegp "image/jpeg"
 	pngp "image/png"
@@ -183,3 +184,34 @@ func (s *CompyTest) TestPngToWebP(c *C) {
 	_, err = webp.Decode(resp.Body)
 	c.Assert(err, IsNil)
 }
+
+func (s *CompyTest) TestAuthentication(c *C) {
+	s.proxy.SetAuthentication("user", "pass")
+	defer s.proxy.SetAuthentication("", "")
+
+	// no password
+	resp, err := s.client.Get(s.server.URL + "/status/200")
+	c.Assert(err, IsNil)
+	defer resp.Body.Close()
+	c.Assert(resp.StatusCode, Equals, 407)
+
+	// incorrect password
+	req, err := http.NewRequest("GET", s.server.URL+"/status/200", nil)
+	c.Assert(err, IsNil)
+	req.Header.Add("Proxy-Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte("wrong:bad")))
+
+	resp, err = s.client.Do(req)
+	c.Assert(err, IsNil)
+	defer resp.Body.Close()
+	c.Assert(resp.StatusCode, Equals, 407)
+
+	// correct password
+	req, err = http.NewRequest("GET", s.server.URL+"/status/200", nil)
+	c.Assert(err, IsNil)
+	req.Header.Add("Proxy-Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte("user:pass")))
+
+	resp, err = s.client.Do(req)
+	c.Assert(err, IsNil)
+	defer resp.Body.Close()
+	c.Assert(resp.StatusCode, Equals, 200)
+}
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 6912460..6ec19d1 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -1,10 +1,12 @@
 package proxy
 
 import (
+	"encoding/base64"
 	"fmt"
 	"log"
 	"net"
 	"net/http"
+	"strings"
 	"sync/atomic"
 )
 
@@ -13,6 +15,8 @@ type Proxy struct {
 	ml          *mitmListener
 	ReadCount   uint64
 	WriteCount  uint64
+	user        string
+	pass        string
 }
 
 type Transcoder interface {
@@ -37,6 +41,11 @@ func (p *Proxy) EnableMitm(ca, key string) error {
 	return nil
 }
 
+func (p *Proxy) SetAuthentication(user, pass string) {
+	p.user = user
+	p.pass = pass
+}
+
 func (p *Proxy) AddTranscoder(contentType string, transcoder Transcoder) {
 	p.transcoders[contentType] = transcoder
 }
@@ -56,10 +65,36 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func (p *Proxy) checkHttpBasicAuth(auth string) bool {
+	prefix := "Basic "
+	if !strings.HasPrefix(auth, prefix) {
+		return false
+	}
+	decoded, err := base64.StdEncoding.DecodeString(auth[len(prefix):])
+	if err != nil {
+		return false
+	}
+	values := strings.SplitN(string(decoded), ":", 2)
+	if len(values) != 2 || values[0] != p.user || values[1] != p.pass {
+		return false
+	}
+	return true
+}
+
 func (p *Proxy) handle(w http.ResponseWriter, r *http.Request) error {
 	if r.Method == "CONNECT" {
 		return p.handleConnect(w, r)
 	}
+
+	// TODO: only HTTPS?
+	if p.user != "" {
+		if !p.checkHttpBasicAuth(r.Header.Get("Proxy-Authorization")) {
+			w.Header().Set("WWW-Authenticate", "Basic realm=\"Compy\"")
+			w.WriteHeader(http.StatusProxyAuthRequired)
+			return nil
+		}
+	}
+
 	resp, err := forward(r)
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)