frp/src/utils/vhost/vhost.go

228 lines
5.0 KiB
Go
Raw Normal View History

2016-04-18 15:16:40 +08:00
// 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
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package vhost
import (
"bytes"
"fmt"
"io"
"net"
"strings"
"sync"
"time"
"github.com/fatedier/frp/src/utils/conn"
2016-04-18 15:16:40 +08:00
)
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)
2016-04-18 15:16:40 +08:00
type VhostMuxer struct {
listener *conn.Listener
timeout time.Duration
vhostFunc muxFunc
authFunc httpAuthFunc
rewriteFunc hostRewriteFunc
registryRouter *VhostRouters
mutex sync.RWMutex
2016-04-18 15:16:40 +08:00
}
func NewVhostMuxer(listener *conn.Listener, vhostFunc muxFunc, authFunc httpAuthFunc, rewriteFunc hostRewriteFunc, timeout time.Duration) (mux *VhostMuxer, err error) {
2016-04-18 15:16:40 +08:00
mux = &VhostMuxer{
listener: listener,
timeout: timeout,
vhostFunc: vhostFunc,
authFunc: authFunc,
rewriteFunc: rewriteFunc,
registryRouter: NewVhostRouters(),
2016-04-18 15:16:40 +08:00
}
go mux.run()
return mux, nil
}
// listen for a new domain name, if rewriteHost is not empty and rewriteFunc is not nil, then rewrite the host header to rewriteHost
2016-12-25 01:53:23 +08:00
func (v *VhostMuxer) Listen(name, location, rewriteHost, userName, passWord string) (l *Listener, err error) {
2016-04-18 15:16:40 +08:00
v.mutex.Lock()
defer v.mutex.Unlock()
2016-12-25 01:53:23 +08:00
l = &Listener{
name: name,
rewriteHost: rewriteHost,
userName: userName,
passWord: passWord,
mux: v,
accept: make(chan *conn.Conn),
2016-04-18 15:16:40 +08:00
}
2016-12-25 01:53:23 +08:00
v.registryRouter.Add(name, location, l)
return l, nil
2016-04-18 15:16:40 +08:00
}
2016-12-25 01:53:23 +08:00
func (v *VhostMuxer) getListener(name, path string) (l *Listener, exist bool) {
2016-04-18 15:16:40 +08:00
v.mutex.RLock()
defer v.mutex.RUnlock()
2016-12-25 01:53:23 +08:00
// first we check the full hostname
// if not exist, then check the wildcard_domain such as *.example.com
vr, found := v.registryRouter.Get(name, path)
if found {
return vr.listener, true
2016-08-23 00:58:51 +08:00
}
2016-08-23 00:58:51 +08:00
domainSplit := strings.Split(name, ".")
if len(domainSplit) < 3 {
return l, false
}
domainSplit[0] = "*"
name = strings.Join(domainSplit, ".")
2016-04-18 15:16:40 +08:00
2016-12-25 01:53:23 +08:00
vr, found = v.registryRouter.Get(name, path)
if !found {
return
}
return vr.listener, true
2016-04-18 15:16:40 +08:00
}
func (v *VhostMuxer) run() {
for {
conn, err := v.listener.Accept()
if err != nil {
return
}
go v.handle(conn)
}
}
func (v *VhostMuxer) handle(c *conn.Conn) {
if err := c.SetDeadline(time.Now().Add(v.timeout)); err != nil {
2016-08-23 00:58:51 +08:00
c.Close()
2016-04-18 15:16:40 +08:00
return
}
sConn, reqInfoMap, err := v.vhostFunc(c)
2016-04-18 15:16:40 +08:00
if err != nil {
2016-08-23 00:58:51 +08:00
c.Close()
2016-04-18 15:16:40 +08:00
return
}
2016-12-25 01:53:23 +08:00
name := strings.ToLower(reqInfoMap["Host"])
path := strings.ToLower(reqInfoMap["Path"])
l, ok := v.getListener(name, path)
2016-04-18 15:16:40 +08:00
if !ok {
2016-08-23 00:58:51 +08:00
c.Close()
2016-04-18 15:16:40 +08:00
return
}
// if authFunc is exist and userName/password is set
// 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 == false || err != nil {
res := noAuthResponse()
res.Write(c.TcpConn)
c.Close()
return
}
}
2016-04-18 15:16:40 +08:00
if err = sConn.SetDeadline(time.Time{}); err != nil {
2016-08-23 00:58:51 +08:00
c.Close()
2016-04-18 15:16:40 +08:00
return
}
c.SetTcpConn(sConn)
2016-04-18 15:16:40 +08:00
l.accept <- c
}
type Listener struct {
name string
rewriteHost string
userName string
passWord string
mux *VhostMuxer // for closing VhostMuxer
accept chan *conn.Conn
2016-04-18 15:16:40 +08:00
}
func (l *Listener) Accept() (*conn.Conn, error) {
conn, ok := <-l.accept
if !ok {
return nil, fmt.Errorf("Listener closed")
}
// if rewriteFunc is exist and rewriteHost is set
// rewrite http requests with a modified host header
if l.mux.rewriteFunc != nil && l.rewriteHost != "" {
sConn, err := l.mux.rewriteFunc(conn, l.rewriteHost)
if err != nil {
return nil, fmt.Errorf("http host header rewrite failed")
}
conn.SetTcpConn(sConn)
}
2016-04-18 15:16:40 +08:00
return conn, nil
}
func (l *Listener) Close() error {
2016-12-25 01:53:23 +08:00
l.mux.registryRouter.Del(l)
2016-04-18 15:16:40 +08:00
close(l.accept)
return nil
}
func (l *Listener) Name() string {
return l.name
}
type sharedConn struct {
net.Conn
sync.Mutex
buff *bytes.Buffer
}
// the bytes you read in io.Reader, will be reserved in sharedConn
2016-04-18 15:16:40 +08:00
func newShareConn(conn net.Conn) (*sharedConn, io.Reader) {
sc := &sharedConn{
Conn: conn,
buff: bytes.NewBuffer(make([]byte, 0, 1024)),
}
return sc, io.TeeReader(conn, sc.buff)
}
func (sc *sharedConn) Read(p []byte) (n int, err error) {
sc.Lock()
if sc.buff == nil {
sc.Unlock()
return sc.Conn.Read(p)
}
sc.Unlock()
2016-04-18 15:16:40 +08:00
n, err = sc.buff.Read(p)
if err == io.EOF {
sc.Lock()
2016-04-18 15:16:40 +08:00
sc.buff = nil
sc.Unlock()
2016-04-18 15:16:40 +08:00
var n2 int
n2, err = sc.Conn.Read(p[n:])
n += n2
}
return
}
func (sc *sharedConn) WriteBuff(buffer []byte) (err error) {
sc.buff.Reset()
_, err = sc.buff.Write(buffer)
return err
}