go-gpt/router/router.go

package router

import (
	"github.com/gogf/gf/v2/errors/gcode"
	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/net/ghttp"
	"github.com/gogf/gf/v2/os/glog"
	"go-gpt/internal/controller"
	"go-gpt/internal/service"
	"net/http"
	"sync"
	"time"
)

type DefaultHandlerResponse struct {
	Code    int         `json:"status"`
	Message string      `json:"message"`
	Data    interface{} `json:"data"`
}

// MiddlewareHandlerResponse is the default middleware handling handler response object and its error.
func MiddlewareHandlerResponse(r *ghttp.Request) {
	r.Middleware.Next()
	// There's custom buffer content, it then exits current handler.
	if r.Response.BufferLength() > 0 {
		return
	}

	var (
		msg  string
		err  = r.GetError()
		res  = r.GetHandlerResponse()
		code = gerror.Code(err)
	)
	if err != nil {
		if code == gcode.CodeNil {
			code = gcode.CodeInternalError
		}
		msg = err.Error()
	} else {
		if r.Response.Status > 0 && r.Response.Status != http.StatusOK {
			msg = http.StatusText(r.Response.Status)
			switch r.Response.Status {
			case http.StatusNotFound:
				code = gcode.CodeNotFound
			case http.StatusForbidden:
				code = gcode.CodeNotAuthorized
			default:
				code = gcode.CodeUnknown
			}
			// It creates error as it can be retrieved by other middlewares.
			err = gerror.NewCode(code, msg)
			r.SetError(err)
		} else {
			code = gcode.CodeOK
		}
	}

	r.Response.WriteJson(DefaultHandlerResponse{
		Code:    code.Code(),
		Message: msg,
		Data:    res,
	})
}

type middlewareService struct{}

func NewTokenMiddleware() *middlewareService {
	return &middlewareService{}
}

func (s *middlewareService) Auth(r *ghttp.Request) {
	service.Auth().MiddlewareFunc()(r)
	r.Middleware.Next()
}

type RateLimitMiddleware struct {
	sync.Mutex
	Counter map[string]int64
}

func NewRateLimitMiddleware() *RateLimitMiddleware {
	return &RateLimitMiddleware{Counter: make(map[string]int64)}
}
func (m *RateLimitMiddleware) Middleware(r *ghttp.Request) {
	ip := r.GetClientIp()
	m.Lock()
	defer m.Unlock()

	now := time.Now().Unix()
	count, ok := m.Counter[ip]
	if !ok {
		count = 1
		m.Counter[ip] = now
	} else {
		if now-m.Counter[ip] < 1 {
			r.Response.WriteStatusExit(429, "Too Many Requests")
		} else {
			m.Counter[ip] = now
			count++
		}
	}

	glog.Debugf(r.Context(), "%s:%d", ip, count)

	r.Middleware.Next()
}
func MiddlewareCORS(r *ghttp.Request) {

	corsOptions := r.Response.DefaultCORSOptions()
	//corsOptions.AllowDomain = []string{"goframe.org", "baidu.com"}
	r.Response.CORS(corsOptions)

	r.Response.CORSDefault()

	r.Middleware.Next()
}
func BindController(group *ghttp.RouterGroup) {
	DomeRouter(group)
	group.Group("/api/v1", func(group *ghttp.RouterGroup) {
		vv := NewRateLimitMiddleware()
		group.Middleware(MiddlewareCORS, vv.Middleware)
		SessionRouter(group)
		ChatRouter(group)
	})
}

func DomeRouter(group *ghttp.RouterGroup) {
	//tokenMi := NewTokenMiddleware()
	//group.Middleware(tokenMi.Auth)
	group.Bind(
		controller.Hello,
	)
}

func ChatRouter(group *ghttp.RouterGroup) {
	//group.Middleware(MiddlewareHandlerResponse)
	group.Bind(
		controller.Chat,
	)
}

func SessionRouter(group *ghttp.RouterGroup) {
	//group.Middleware(MiddlewareHandlerResponse)
	group.Group("/session", func(group *ghttp.RouterGroup) {
		group.Bind(
			controller.Session,
		)
	})
}