Changed: add possibility to override authorize function

README.md

@@ -144,6 +144,30 @@ proxy := gomitmproxy.NewProxy(gomitmproxy.Config{
 })
 ```
 
+Or implement using `OnAuthorize` handler.
+
+```go
+proxy := gomitmproxy.NewProxy(gomitmproxy.Config{
+    ListenAddr: &net.TCPAddr{
+        IP:   net.IPv4(0, 0, 0, 0),
+        Port: 8080,
+    },
+    OnAuthorize: func(session *gomitmproxy.Session) (bool, *http.Response) {
+		if session.Ctx().HasParent() {
+			// If we're here, it means the connection is authorized already.
+			return true, nil
+		}
+
+		username, password, ok := session.Request().BasicAuth()
+		if !ok {
+			return false, gomitmproxy.NewNotAuthorizedResponse(session)
+		}
+
+		return true, nil
+    },
+})
+```
+
 ### HTTP over TLS (HTTPS) proxy
 
 If you want to protect yourself from eavesdropping on your traffic to proxy, you can configure

auth.go

@@ -8,20 +8,20 @@ import (
 	"github.com/AdguardTeam/gomitmproxy/proxyutil"
 )
 
-// basicAuth returns an HTTP authorization header value according to RFC2617.
+// BasicAuth returns an HTTP authorization header value according to RFC2617.
 // See 2 (end of page 4) https://www.ietf.org/rfc/rfc2617.txt:
 // "To receive authorization, the client sends the userid and password,
 // separated by a single colon (":") character, within a base64 encoded string
 // in the credentials."
 // It is not meant to be urlencoded.
-func basicAuth(username, password string) string {
+func BasicAuth(username, password string) string {
 	auth := username + ":" + password
 	return base64.StdEncoding.EncodeToString([]byte(auth))
 }
 
-// newNotAuthorizedResponse creates a new "407 (Proxy Authentication Required)"
+// NewNotAuthorizedResponse creates a new "407 (Proxy Authentication Required)"
 // response.
-func newNotAuthorizedResponse(session *Session) *http.Response {
+func NewNotAuthorizedResponse(session *Session) *http.Response {
 	res := proxyutil.NewResponse(http.StatusProxyAuthRequired, nil, session.req)
 
 	// See https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Proxy-Authenticate.
@@ -34,7 +34,7 @@ func newNotAuthorizedResponse(session *Session) *http.Response {
 // request is authorized. If it returns false, it also returns the response that
 // should be written to the client.
 func (p *Proxy) authorize(session *Session) (bool, *http.Response) {
-	if session.ctx.parent != nil {
+	if session.Ctx().HasParent() {
 		// If we're here, it means the connection is authorized already.
 		return true, nil
 	}
@@ -44,14 +44,14 @@ func (p *Proxy) authorize(session *Session) (bool, *http.Response) {
 	}
 
 	// See https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Proxy-Authorization.
-	proxyAuth := session.req.Header.Get("Proxy-Authorization")
+	proxyAuth := session.Request().Header.Get("Proxy-Authorization")
 	if strings.Index(proxyAuth, "Basic ") != 0 {
-		return false, newNotAuthorizedResponse(session)
+		return false, NewNotAuthorizedResponse(session)
 	}
 
 	authHeader := proxyAuth[len("Basic "):]
-	if authHeader != basicAuth(p.Username, p.Password) {
-		return false, newNotAuthorizedResponse(session)
+	if authHeader != BasicAuth(p.Username, p.Password) {
+		return false, NewNotAuthorizedResponse(session)
 	}
 
 	return true, nil

config.go

@@ -20,6 +20,9 @@ type OnResponseFunc func(session *Session) (resp *http.Response)
 // OnErrorFunc is a declaration of the Config.OnError handler.
 type OnErrorFunc func(session *Session, err error)
 
+// OnAuthorizeFunc is a declaration of the Config.OnAuthorize handler.
+type OnAuthorizeFunc func(session *Session) (bool, *http.Response)
+
 // Config is the configuration of the Proxy.
 type Config struct {
 	// ListenAddr is the TCP address the proxy should listen to.
@@ -83,4 +86,7 @@ type Config struct {
 	// OnError is called if there's an issue with retrieving the response from
 	// the remote server.
 	OnError OnErrorFunc
+
+	// OnAuthorize is called for authorize request.
+	OnAuthorize OnAuthorizeFunc
 }

context.go

@@ -102,6 +102,14 @@ func (c *Context) SetProp(key string, val interface{}) {
 	c.props[key] = val
 }
 
+// HasParent check if parent is not nil.
+func (c *Context) HasParent() bool {
+	if c.parent != nil {
+		return true
+	}
+	return false
+}
+
 // Session contains all the necessary information about the request-response
 // pair that is currently being processed.
 type Session struct {

proxy.go

@@ -258,7 +258,23 @@ func (p *Proxy) handleRequest(ctx *Context) (err error) {
 
 	if !customRes {
 		// check proxy authorization first.
-		if p.Username != "" {
+		if p.OnAuthorize != nil {
+			auth, res := p.OnAuthorize(session)
+			if !auth {
+				log.Debug("id=%s: proxy auth required", session.ID())
+				session.res = res
+
+				defer log.OnCloserError(res.Body, log.DEBUG)
+
+				_ = p.writeResponse(session)
+
+				// Do not return any error here as we must keep the connection
+				// alive. When the client receives 407 error, it can write
+				// another request with user credentials to the same connection.
+				// See https://github.com/AdguardTeam/gomitmproxy/pull/19.
+				return nil
+			}
+		} else if p.Username != "" {
 			auth, res := p.authorize(session)
 			if !auth {
 				log.Debug("id=%s: proxy auth required", session.ID())