supprt proxy auth

isayme · May 22, 2024 · 7b9453a · 7b9453a
1 parent 5c7e5d2
commit 7b9453a
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 6 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -16,6 +16,8 @@ var showVersion bool
 var logFormat string
 var logLevel string
 var listenPort uint16
+var username string
+var password string
 var certFile string
 var keyFile string
 var proxyAddress string
@@ -32,6 +34,8 @@ func init() {
 	rootCmd.Flags().StringVarP(&logFormat, "log-format", "", "console", "log format")
 	rootCmd.Flags().StringVarP(&logLevel, "log-level", "", "info", "log level")
 	rootCmd.Flags().Uint16VarP(&listenPort, "port", "p", 1087, "listen port")
+	rootCmd.Flags().StringVarP(&username, "username", "", "", "proxy server auth username")
+	rootCmd.Flags().StringVarP(&password, "password", "", "", "proxy server auth password")
 	rootCmd.Flags().StringVarP(&certFile, "cert-file", "", "", "cert file")
 	rootCmd.Flags().StringVarP(&keyFile, "key-file", "", "", "key file")
 	rootCmd.Flags().StringVar(&proxyAddress, "proxy", "", "use proxy, format: 'socks5://host:port' or 'http://host:port' or 'https://host:port'")
@@ -57,18 +61,27 @@ var rootCmd = &cobra.Command{
 		address := fmt.Sprintf(":%d", listenPort)
 
 		options := []httpproxy.ServerOption{}
-		if proxyAddress != "" {
-			options = append(options, httpproxy.WithProxy(proxyAddress))
-			logger.Debugw("option", "proxy", proxyAddress)
+
+		if username != "" && password != "" {
+			maskPassword := password
+			if len(maskPassword) > 1 {
+				maskPassword = password[:1] + "***" + password[len(password)-1:]
+			}
+			logger.Debugw("option", "username", username, "password", maskPassword)
+			options = append(options, httpproxy.WithUsername(username))
+			options = append(options, httpproxy.WithPassword(password))
 		}
 		if connectTimeout > 0 {
-			options = append(options, httpproxy.WithConnectTimeout(connectTimeout))
 			logger.Debugw("option", "connect-timeout", connectTimeout.String())
-
+			options = append(options, httpproxy.WithConnectTimeout(connectTimeout))
 		}
 		if timeout > 0 {
-			options = append(options, httpproxy.WithTimeout(timeout))
 			logger.Debugw("option", "timeout", timeout.String())
+			options = append(options, httpproxy.WithTimeout(timeout))
+		}
+		if proxyAddress != "" {
+			logger.Debugw("option", "proxy", proxyAddress)
+			options = append(options, httpproxy.WithProxy(proxyAddress))
 		}
 
 		server, err := httpproxy.NewServer(address, options...)

diff --git a/httpproxy/option.go b/httpproxy/option.go
@@ -3,6 +3,9 @@ package httpproxy
 import "time"
 
 type serverOptions struct {
+	username string
+	password string
+
 	proxy          string
 	connectTimeout time.Duration
 	timeout        time.Duration
@@ -26,6 +29,18 @@ func newFuncServerOption(f func(*serverOptions)) *funcServerOption {
 	}
 }
 
+func WithUsername(username string) ServerOption {
+	return newFuncServerOption(func(o *serverOptions) {
+		o.username = username
+	})
+}
+
+func WithPassword(password string) ServerOption {
+	return newFuncServerOption(func(o *serverOptions) {
+		o.password = password
+	})
+}
+
 func WithProxy(addr string) ServerOption {
 	return newFuncServerOption(func(o *serverOptions) {
 		o.proxy = addr

diff --git a/httpproxy/server.go b/httpproxy/server.go
@@ -4,11 +4,13 @@ import (
 	"bufio"
 	"context"
 	"crypto/tls"
+	"encoding/base64"
 	"fmt"
 	"io"
 	"net"
 	"net/http"
 	"net/url"
+	"strings"
 	"sync"
 	"time"
 
@@ -141,6 +143,19 @@ func (s *Server) handleConnection(conn net.Conn) {
 		return
 	}
 
+	// auth
+	if s.options.username != "" && s.options.password != "" {
+		authorization := req.Header.Get("Proxy-Authorization")
+		username, password, ok := parseBasicAuth(authorization)
+		if !ok || username != s.options.username || password != s.options.password {
+			conn.Write([]byte("HTTP/1.1 407 Proxy Authentication Required\r\n"))
+			conn.Write([]byte("Content-Type: text/plain\r\n"))
+			conn.Write([]byte(fmt.Sprintf("Server: %s\r\n\r\n", Name)))
+			conn.Write([]byte(fmt.Sprintf("%s %s\r\n\r\n", Name, Version)))
+			return
+		}
+	}
+
 	logger.Infow("newRequest", "url", req.URL.String())
 	start := time.Now()
 	defer func() {
@@ -202,3 +217,22 @@ func (s *Server) handleConnection(conn net.Conn) {
 
 	wg.Wait()
 }
+
+// from package http
+func parseBasicAuth(auth string) (username, password string, ok bool) {
+	const prefix = "Basic "
+	// Case insensitive prefix match. See Issue 22736.
+	if len(auth) < len(prefix) || !strings.EqualFold(auth[:len(prefix)], prefix) {
+		return "", "", false
+	}
+	c, err := base64.StdEncoding.DecodeString(auth[len(prefix):])
+	if err != nil {
+		return "", "", false
+	}
+	cs := string(c)
+	username, password, ok = strings.Cut(cs, ":")
+	if !ok {
+		return "", "", false
+	}
+	return username, password, true
+}