tcpmux: support authentication (#3345)

pkg/util/vhost/https.go

@@ -28,7 +28,10 @@ type HTTPSMuxer struct {
 }
 
 func NewHTTPSMuxer(listener net.Listener, timeout time.Duration) (*HTTPSMuxer, error) {
-	mux, err := NewMuxer(listener, GetHTTPSHostname, nil, nil, nil, timeout)
+	mux, err := NewMuxer(listener, GetHTTPSHostname, timeout)
+	if err != nil {
+		return nil, err
+	}
 	return &HTTPSMuxer{mux}, err
 }
 

pkg/util/vhost/resource.go

@@ -85,17 +85,3 @@ func notFoundResponse() *http.Response {
 	}
 	return res
 }
-
-func noAuthResponse() *http.Response {
-	header := make(map[string][]string)
-	header["WWW-Authenticate"] = []string{`Basic realm="Restricted"`}
-	res := &http.Response{
-		Status:     "401 Not authorized",
-		StatusCode: 401,
-		Proto:      "HTTP/1.1",
-		ProtoMajor: 1,
-		ProtoMinor: 1,
-		Header:     header,
-	}
-	return res
-}

pkg/util/vhost/vhost.go

@@ -43,43 +43,55 @@ type RequestRouteInfo struct {
 
 type (
 	muxFunc         func(net.Conn) (net.Conn, map[string]string, error)
-	httpAuthFunc    func(net.Conn, string, string, string) (bool, error)
+	authFunc        func(conn net.Conn, username, password string, reqInfoMap map[string]string) (bool, error)
 	hostRewriteFunc func(net.Conn, string) (net.Conn, error)
-	successFunc     func(net.Conn, map[string]string) error
+	successHookFunc func(net.Conn, map[string]string) error
 )
 
-// Muxer is only used for https and tcpmux proxy.
+// Muxer is a functional component used for https and tcpmux proxies.
+// It accepts connections and extracts vhost information from the beginning of the connection data.
+// It then routes the connection to its appropriate listener.
 type Muxer struct {
-	listener       net.Listener
-	timeout        time.Duration
+	listener net.Listener
+	timeout  time.Duration
+
 	vhostFunc      muxFunc
-	authFunc       httpAuthFunc
-	successFunc    successFunc
-	rewriteFunc    hostRewriteFunc
+	checkAuth      authFunc
+	successHook    successHookFunc
+	rewriteHost    hostRewriteFunc
 	registryRouter *Routers
 }
 
 func NewMuxer(
 	listener net.Listener,
 	vhostFunc muxFunc,
-	authFunc httpAuthFunc,
-	successFunc successFunc,
-	rewriteFunc hostRewriteFunc,
 	timeout time.Duration,
 ) (mux *Muxer, err error) {
 	mux = &Muxer{
 		listener:       listener,
 		timeout:        timeout,
 		vhostFunc:      vhostFunc,
-		authFunc:       authFunc,
-		successFunc:    successFunc,
-		rewriteFunc:    rewriteFunc,
 		registryRouter: NewRouters(),
 	}
 	go mux.run()
 	return mux, nil
 }
 
+func (v *Muxer) SetCheckAuthFunc(f authFunc) *Muxer {
+	v.checkAuth = f
+	return v
+}
+
+func (v *Muxer) SetSuccessHookFunc(f successHookFunc) *Muxer {
+	v.successHook = f
+	return v
+}
+
+func (v *Muxer) SetRewriteHostFunc(f hostRewriteFunc) *Muxer {
+	v.rewriteHost = f
+	return v
+}
+
 type ChooseEndpointFunc func() (string, error)
 
 type CreateConnFunc func(remoteAddr string) (net.Conn, error)
@@ -101,7 +113,7 @@ type RouteConfig struct {
 	CreateConnByEndpointFn CreateConnByEndpointFunc
 }
 
-// listen for a new domain name, if rewriteHost is not empty  and rewriteFunc is not nil
+// listen for a new domain name, if rewriteHost is not empty and rewriteHost func is not nil,
 // then rewrite the host header to rewriteHost
 func (v *Muxer) Listen(ctx context.Context, cfg *RouteConfig) (l *Listener, err error) {
 	l = &Listener{
@@ -109,8 +121,8 @@ func (v *Muxer) Listen(ctx context.Context, cfg *RouteConfig) (l *Listener, err
 		location:        cfg.Location,
 		routeByHTTPUser: cfg.RouteByHTTPUser,
 		rewriteHost:     cfg.RewriteHost,
-		userName:        cfg.Username,
-		passWord:        cfg.Password,
+		username:        cfg.Username,
+		password:        cfg.Password,
 		mux:             v,
 		accept:          make(chan net.Conn),
 		ctx:             ctx,
@@ -205,25 +217,20 @@ func (v *Muxer) handle(c net.Conn) {
 	}
 
 	xl := xlog.FromContextSafe(l.ctx)
-	if v.successFunc != nil {
-		if err := v.successFunc(c, reqInfoMap); err != nil {
+	if v.successHook != nil {
+		if err := v.successHook(c, reqInfoMap); err != nil {
 			xl.Info("success func failure on vhost connection: %v", err)
 			_ = c.Close()
 			return
 		}
 	}
 
-	// if authFunc is exist and username/password is set
+	// if checkAuth func is exist and username/password is set
 	// then verify user access
-	if l.mux.authFunc != nil && l.userName != "" && l.passWord != "" {
-		bAccess, err := l.mux.authFunc(c, l.userName, l.passWord, reqInfoMap["Authorization"])
-		if !bAccess || err != nil {
-			xl.Debug("check http Authorization failed")
-			res := noAuthResponse()
-			if res.Body != nil {
-				defer res.Body.Close()
-			}
-			_ = res.Write(c)
+	if l.mux.checkAuth != nil && l.username != "" {
+		ok, err := l.mux.checkAuth(c, l.username, l.password, reqInfoMap)
+		if !ok || err != nil {
+			xl.Debug("auth failed for user: %s", l.username)
 			_ = c.Close()
 			return
 		}
@@ -249,8 +256,8 @@ type Listener struct {
 	location        string
 	routeByHTTPUser string
 	rewriteHost     string
-	userName        string
-	passWord        string
+	username        string
+	password        string
 	mux             *Muxer // for closing Muxer
 	accept          chan net.Conn
 	ctx             context.Context
@@ -263,11 +270,11 @@ func (l *Listener) Accept() (net.Conn, error) {
 		return nil, fmt.Errorf("Listener closed")
 	}
 
-	// if rewriteFunc is exist
+	// if rewriteHost func is exist
 	// rewrite http requests with a modified host header
 	// if l.rewriteHost is empty, nothing to do
-	if l.mux.rewriteFunc != nil {
-		sConn, err := l.mux.rewriteFunc(conn, l.rewriteHost)
+	if l.mux.rewriteHost != nil {
+		sConn, err := l.mux.rewriteHost(conn, l.rewriteHost)
 		if err != nil {
 			xl.Warn("host header rewrite failed: %v", err)
 			return nil, fmt.Errorf("host header rewrite failed")