revsuit/pkg/rhttp/http.go

package rhttp

import (
	"math/rand"
	"net/http"
	"net/http/httputil"
	"regexp"
	"strconv"
	"strings"
	"sync"

	"github.com/gin-gonic/gin"
	"github.com/li4n0/revsuit/internal/database"
	"github.com/li4n0/revsuit/internal/qqwry"
	log "unknwon.dev/clog/v2"
)

type Server struct {
	Addr      string
	Token     string
	IpHeader  string
	Router    *gin.Engine
	ApiGroup  *gin.RouterGroup
	rules     []*Rule
	rulesLock sync.RWMutex
}

const (
	queryVar  = `\$\{query\.(.+?)\}`
	bodyVar   = `\$\{body\.(.+?)\}`
	headerVar = `\$\{header\.(.+?)\}`
)

var (
	server           *Server
	once             sync.Once
	queryVarMatcher  = regexp.MustCompile(queryVar)
	bodyVarMatcher   = regexp.MustCompile(bodyVar)
	headerVarMatcher = regexp.MustCompile(headerVar)
)

func GetServer() *Server {
	once.Do(func() {
		letterBytes := "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
		randString := func(n int) string {
			b := make([]byte, n)
			for i := range b {
				b[i] = letterBytes[rand.Intn(len(letterBytes))]
			}
			return string(b)
		}
		server = &Server{
			Addr:  ":80",
			Token: randString(9),
			rules: make([]*Rule, 0),
		}
	})
	return server
}

func (s *Server) SetAddr(addr string) *Server {
	s.Addr = addr
	return s
}

func (s *Server) SetToken(token string) *Server {
	s.Token = token
	return s
}

func (s *Server) SetIpHeader(header string) *Server {
	s.IpHeader = header
	return s
}

func (s *Server) getRules() []*Rule {
	defer s.rulesLock.RUnlock()
	s.rulesLock.RLock()
	return s.rules
}

func (s *Server) updateRules() error {
	db := database.DB.Model(new(Rule))
	s.rulesLock.Lock()
	db.Order("rank desc").Find(&s.rules)
	s.rulesLock.Unlock()
	return nil
}

func (s *Server) Run() {
	if err := s.updateRules(); err != nil {
		log.Error(err.Error())
	}
	log.Info("Starting HTTP Server at %s, token:%s", s.Addr, s.Token)
	err := s.Router.Run(s.Addr)
	if err != nil {
		log.Fatal(err.Error())
	}
}

func getRawRequest(r *http.Request) ([]byte, error) {
	// to resolve httputil.DumpRequestOut error with "" protocol
	r.URL.Scheme = "http"
	r.URL.Host = "revsuit"

	return httputil.DumpRequestOut(r, true)
}

func compileTpl(c *gin.Context, tpl string) (compiled string) {
	compiled = tpl
	if queryVarMatcher.FindString(tpl) != "" {
		compiled = queryVarMatcher.ReplaceAllString(compiled, c.Query(queryVarMatcher.FindStringSubmatch(tpl)[1]))
	}

	if bodyVarMatcher.FindString(tpl) != "" {
		compiled = bodyVarMatcher.ReplaceAllString(compiled, c.PostForm(bodyVarMatcher.FindStringSubmatch(tpl)[1]))
	}

	if headerVarMatcher.FindString(tpl) != "" {
		compiled = headerVarMatcher.ReplaceAllString(compiled, c.GetHeader(headerVarMatcher.FindStringSubmatch(tpl)[1]))
	}
	return
}

func (s *Server) Receive(c *gin.Context) {
	u := c.Request.URL.String()
	for _, _rule := range s.getRules() {
		flag, flagGroup := _rule.Match(u)
		if flag == "" {
			continue
		}

		var (
			ip   = strings.Split(c.Request.RemoteAddr, ":")[0]
			area = qqwry.Area(ip)
		)

		if ip1 := c.Request.Header.Get(s.IpHeader); s.IpHeader != "" && ip1 != "" {
			ip = ip1
			delete(c.Request.Header, s.IpHeader)
		}

		raw, err := getRawRequest(c.Request)
		if err != nil {
			log.Warn(err.Error())
		}

		// create new record
		r, err := NewRecord(_rule, flag, c.Request.Method, u, ip, area, string(raw))
		if err != nil {
			log.Error("HTTP record(rule_id:%d) created failed :%s", _rule.ID, err.Error())
			code, err := strconv.Atoi(compileTpl(c, _rule.ResponseStatusCode))
			if err != nil || code < 100 || code > 600 {
				code = 400
			}

			c.String(code, compileTpl(c, _rule.ResponseBody))
			return
		}
		log.Trace("HTTP record(id:%d) has been created", r.ID)

		//only send to client when this connection recorded first time.
		if _rule.PushToClient {
			if flagGroup != "" {
				var count int64
				database.DB.Where("rule_name=? and raw like ?", _rule.Name, "%"+flagGroup+"%").Model(&Record{}).Count(&count)
				if count <= 1 {
					r.PushToClient()
					log.Trace("HTTP record(id:%d) has been put to client message queue", r.ID)
				}
			}
			r.PushToClient()
			log.Trace("HTTP record(id:%d) has been put to client message queue", r.ID)
		}

		//send notice
		if _rule.Notice {
			go func() {
				r.Notice()
				log.Trace("HTTP record(id:%d) notice has been sent", r.ID)
			}()
		}

		for header, value := range _rule.ResponseHeaders {
			c.Header(compileTpl(c, header), compileTpl(c, value))
		}

		code, err := strconv.Atoi(compileTpl(c, _rule.ResponseStatusCode))
		if err != nil || code < 100 || code > 600 {
			code = 400
		}

		c.String(code, compileTpl(c, _rule.ResponseBody))
		return
	}

}