cyonjan/frp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Add success function to vhost

Browse Source
This commit is contained in:
Guy Lewin 2020-03-03 17:42:42 -05:00
parent 6fc7135cbd
commit a537bbd585
3 changed files with 16 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
utils/tcpmux/httpconnect.go
View File

@ -31,7 +31,7 @@ type HttpConnectTcpMuxer struct {
}
func NewHttpConnectTcpMuxer(listener net.Listener, timeout time.Duration) (*HttpConnectTcpMuxer, error) {
mux, err := vhost.NewVhostMuxer(listener, getHostFromHttpConnect, nil, sendHttpOk, timeout)
mux, err := vhost.NewVhostMuxer(listener, getHostFromHttpConnect, nil, sendHttpOk, nil, timeout)
return &HttpConnectTcpMuxer{mux}, err
}
@ -52,10 +52,8 @@ func readHttpConnectRequest(rd io.Reader) (host string, err error) {
return
}
func sendHttpOk(c net.Conn, _ string) (_ net.Conn, err error) {
okResp := util.OkResponse()
err = okResp.Write(c)
return c, err
func sendHttpOk(c net.Conn) error {
return util.OkResponse().Write(c)
}
func getHostFromHttpConnect(c net.Conn) (_ net.Conn, _ map[string]string, err error) {

2
utils/vhost/https.go
View File

@ -48,7 +48,7 @@ type HttpsMuxer struct {
}
func NewHttpsMuxer(listener net.Listener, timeout time.Duration) (*HttpsMuxer, error) {
mux, err := NewVhostMuxer(listener, GetHttpsHostname, nil, nil, timeout)
mux, err := NewVhostMuxer(listener, GetHttpsHostname, nil, nil, nil, timeout)
return &HttpsMuxer{mux}, err
}

13
utils/vhost/vhost.go
View File

@ -29,22 +29,25 @@ import (
type muxFunc func(net.Conn) (net.Conn, map[string]string, error)
type httpAuthFunc func(net.Conn, string, string, string) (bool, error)
type hostRewriteFunc func(net.Conn, string) (net.Conn, error)
type successFunc func(net.Conn) error
type VhostMuxer struct {
listener net.Listener
timeout time.Duration
vhostFunc muxFunc
authFunc httpAuthFunc
successFunc successFunc
rewriteFunc hostRewriteFunc
registryRouter *VhostRouters
}
func NewVhostMuxer(listener net.Listener, vhostFunc muxFunc, authFunc httpAuthFunc, rewriteFunc hostRewriteFunc, timeout time.Duration) (mux *VhostMuxer, err error) {
func NewVhostMuxer(listener net.Listener, vhostFunc muxFunc, authFunc httpAuthFunc, successFunc successFunc, rewriteFunc hostRewriteFunc, timeout time.Duration) (mux *VhostMuxer, err error) {
mux = &VhostMuxer{
listener: listener,
timeout: timeout,
vhostFunc: vhostFunc,
authFunc: authFunc,
successFunc: successFunc,
rewriteFunc: rewriteFunc,
registryRouter: NewVhostRouters(),
}
@ -149,7 +152,15 @@ func (v *VhostMuxer) handle(c net.Conn) {
c.Close()
return
}
xl := xlog.FromContextSafe(l.ctx)
if v.successFunc != nil {
if err := v.successFunc(c); err != nil {
xl.Info("success func failure on vhost connection: %v", err)
c.Close()
return
}
}
// if authFunc is exist and userName/password is set
// then verify user access