newhttp.go 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. // Copyright 2017 fatedier, fatedier@gmail.com
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package vhost
  15. import (
  16. "bytes"
  17. "context"
  18. "errors"
  19. "log"
  20. "net"
  21. "net/http"
  22. "strings"
  23. "sync"
  24. "time"
  25. frpLog "github.com/fatedier/frp/utils/log"
  26. "github.com/fatedier/frp/utils/pool"
  27. )
  28. var (
  29. responseHeaderTimeout = time.Duration(30) * time.Second
  30. ErrRouterConfigConflict = errors.New("router config conflict")
  31. ErrNoDomain = errors.New("no such domain")
  32. )
  33. func getHostFromAddr(addr string) (host string) {
  34. strs := strings.Split(addr, ":")
  35. if len(strs) > 1 {
  36. host = strs[0]
  37. } else {
  38. host = addr
  39. }
  40. return
  41. }
  42. type HttpReverseProxy struct {
  43. proxy *ReverseProxy
  44. tr *http.Transport
  45. vhostRouter *VhostRouters
  46. cfgMu sync.RWMutex
  47. }
  48. func NewHttpReverseProxy() *HttpReverseProxy {
  49. rp := &HttpReverseProxy{
  50. vhostRouter: NewVhostRouters(),
  51. }
  52. proxy := &ReverseProxy{
  53. Director: func(req *http.Request) {
  54. req.URL.Scheme = "http"
  55. url := req.Context().Value("url").(string)
  56. host := getHostFromAddr(req.Context().Value("host").(string))
  57. host = rp.GetRealHost(host, url)
  58. if host != "" {
  59. req.Host = host
  60. }
  61. req.URL.Host = req.Host
  62. },
  63. Transport: &http.Transport{
  64. ResponseHeaderTimeout: responseHeaderTimeout,
  65. DisableKeepAlives: true,
  66. DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
  67. url := ctx.Value("url").(string)
  68. host := getHostFromAddr(ctx.Value("host").(string))
  69. return rp.CreateConnection(host, url)
  70. },
  71. },
  72. WebSocketDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
  73. url := ctx.Value("url").(string)
  74. host := getHostFromAddr(ctx.Value("host").(string))
  75. return rp.CreateConnection(host, url)
  76. },
  77. BufferPool: newWrapPool(),
  78. ErrorLog: log.New(newWrapLogger(), "", 0),
  79. }
  80. rp.proxy = proxy
  81. return rp
  82. }
  83. func (rp *HttpReverseProxy) Register(routeCfg VhostRouteConfig) error {
  84. rp.cfgMu.Lock()
  85. defer rp.cfgMu.Unlock()
  86. _, ok := rp.vhostRouter.Exist(routeCfg.Domain, routeCfg.Location)
  87. if ok {
  88. return ErrRouterConfigConflict
  89. } else {
  90. rp.vhostRouter.Add(routeCfg.Domain, routeCfg.Location, &routeCfg)
  91. }
  92. return nil
  93. }
  94. func (rp *HttpReverseProxy) UnRegister(domain string, location string) {
  95. rp.cfgMu.Lock()
  96. defer rp.cfgMu.Unlock()
  97. rp.vhostRouter.Del(domain, location)
  98. }
  99. func (rp *HttpReverseProxy) GetRealHost(domain string, location string) (host string) {
  100. vr, ok := rp.getVhost(domain, location)
  101. if ok {
  102. host = vr.payload.(*VhostRouteConfig).RewriteHost
  103. }
  104. return
  105. }
  106. func (rp *HttpReverseProxy) CreateConnection(domain string, location string) (net.Conn, error) {
  107. vr, ok := rp.getVhost(domain, location)
  108. if ok {
  109. fn := vr.payload.(*VhostRouteConfig).CreateConnFn
  110. if fn != nil {
  111. return fn()
  112. }
  113. }
  114. return nil, ErrNoDomain
  115. }
  116. func (rp *HttpReverseProxy) CheckAuth(domain, location, user, passwd string) bool {
  117. vr, ok := rp.getVhost(domain, location)
  118. if ok {
  119. checkUser := vr.payload.(*VhostRouteConfig).Username
  120. checkPasswd := vr.payload.(*VhostRouteConfig).Password
  121. if (checkUser != "" || checkPasswd != "") && (checkUser != user || checkPasswd != passwd) {
  122. return false
  123. }
  124. }
  125. return true
  126. }
  127. func (rp *HttpReverseProxy) getVhost(domain string, location string) (vr *VhostRouter, ok bool) {
  128. rp.cfgMu.RLock()
  129. defer rp.cfgMu.RUnlock()
  130. // first we check the full hostname
  131. // if not exist, then check the wildcard_domain such as *.example.com
  132. vr, ok = rp.vhostRouter.Get(domain, location)
  133. if ok {
  134. return
  135. }
  136. domainSplit := strings.Split(domain, ".")
  137. if len(domainSplit) < 3 {
  138. return vr, false
  139. }
  140. domainSplit[0] = "*"
  141. domain = strings.Join(domainSplit, ".")
  142. vr, ok = rp.vhostRouter.Get(domain, location)
  143. return
  144. }
  145. func (rp *HttpReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
  146. domain := getHostFromAddr(req.Host)
  147. location := req.URL.Path
  148. user, passwd, _ := req.BasicAuth()
  149. if !rp.CheckAuth(domain, location, user, passwd) {
  150. rw.Header().Set("WWW-Authenticate", `Basic realm="Restricted"`)
  151. http.Error(rw, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
  152. return
  153. }
  154. rp.proxy.ServeHTTP(rw, req)
  155. }
  156. type wrapPool struct{}
  157. func newWrapPool() *wrapPool { return &wrapPool{} }
  158. func (p *wrapPool) Get() []byte { return pool.GetBuf(32 * 1024) }
  159. func (p *wrapPool) Put(buf []byte) { pool.PutBuf(buf) }
  160. type wrapLogger struct{}
  161. func newWrapLogger() *wrapLogger { return &wrapLogger{} }
  162. func (l *wrapLogger) Write(p []byte) (n int, err error) {
  163. frpLog.Warn("%s", string(bytes.TrimRight(p, "\n")))
  164. return len(p), nil
  165. }