mirror of
https://github.com/dstotijn/hetty.git
synced 2025-07-01 18:47:29 -04:00
Replace SQLite with BadgerDB
This commit is contained in:
@ -4,15 +4,18 @@ import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
"math/rand"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"time"
|
||||
|
||||
"github.com/dstotijn/hetty/pkg/proj"
|
||||
"github.com/oklog/ulid"
|
||||
|
||||
"github.com/dstotijn/hetty/pkg/proxy"
|
||||
"github.com/dstotijn/hetty/pkg/scope"
|
||||
"github.com/dstotijn/hetty/pkg/search"
|
||||
@ -22,127 +25,109 @@ type contextKey int
|
||||
|
||||
const LogBypassedKey contextKey = 0
|
||||
|
||||
const moduleName = "reqlog"
|
||||
var (
|
||||
ErrRequestNotFound = errors.New("reqlog: request not found")
|
||||
ErrProjectIDMustBeSet = errors.New("reqlog: project ID must be set")
|
||||
)
|
||||
|
||||
var ErrRequestNotFound = errors.New("reqlog: request not found")
|
||||
//nolint:gosec
|
||||
var ulidEntropy = rand.New(rand.NewSource(time.Now().UnixNano()))
|
||||
|
||||
type Request struct {
|
||||
ID int64
|
||||
Request http.Request
|
||||
Body []byte
|
||||
Timestamp time.Time
|
||||
Response *Response
|
||||
type RequestLog struct {
|
||||
ID ulid.ULID
|
||||
ProjectID ulid.ULID
|
||||
|
||||
URL *url.URL
|
||||
Method string
|
||||
Proto string
|
||||
Header http.Header
|
||||
Body []byte
|
||||
|
||||
Response *ResponseLog
|
||||
}
|
||||
|
||||
type Response struct {
|
||||
ID int64
|
||||
RequestID int64
|
||||
Response http.Response
|
||||
Body []byte
|
||||
Timestamp time.Time
|
||||
type ResponseLog struct {
|
||||
Proto string
|
||||
StatusCode int
|
||||
Status string
|
||||
Header http.Header
|
||||
Body []byte
|
||||
}
|
||||
|
||||
type Service struct {
|
||||
BypassOutOfScopeRequests bool
|
||||
FindReqsFilter FindRequestsFilter
|
||||
ActiveProjectID ulid.ULID
|
||||
|
||||
scope *scope.Scope
|
||||
repo Repository
|
||||
}
|
||||
|
||||
type FindRequestsFilter struct {
|
||||
OnlyInScope bool
|
||||
SearchExpr search.Expression `json:"-"`
|
||||
RawSearchExpr string
|
||||
ProjectID ulid.ULID
|
||||
OnlyInScope bool
|
||||
SearchExpr search.Expression
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
Scope *scope.Scope
|
||||
Repository Repository
|
||||
ProjectService proj.Service
|
||||
BypassOutOfScopeRequests bool
|
||||
Scope *scope.Scope
|
||||
Repository Repository
|
||||
}
|
||||
|
||||
func NewService(cfg Config) *Service {
|
||||
svc := &Service{
|
||||
scope: cfg.Scope,
|
||||
repo: cfg.Repository,
|
||||
BypassOutOfScopeRequests: cfg.BypassOutOfScopeRequests,
|
||||
return &Service{
|
||||
repo: cfg.Repository,
|
||||
scope: cfg.Scope,
|
||||
}
|
||||
|
||||
cfg.ProjectService.OnProjectOpen(func(_ string) error {
|
||||
err := svc.repo.FindSettingsByModule(context.Background(), moduleName, svc)
|
||||
if errors.Is(err, proj.ErrNoSettings) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("reqlog: could not load settings: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
cfg.ProjectService.OnProjectClose(func(_ string) error {
|
||||
svc.BypassOutOfScopeRequests = false
|
||||
svc.FindReqsFilter = FindRequestsFilter{}
|
||||
return nil
|
||||
})
|
||||
|
||||
return svc
|
||||
}
|
||||
|
||||
func (svc *Service) FindRequests(ctx context.Context) ([]Request, error) {
|
||||
func (svc *Service) FindRequests(ctx context.Context) ([]RequestLog, error) {
|
||||
return svc.repo.FindRequestLogs(ctx, svc.FindReqsFilter, svc.scope)
|
||||
}
|
||||
|
||||
func (svc *Service) FindRequestLogByID(ctx context.Context, id int64) (Request, error) {
|
||||
func (svc *Service) FindRequestLogByID(ctx context.Context, id ulid.ULID) (RequestLog, error) {
|
||||
return svc.repo.FindRequestLogByID(ctx, id)
|
||||
}
|
||||
|
||||
func (svc *Service) SetRequestLogFilter(ctx context.Context, filter FindRequestsFilter) error {
|
||||
svc.FindReqsFilter = filter
|
||||
return svc.repo.UpsertSettings(ctx, "reqlog", svc)
|
||||
func (svc *Service) ClearRequests(ctx context.Context, projectID ulid.ULID) error {
|
||||
return svc.repo.ClearRequestLogs(ctx, projectID)
|
||||
}
|
||||
|
||||
func (svc *Service) ClearRequests(ctx context.Context) error {
|
||||
return svc.repo.ClearRequestLogs(ctx)
|
||||
}
|
||||
|
||||
func (svc *Service) addRequest(
|
||||
ctx context.Context,
|
||||
req http.Request,
|
||||
body []byte,
|
||||
timestamp time.Time,
|
||||
) (*Request, error) {
|
||||
return svc.repo.AddRequestLog(ctx, req, body, timestamp)
|
||||
}
|
||||
|
||||
func (svc *Service) addResponse(
|
||||
ctx context.Context,
|
||||
reqID int64,
|
||||
res http.Response,
|
||||
body []byte,
|
||||
timestamp time.Time,
|
||||
) (*Response, error) {
|
||||
func (svc *Service) storeResponse(ctx context.Context, reqLogID ulid.ULID, res *http.Response) error {
|
||||
if res.Header.Get("Content-Encoding") == "gzip" {
|
||||
gzipReader, err := gzip.NewReader(bytes.NewBuffer(body))
|
||||
gzipReader, err := gzip.NewReader(res.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("reqlog: could not create gzip reader: %w", err)
|
||||
return fmt.Errorf("could not create gzip reader: %w", err)
|
||||
}
|
||||
defer gzipReader.Close()
|
||||
|
||||
body, err = ioutil.ReadAll(gzipReader)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("reqlog: could not read gzipped response body: %w", err)
|
||||
buf := &bytes.Buffer{}
|
||||
|
||||
if _, err := io.Copy(buf, gzipReader); err != nil {
|
||||
return fmt.Errorf("could not read gzipped response body: %w", err)
|
||||
}
|
||||
|
||||
res.Body = io.NopCloser(buf)
|
||||
}
|
||||
|
||||
return svc.repo.AddResponseLog(ctx, reqID, res, body, timestamp)
|
||||
body, err := io.ReadAll(res.Body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("could not read body: %w", err)
|
||||
}
|
||||
|
||||
resLog := ResponseLog{
|
||||
Proto: res.Proto,
|
||||
StatusCode: res.StatusCode,
|
||||
Status: res.Status,
|
||||
Header: res.Header,
|
||||
Body: body,
|
||||
}
|
||||
|
||||
return svc.repo.StoreResponseLog(ctx, reqLogID, resLog)
|
||||
}
|
||||
|
||||
func (svc *Service) RequestModifier(next proxy.RequestModifyFunc) proxy.RequestModifyFunc {
|
||||
return func(req *http.Request) {
|
||||
now := time.Now()
|
||||
|
||||
next(req)
|
||||
|
||||
clone := req.Clone(req.Context())
|
||||
@ -160,10 +145,19 @@ func (svc *Service) RequestModifier(next proxy.RequestModifyFunc) proxy.RequestM
|
||||
}
|
||||
|
||||
req.Body = ioutil.NopCloser(bytes.NewBuffer(body))
|
||||
clone.Body = ioutil.NopCloser(bytes.NewBuffer(body))
|
||||
}
|
||||
|
||||
// Bypass logging if no project is active.
|
||||
if svc.ActiveProjectID.Compare(ulid.ULID{}) == 0 {
|
||||
ctx := context.WithValue(req.Context(), LogBypassedKey, true)
|
||||
*req = *req.WithContext(ctx)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// Bypass logging if this setting is enabled and the incoming request
|
||||
// doens't match any rules of the scope.
|
||||
// doesn't match any scope rules.
|
||||
if svc.BypassOutOfScopeRequests && !svc.scope.Match(clone, body) {
|
||||
ctx := context.WithValue(req.Context(), LogBypassedKey, true)
|
||||
*req = *req.WithContext(ctx)
|
||||
@ -171,26 +165,29 @@ func (svc *Service) RequestModifier(next proxy.RequestModifyFunc) proxy.RequestM
|
||||
return
|
||||
}
|
||||
|
||||
reqLog, err := svc.addRequest(req.Context(), *clone, body, now)
|
||||
if errors.Is(err, proj.ErrNoProject) {
|
||||
ctx := context.WithValue(req.Context(), LogBypassedKey, true)
|
||||
*req = *req.WithContext(ctx)
|
||||
reqLog := RequestLog{
|
||||
ID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
|
||||
ProjectID: svc.ActiveProjectID,
|
||||
Method: clone.Method,
|
||||
URL: clone.URL,
|
||||
Proto: clone.Proto,
|
||||
Header: clone.Header,
|
||||
Body: body,
|
||||
}
|
||||
|
||||
return
|
||||
} else if err != nil {
|
||||
err := svc.repo.StoreRequestLog(req.Context(), reqLog)
|
||||
if err != nil {
|
||||
log.Printf("[ERROR] Could not store request log: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
ctx := context.WithValue(req.Context(), proxy.ReqIDKey, reqLog.ID)
|
||||
ctx := context.WithValue(req.Context(), proxy.ReqLogIDKey, reqLog.ID)
|
||||
*req = *req.WithContext(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
func (svc *Service) ResponseModifier(next proxy.ResponseModifyFunc) proxy.ResponseModifyFunc {
|
||||
return func(res *http.Response) error {
|
||||
now := time.Now()
|
||||
|
||||
if err := next(res); err != nil {
|
||||
return err
|
||||
}
|
||||
@ -199,8 +196,8 @@ func (svc *Service) ResponseModifier(next proxy.ResponseModifyFunc) proxy.Respon
|
||||
return nil
|
||||
}
|
||||
|
||||
reqID, _ := res.Request.Context().Value(proxy.ReqIDKey).(int64)
|
||||
if reqID == 0 {
|
||||
reqLogID, ok := res.Request.Context().Value(proxy.ReqLogIDKey).(ulid.ULID)
|
||||
if !ok {
|
||||
return errors.New("reqlog: request is missing ID")
|
||||
}
|
||||
|
||||
@ -213,9 +210,10 @@ func (svc *Service) ResponseModifier(next proxy.ResponseModifyFunc) proxy.Respon
|
||||
}
|
||||
|
||||
res.Body = ioutil.NopCloser(bytes.NewBuffer(body))
|
||||
clone.Body = ioutil.NopCloser(bytes.NewBuffer(body))
|
||||
|
||||
go func() {
|
||||
if _, err := svc.addResponse(context.Background(), reqID, clone, body, now); err != nil {
|
||||
if err := svc.storeResponse(context.Background(), reqLogID, &clone); err != nil {
|
||||
log.Printf("[ERROR] Could not store response log: %v", err)
|
||||
}
|
||||
}()
|
||||
@ -223,33 +221,3 @@ func (svc *Service) ResponseModifier(next proxy.ResponseModifyFunc) proxy.Respon
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// UnmarshalJSON implements json.Unmarshaler.
|
||||
func (f *FindRequestsFilter) UnmarshalJSON(b []byte) error {
|
||||
var dto struct {
|
||||
OnlyInScope bool
|
||||
RawSearchExpr string
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(b, &dto); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
filter := FindRequestsFilter{
|
||||
OnlyInScope: dto.OnlyInScope,
|
||||
RawSearchExpr: dto.RawSearchExpr,
|
||||
}
|
||||
|
||||
if dto.RawSearchExpr != "" {
|
||||
expr, err := search.ParseQuery(dto.RawSearchExpr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
filter.SearchExpr = expr
|
||||
}
|
||||
|
||||
*f = filter
|
||||
|
||||
return nil
|
||||
}
|
||||
|
Reference in New Issue
Block a user