github.com/kaydxh/golang@v0.0.131/pkg/proxy/proxy.go (about)

     1  /*
     2   *Copyright (c) 2022, kaydxh
     3   *
     4   *Permission is hereby granted, free of charge, to any person obtaining a copy
     5   *of this software and associated documentation files (the "Software"), to deal
     6   *in the Software without restriction, including without limitation the rights
     7   *to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     8   *copies of the Software, and to permit persons to whom the Software is
     9   *furnished to do so, subject to the following conditions:
    10   *
    11   *The above copyright notice and this permission notice shall be included in all
    12   *copies or substantial portions of the Software.
    13   *
    14   *THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    15   *IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    16   *FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    17   *AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    18   *LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    19   *OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    20   *SOFTWARE.
    21   */
    22  package proxy
    23  
    24  import (
    25  	"fmt"
    26  	"net/http"
    27  	"net/http/httputil"
    28  
    29  	"github.com/gin-gonic/gin"
    30  	"github.com/gin-gonic/gin/render"
    31  	http_ "github.com/kaydxh/golang/go/net/http"
    32  	time_ "github.com/kaydxh/golang/go/time"
    33  	logs_ "github.com/kaydxh/golang/pkg/logs"
    34  )
    35  
    36  type ProxyMode int32
    37  
    38  const (
    39  	Reverse_ProxyMode  ProxyMode = 0
    40  	Forward_ProxyMode  ProxyMode = 1
    41  	Redirect_ProxyMode ProxyMode = 2
    42  )
    43  
    44  type ProxyMatchedFunc func(c *gin.Context) bool
    45  type ProxyTargetResolverFunc func(c *gin.Context) (host string, err error)
    46  
    47  type proxyOptions struct {
    48  	proxyMatchedFunc        ProxyMatchedFunc
    49  	proxyTargetResolverFunc ProxyTargetResolverFunc
    50  	proxyMode               ProxyMode
    51  }
    52  
    53  type Proxy struct {
    54  	router gin.IRouter
    55  	opts   proxyOptions
    56  }
    57  
    58  func defaultProxyOptions() proxyOptions {
    59  	return proxyOptions{
    60  		proxyMode: Reverse_ProxyMode,
    61  	}
    62  }
    63  
    64  func NewProxy(router gin.IRouter, options ...ProxyOption) (*Proxy, error) {
    65  	p := &Proxy{
    66  		router: router,
    67  	}
    68  	p.opts = defaultProxyOptions()
    69  	p.ApplyOptions(options...)
    70  
    71  	p.setProxy()
    72  	return p, nil
    73  }
    74  
    75  func (p *Proxy) ProxyHandler() gin.HandlerFunc {
    76  
    77  	return func(c *gin.Context) {
    78  		tc := time_.New(true)
    79  		logger := logs_.GetLogger(c.Request.Context())
    80  		summary := func() {
    81  			tc.Tick("proxy handler")
    82  			logger.Infof(tc.String())
    83  		}
    84  		defer summary()
    85  
    86  		if p.opts.proxyMatchedFunc != nil {
    87  			// not apply proxy, process inplace
    88  			if !p.opts.proxyMatchedFunc(c) {
    89  				return
    90  			}
    91  		}
    92  
    93  		if p.opts.proxyTargetResolverFunc == nil {
    94  			// proxy target resolver func is nil, process inplace
    95  			return
    96  		}
    97  		targetAddr, err := p.opts.proxyTargetResolverFunc(c)
    98  		if err != nil {
    99  			c.Render(http.StatusOK, render.JSON{Data: fmt.Errorf("resolve proxy target err: %v", err)})
   100  			return
   101  		}
   102  		if targetAddr == "" {
   103  			// proxy target is empty, process inplace
   104  			return
   105  		}
   106  
   107  		targetUrl := http_.CloneURL(c.Request.URL)
   108  		if targetUrl.Scheme == "" {
   109  			targetUrl.Scheme = "http"
   110  		}
   111  		targetUrl.Host = targetAddr
   112  		targetUrl.Path = "/"
   113  
   114  		switch p.opts.proxyMode {
   115  		case Redirect_ProxyMode:
   116  			c.Redirect(http.StatusTemporaryRedirect, targetUrl.String())
   117  			c.Abort()
   118  			return
   119  
   120  		case Reverse_ProxyMode:
   121  			c.Request.Host = targetUrl.Host
   122  
   123  		case Forward_ProxyMode:
   124  		}
   125  
   126  		rp := httputil.NewSingleHostReverseProxy(targetUrl)
   127  		rp.ServeHTTP(c.Writer, c.Request)
   128  		c.Abort()
   129  	}
   130  }
   131  
   132  func (p *Proxy) setProxy() {
   133  	p.router.Use(p.ProxyHandler())
   134  	/*
   135  		for _, pattern := range p.opts.routerPatterns {
   136  			p.router.Any(pattern, p.ProxyHandler())
   137  		}
   138  	*/
   139  }