utils/vhost: supprot http authentication

src/utils/vhost/http.go

@@ -17,6 +17,7 @@ package vhost
 import (
 	"bufio"
 	"bytes"
+	"encoding/base64"
 	"fmt"
 	"io"
 	"net"
@@ -33,21 +34,29 @@ type HttpMuxer struct {
 	*VhostMuxer
 }
 
-func GetHttpHostname(c *conn.Conn) (_ net.Conn, routerName string, err error) {
+func GetHttpRequestInfo(c *conn.Conn) (_ net.Conn, _ map[string]string, err error) {
+	reqInfoMap := make(map[string]string, 0)
 	sc, rd := newShareConn(c.TcpConn)
 
 	request, err := http.ReadRequest(bufio.NewReader(rd))
 	if err != nil {
-		return sc, "", err
+		return sc, reqInfoMap, err
 	}
+	// hostName
 	tmpArr := strings.Split(request.Host, ":")
-	routerName = tmpArr[0]
+	reqInfoMap["Host"] = tmpArr[0]
+
+	// Authorization
+	authStr := request.Header.Get("Authorization")
+	if authStr != "" {
+		reqInfoMap["Authorization"] = authStr
+	}
 	request.Body.Close()
-	return sc, routerName, nil
+	return sc, reqInfoMap, nil
 }
 
 func NewHttpMuxer(listener *conn.Listener, timeout time.Duration) (*HttpMuxer, error) {
-	mux, err := NewVhostMuxer(listener, GetHttpHostname, HttpHostNameRewrite, timeout)
+	mux, err := NewVhostMuxer(listener, GetHttpRequestInfo, HttpAuthFunc, HttpHostNameRewrite, timeout)
 	return &HttpMuxer{mux}, err
 }
 
@@ -169,3 +178,38 @@ func changeHostName(buff *bytes.Buffer, rewriteHost string) (_ []byte, err error
 	retBuf.Write(peek)
 	return retBuf.Bytes(), err
 }
+
+func HttpAuthFunc(c *conn.Conn, userName, passWord, authorization string) (bAccess bool, err error) {
+	s := strings.SplitN(authorization, " ", 2)
+	if len(s) != 2 {
+		res := noAuthResponse()
+		res.Write(c.TcpConn)
+		return
+	}
+	b, err := base64.StdEncoding.DecodeString(s[1])
+	if err != nil {
+		return
+	}
+	pair := strings.SplitN(string(b), ":", 2)
+	if len(pair) != 2 {
+		return
+	}
+	if pair[0] != userName || pair[1] != passWord {
+		return
+	}
+	return true, nil
+}
+
+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
+}

src/utils/vhost/https.go

@@ -48,7 +48,7 @@ type HttpsMuxer struct {
 }
 
 func NewHttpsMuxer(listener *conn.Listener, timeout time.Duration) (*HttpsMuxer, error) {
-	mux, err := NewVhostMuxer(listener, GetHttpsHostname, nil, timeout)
+	mux, err := NewVhostMuxer(listener, GetHttpsHostname, nil, nil, timeout)
 	return &HttpsMuxer{mux}, err
 }
 
@@ -178,11 +178,13 @@ func readHandshake(rd io.Reader) (host string, err error) {
 	return
 }
 
-func GetHttpsHostname(c *conn.Conn) (sc net.Conn, routerName string, err error) {
+func GetHttpsHostname(c *conn.Conn) (sc net.Conn, _ map[string]string, err error) {
+	reqInfoMap := make(map[string]string, 0)
 	sc, rd := newShareConn(c.TcpConn)
 	host, err := readHandshake(rd)
 	if err != nil {
-		return sc, "", err
+		return sc, reqInfoMap, err
 	}
-	return sc, host, nil
+	reqInfoMap["Host"] = host
+	return sc, reqInfoMap, nil
 }

src/utils/vhost/vhost.go

@@ -1,5 +1,3 @@
-// Copyright 2016 fatedier, fatedier@gmail.com
-//
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
@@ -26,23 +24,26 @@ import (
 	"github.com/fatedier/frp/src/utils/conn"
 )
 
-type muxFunc func(*conn.Conn) (net.Conn, string, error)
+type muxFunc func(*conn.Conn) (net.Conn, map[string]string, error)
+type httpAuthFunc func(*conn.Conn, string, string, string) (bool, error)
 type hostRewriteFunc func(*conn.Conn, string) (net.Conn, error)
 
 type VhostMuxer struct {
 	listener    *conn.Listener
 	timeout     time.Duration
 	vhostFunc   muxFunc
+	authFunc    httpAuthFunc
 	rewriteFunc hostRewriteFunc
 	registryMap map[string]*Listener
 	mutex       sync.RWMutex
 }
 
-func NewVhostMuxer(listener *conn.Listener, vhostFunc muxFunc, rewriteFunc hostRewriteFunc, timeout time.Duration) (mux *VhostMuxer, err error) {
+func NewVhostMuxer(listener *conn.Listener, vhostFunc muxFunc, authFunc httpAuthFunc, rewriteFunc hostRewriteFunc, timeout time.Duration) (mux *VhostMuxer, err error) {
 	mux = &VhostMuxer{
 		listener:    listener,
 		timeout:     timeout,
 		vhostFunc:   vhostFunc,
+		authFunc:    authFunc,
 		rewriteFunc: rewriteFunc,
 		registryMap: make(map[string]*Listener),
 	}
@@ -51,7 +52,7 @@ func NewVhostMuxer(listener *conn.Listener, vhostFunc muxFunc, rewriteFunc hostR
 }
 
 // listen for a new domain name, if rewriteHost is not empty  and rewriteFunc is not nil, then rewrite the host header to rewriteHost
-func (v *VhostMuxer) Listen(name string, rewriteHost string) (l *Listener, err error) {
+func (v *VhostMuxer) Listen(name string, rewriteHost, userName, passWord string) (l *Listener, err error) {
 	v.mutex.Lock()
 	defer v.mutex.Unlock()
 	if _, exist := v.registryMap[name]; exist {
@@ -61,6 +62,8 @@ func (v *VhostMuxer) Listen(name string, rewriteHost string) (l *Listener, err e
 	l = &Listener{
 		name:        name,
 		rewriteHost: rewriteHost,
+		userName:    userName,
+		passWord:    passWord,
 		mux:         v,
 		accept:      make(chan *conn.Conn),
 	}
@@ -109,13 +112,13 @@ func (v *VhostMuxer) handle(c *conn.Conn) {
 		return
 	}
 
-	sConn, name, err := v.vhostFunc(c)
+	sConn, reqInfoMap, err := v.vhostFunc(c)
 	if err != nil {
 		c.Close()
 		return
 	}
 
-	name = strings.ToLower(name)
+	name := strings.ToLower(reqInfoMap["Host"])
 	// get listener by hostname
 	l, ok := v.getListener(name)
 	if !ok {
@@ -123,6 +126,20 @@ func (v *VhostMuxer) handle(c *conn.Conn) {
 		return
 	}
 
+	// if authFunc is exist and userName/password is set
+	// verify user access
+	fmt.Printf("reqInfo: %+v\n", reqInfoMap)
+	if l.mux.authFunc != nil &&
+		l.userName != "" && l.passWord != "" {
+		bAccess, err := l.mux.authFunc(c, l.userName, l.passWord, reqInfoMap["Authorization"])
+		if bAccess == false || err != nil {
+			res := noAuthResponse()
+			res.Write(c.TcpConn)
+			c.Close()
+			return
+		}
+	}
+
 	if err = sConn.SetDeadline(time.Time{}); err != nil {
 		c.Close()
 		return
@@ -135,6 +152,8 @@ func (v *VhostMuxer) handle(c *conn.Conn) {
 type Listener struct {
 	name        string
 	rewriteHost string
+	userName    string
+	passWord    string
 	mux         *VhostMuxer // for closing VhostMuxer
 	accept      chan *conn.Conn
 }
@@ -154,6 +173,7 @@ func (l *Listener) Accept() (*conn.Conn, error) {
 		}
 		conn.SetTcpConn(sConn)
 	}
+
 	return conn, nil
 }