Add Sender module

pkg/db/badger/badger.go

@@ -8,12 +8,16 @@ import (
 
 const (
 	// Key prefixes. Each prefix value should be unique.
-	projectPrefix = 0x00
-	reqLogPrefix  = 0x01
-	resLogPrefix  = 0x02
+	projectPrefix   = 0x00
+	reqLogPrefix    = 0x01
+	resLogPrefix    = 0x02
+	senderReqPrefix = 0x03
 
 	// Request log indices.
 	reqLogProjectIDIndex = 0x00
+
+	// Sender request indices.
+	senderReqProjectIDIndex = 0x00
 )
 
 // Database is used to store and retrieve data from an underlying Badger database.

pkg/db/badger/proj.go

@@ -65,6 +65,11 @@ func (db *Database) DeleteProject(ctx context.Context, projectID ulid.ULID) erro
 		return fmt.Errorf("badger: failed to delete project request logs: %w", err)
 	}
 
+	err = db.DeleteSenderRequests(ctx, projectID)
+	if err != nil {
+		return fmt.Errorf("badger: failed to delete project sender requests: %w", err)
+	}
+
 	err = db.badger.Update(func(txn *badger.Txn) error {
 		return txn.Delete(entryKey(projectPrefix, 0, projectID[:]))
 	})

pkg/db/badger/proj_test.go

@@ -56,6 +56,7 @@ func TestUpsertProject(t *testing.T) {
 		Settings: proj.Settings{
 			ReqLogBypassOutOfScope: true,
 			ReqLogOnlyFindInScope:  true,
+			ReqLogSearchExpr:       searchExpr,
 			ScopeRules: []scope.Rule{
 				{
 					URL: regexp.MustCompile("^https://(.*)example.com(.*)$"),
@@ -66,7 +67,6 @@ func TestUpsertProject(t *testing.T) {
 					Body: regexp.MustCompile("^foo(.*)"),
 				},
 			},
-			SearchExpr: searchExpr,
 		},
 	}
 
@@ -178,21 +178,38 @@ func TestDeleteProject(t *testing.T) {
 	// Store fixtures.
 	projectID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
 	reqLogID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
+	senderReqID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
 
 	err = badgerDB.Update(func(txn *badgerdb.Txn) error {
+		// Project item.
 		if err := txn.Set(entryKey(projectPrefix, 0, projectID[:]), nil); err != nil {
 			return err
 		}
+
+		// Sender request items.
+		if err := txn.Set(entryKey(senderReqPrefix, 0, senderReqID[:]), nil); err != nil {
+			return err
+		}
+		if err := txn.Set(entryKey(resLogPrefix, 0, senderReqID[:]), nil); err != nil {
+			return err
+		}
+		err := txn.Set(entryKey(senderReqPrefix, senderReqProjectIDIndex, append(projectID[:], senderReqID[:]...)), nil)
+		if err != nil {
+			return err
+		}
+
+		// Request log items.
 		if err := txn.Set(entryKey(reqLogPrefix, 0, reqLogID[:]), nil); err != nil {
 			return err
 		}
 		if err := txn.Set(entryKey(resLogPrefix, 0, reqLogID[:]), nil); err != nil {
 			return err
 		}
-		err := txn.Set(entryKey(reqLogPrefix, reqLogProjectIDIndex, append(projectID[:], reqLogID[:]...)), nil)
+		err = txn.Set(entryKey(reqLogPrefix, reqLogProjectIDIndex, append(projectID[:], reqLogID[:]...)), nil)
 		if err != nil {
 			return err
 		}
+
 		return nil
 	})
 	if err != nil {
@@ -222,7 +239,7 @@ func TestDeleteProject(t *testing.T) {
 		t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
 	}
 
-	// Assert response log item was deleted.
+	// Assert response log item related to request log was deleted.
 	err = badgerDB.View(func(txn *badgerdb.Txn) error {
 		_, err := txn.Get(entryKey(resLogPrefix, 0, reqLogID[:]))
 		return err
@@ -239,6 +256,33 @@ func TestDeleteProject(t *testing.T) {
 	if !errors.Is(err, badgerdb.ErrKeyNotFound) {
 		t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
 	}
+
+	// Assert sender request item was deleted.
+	err = badgerDB.View(func(txn *badgerdb.Txn) error {
+		_, err := txn.Get(entryKey(senderReqPrefix, 0, senderReqID[:]))
+		return err
+	})
+	if !errors.Is(err, badgerdb.ErrKeyNotFound) {
+		t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
+	}
+
+	// Assert response log item related to sender request was deleted.
+	err = badgerDB.View(func(txn *badgerdb.Txn) error {
+		_, err := txn.Get(entryKey(resLogPrefix, 0, senderReqID[:]))
+		return err
+	})
+	if !errors.Is(err, badgerdb.ErrKeyNotFound) {
+		t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
+	}
+
+	// Assert sender request project ID index key was deleted.
+	err = badgerDB.View(func(txn *badgerdb.Txn) error {
+		_, err := txn.Get(entryKey(senderReqPrefix, senderReqProjectIDIndex, append(projectID[:], senderReqID[:]...)))
+		return err
+	})
+	if !errors.Is(err, badgerdb.ErrKeyNotFound) {
+		t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
+	}
 }
 
 func TestProjects(t *testing.T) {

pkg/db/badger/reqlog.go

@@ -66,7 +66,11 @@ func (db *Database) FindRequestLogs(ctx context.Context, filter reqlog.FindReque
 
 func getRequestLogWithResponse(txn *badger.Txn, reqLogID ulid.ULID) (reqlog.RequestLog, error) {
 	item, err := txn.Get(entryKey(reqLogPrefix, 0, reqLogID[:]))
-	if err != nil {
+
+	switch {
+	case errors.Is(err, badger.ErrKeyNotFound):
+		return reqlog.RequestLog{}, reqlog.ErrRequestNotFound
+	case err != nil:
 		return reqlog.RequestLog{}, fmt.Errorf("failed to lookup request log item: %w", err)
 	}
 

pkg/db/badger/sender.go

@@ -0,0 +1,236 @@
+package badger
+
+import (
+	"bytes"
+	"context"
+	"encoding/gob"
+	"errors"
+	"fmt"
+
+	"github.com/dgraph-io/badger/v3"
+	"github.com/oklog/ulid"
+
+	"github.com/dstotijn/hetty/pkg/reqlog"
+	"github.com/dstotijn/hetty/pkg/scope"
+	"github.com/dstotijn/hetty/pkg/sender"
+)
+
+func (db *Database) StoreSenderRequest(ctx context.Context, req sender.Request) error {
+	buf := bytes.Buffer{}
+
+	err := gob.NewEncoder(&buf).Encode(req)
+	if err != nil {
+		return fmt.Errorf("badger: failed to encode sender request: %w", err)
+	}
+
+	entries := []*badger.Entry{
+		// Sender request itself.
+		{
+			Key:   entryKey(senderReqPrefix, 0, req.ID[:]),
+			Value: buf.Bytes(),
+		},
+		// Index by project ID.
+		{
+			Key: entryKey(senderReqPrefix, senderReqProjectIDIndex, append(req.ProjectID[:], req.ID[:]...)),
+		},
+	}
+
+	err = db.badger.Update(func(txn *badger.Txn) error {
+		for i := range entries {
+			err := txn.SetEntry(entries[i])
+			if err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		return fmt.Errorf("badger: failed to commit transaction: %w", err)
+	}
+
+	return nil
+}
+
+func (db *Database) FindSenderRequestByID(ctx context.Context, senderReqID ulid.ULID) (sender.Request, error) {
+	txn := db.badger.NewTransaction(false)
+	defer txn.Discard()
+
+	req, err := getSenderRequestWithResponseLog(txn, senderReqID)
+	if err != nil {
+		return sender.Request{}, fmt.Errorf("badger: failed to get sender request: %w", err)
+	}
+
+	return req, nil
+}
+
+func (db *Database) FindSenderRequests(ctx context.Context, filter sender.FindRequestsFilter, scope *scope.Scope) ([]sender.Request, error) {
+	if filter.ProjectID.Compare(ulid.ULID{}) == 0 {
+		return nil, sender.ErrProjectIDMustBeSet
+	}
+
+	txn := db.badger.NewTransaction(false)
+	defer txn.Discard()
+
+	senderReqIDs, err := findSenderRequestIDsByProjectID(txn, filter.ProjectID)
+	if err != nil {
+		return nil, fmt.Errorf("badger: failed to find sender request IDs: %w", err)
+	}
+
+	senderReqs := make([]sender.Request, 0, len(senderReqIDs))
+
+	for _, id := range senderReqIDs {
+		senderReq, err := getSenderRequestWithResponseLog(txn, id)
+		if err != nil {
+			return nil, fmt.Errorf("badger: failed to get sender request (id: %v): %w", id.String(), err)
+		}
+
+		if filter.OnlyInScope {
+			if !senderReq.MatchScope(scope) {
+				continue
+			}
+		}
+
+		// Filter by search expression.
+		// TODO: Once pagination is introduced, this filter logic should be done
+		// as items are retrieved (e.g. when using a `badger.Iterator`).
+		if filter.SearchExpr != nil {
+			match, err := senderReq.Matches(filter.SearchExpr)
+			if err != nil {
+				return nil, fmt.Errorf(
+					"badger: failed to match search expression for sender request (id: %v): %w",
+					id.String(), err,
+				)
+			}
+
+			if !match {
+				continue
+			}
+		}
+
+		senderReqs = append(senderReqs, senderReq)
+	}
+
+	return senderReqs, nil
+}
+
+func (db *Database) DeleteSenderRequests(ctx context.Context, projectID ulid.ULID) error {
+	// Note: this transaction is used just for reading; we use the `badger.WriteBatch`
+	// API to bulk delete items.
+	txn := db.badger.NewTransaction(false)
+	defer txn.Discard()
+
+	senderReqIDs, err := findSenderRequestIDsByProjectID(txn, projectID)
+	if err != nil {
+		return fmt.Errorf("badger: failed to find sender request IDs: %w", err)
+	}
+
+	writeBatch := db.badger.NewWriteBatch()
+	defer writeBatch.Cancel()
+
+	for _, senderReqID := range senderReqIDs {
+		// Delete sender requests.
+		err := writeBatch.Delete(entryKey(senderReqPrefix, 0, senderReqID[:]))
+		if err != nil {
+			return fmt.Errorf("badger: failed to delete sender requests: %w", err)
+		}
+
+		// Delete related response log.
+		err = writeBatch.Delete(entryKey(resLogPrefix, 0, senderReqID[:]))
+		if err != nil {
+			return fmt.Errorf("badger: failed to delete request log: %w", err)
+		}
+	}
+
+	if err := writeBatch.Flush(); err != nil {
+		return fmt.Errorf("badger: failed to commit batch write: %w", err)
+	}
+
+	err = db.badger.DropPrefix(entryKey(senderReqPrefix, senderReqProjectIDIndex, projectID[:]))
+	if err != nil {
+		return fmt.Errorf("badger: failed to drop sender request project ID index items: %w", err)
+	}
+
+	return nil
+}
+
+func getSenderRequestWithResponseLog(txn *badger.Txn, senderReqID ulid.ULID) (sender.Request, error) {
+	item, err := txn.Get(entryKey(senderReqPrefix, 0, senderReqID[:]))
+
+	switch {
+	case errors.Is(err, badger.ErrKeyNotFound):
+		return sender.Request{}, sender.ErrRequestNotFound
+	case err != nil:
+		return sender.Request{}, fmt.Errorf("failed to lookup sender request item: %w", err)
+	}
+
+	req := sender.Request{
+		ID: senderReqID,
+	}
+
+	err = item.Value(func(rawSenderReq []byte) error {
+		err = gob.NewDecoder(bytes.NewReader(rawSenderReq)).Decode(&req)
+		if err != nil {
+			return fmt.Errorf("failed to decode sender request: %w", err)
+		}
+
+		return nil
+	})
+	if err != nil {
+		return sender.Request{}, fmt.Errorf("failed to retrieve or parse sender request value: %w", err)
+	}
+
+	item, err = txn.Get(entryKey(resLogPrefix, 0, senderReqID[:]))
+
+	if errors.Is(err, badger.ErrKeyNotFound) {
+		return req, nil
+	}
+
+	if err != nil {
+		return sender.Request{}, fmt.Errorf("failed to get response log: %w", err)
+	}
+
+	err = item.Value(func(rawReslog []byte) error {
+		var resLog reqlog.ResponseLog
+		err = gob.NewDecoder(bytes.NewReader(rawReslog)).Decode(&resLog)
+		if err != nil {
+			return fmt.Errorf("failed to decode response log: %w", err)
+		}
+
+		req.Response = &resLog
+
+		return nil
+	})
+	if err != nil {
+		return sender.Request{}, fmt.Errorf("failed to retrieve or parse response log value: %w", err)
+	}
+
+	return req, nil
+}
+
+func findSenderRequestIDsByProjectID(txn *badger.Txn, projectID ulid.ULID) ([]ulid.ULID, error) {
+	senderReqIDs := make([]ulid.ULID, 0)
+	opts := badger.DefaultIteratorOptions
+	opts.PrefetchValues = false
+	opts.Reverse = true
+	iterator := txn.NewIterator(opts)
+	defer iterator.Close()
+
+	var projectIndexKey []byte
+
+	prefix := entryKey(senderReqPrefix, senderReqProjectIDIndex, projectID[:])
+
+	for iterator.Seek(append(prefix, 255)); iterator.ValidForPrefix(prefix); iterator.Next() {
+		projectIndexKey = iterator.Item().KeyCopy(projectIndexKey)
+
+		var id ulid.ULID
+		// The request log ID starts *after* the first 2 prefix and index bytes
+		// and the 16 byte project ID.
+		if err := id.UnmarshalBinary(projectIndexKey[18:]); err != nil {
+			return nil, fmt.Errorf("failed to parse sender request ID: %w", err)
+		}
+
+		senderReqIDs = append(senderReqIDs, id)
+	}
+
+	return senderReqIDs, nil
+}

pkg/db/badger/sender_test.go

@@ -0,0 +1,204 @@
+package badger_test
+
+import (
+	"context"
+	"errors"
+	"math/rand"
+	"net/http"
+	"net/url"
+	"testing"
+	"time"
+
+	badgerdb "github.com/dgraph-io/badger/v3"
+	"github.com/google/go-cmp/cmp"
+	"github.com/oklog/ulid"
+
+	"github.com/dstotijn/hetty/pkg/db/badger"
+	"github.com/dstotijn/hetty/pkg/reqlog"
+	"github.com/dstotijn/hetty/pkg/sender"
+)
+
+//nolint:gosec
+var ulidEntropy = rand.New(rand.NewSource(time.Now().UnixNano()))
+
+var exampleURL = func() *url.URL {
+	u, err := url.Parse("https://example.com/foobar")
+	if err != nil {
+		panic(err)
+	}
+
+	return u
+}()
+
+func TestFindRequestByID(t *testing.T) {
+	t.Parallel()
+
+	database, err := badger.OpenDatabase(badgerdb.DefaultOptions("").WithInMemory(true))
+	if err != nil {
+		t.Fatalf("failed to open badger database: %v", err)
+	}
+	defer database.Close()
+
+	// See: https://go.dev/blog/subtests#cleaning-up-after-a-group-of-parallel-tests
+	t.Run("group", func(t *testing.T) {
+		t.Run("sender request not found", func(t *testing.T) {
+			t.Parallel()
+
+			_, err := database.FindSenderRequestByID(context.Background(), ulid.ULID{})
+			if !errors.Is(err, sender.ErrRequestNotFound) {
+				t.Fatalf("expected `sender.ErrRequestNotFound`, got: %v", err)
+			}
+		})
+
+		t.Run("sender request found", func(t *testing.T) {
+			t.Parallel()
+
+			exp := sender.Request{
+				ID:                 ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
+				ProjectID:          ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
+				SourceRequestLogID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
+
+				URL:    exampleURL,
+				Method: http.MethodGet,
+				Proto:  sender.HTTPProto2,
+				Header: http.Header{
+					"X-Foo": []string{"bar"},
+				},
+				Body: []byte("foo"),
+			}
+
+			err := database.StoreSenderRequest(context.Background(), exp)
+			if err != nil {
+				t.Fatalf("unexpected error (expected: nil, got: %v)", err)
+			}
+
+			resLog := reqlog.ResponseLog{
+				Proto:      "HTTP/2.0",
+				Status:     "200 OK",
+				StatusCode: 200,
+				Header: http.Header{
+					"X-Yolo": []string{"swag"},
+				},
+				Body: []byte("bar"),
+			}
+
+			err = database.StoreResponseLog(context.Background(), exp.ID, resLog)
+			if err != nil {
+				t.Fatalf("unexpected error (expected: nil, got: %v)", err)
+			}
+
+			exp.Response = &resLog
+
+			got, err := database.FindSenderRequestByID(context.Background(), exp.ID)
+			if err != nil {
+				t.Fatalf("unexpected error (expected: nil, got: %v)", err)
+			}
+
+			if diff := cmp.Diff(exp, got); diff != "" {
+				t.Fatalf("sender request not equal (-exp, +got):\n%v", diff)
+			}
+		})
+	})
+}
+
+func TestFindSenderRequests(t *testing.T) {
+	t.Parallel()
+
+	t.Run("without project ID in filter", func(t *testing.T) {
+		t.Parallel()
+
+		database, err := badger.OpenDatabase(badgerdb.DefaultOptions("").WithInMemory(true))
+		if err != nil {
+			t.Fatalf("failed to open badger database: %v", err)
+		}
+		defer database.Close()
+
+		filter := sender.FindRequestsFilter{}
+
+		_, err = database.FindSenderRequests(context.Background(), filter, nil)
+		if !errors.Is(err, sender.ErrProjectIDMustBeSet) {
+			t.Fatalf("expected `sender.ErrProjectIDMustBeSet`, got: %v", err)
+		}
+	})
+
+	t.Run("returns sender requests and related response logs", func(t *testing.T) {
+		t.Parallel()
+
+		database, err := badger.OpenDatabase(badgerdb.DefaultOptions("").WithInMemory(true))
+		if err != nil {
+			t.Fatalf("failed to open badger database: %v", err)
+		}
+		defer database.Close()
+
+		projectID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
+
+		fixtures := []sender.Request{
+			{
+				ID:                 ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
+				ProjectID:          projectID,
+				SourceRequestLogID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
+				URL:                exampleURL,
+				Method:             http.MethodPost,
+				Proto:              "HTTP/1.1",
+				Header: http.Header{
+					"X-Foo": []string{"baz"},
+				},
+				Body: []byte("foo"),
+				Response: &reqlog.ResponseLog{
+					Proto:      "HTTP/1.1",
+					Status:     "200 OK",
+					StatusCode: 200,
+					Header: http.Header{
+						"X-Yolo": []string{"swag"},
+					},
+					Body: []byte("bar"),
+				},
+			},
+			{
+				ID:                 ulid.MustNew(ulid.Timestamp(time.Now())+100, ulidEntropy),
+				ProjectID:          projectID,
+				SourceRequestLogID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
+				URL:                exampleURL,
+				Method:             http.MethodGet,
+				Proto:              "HTTP/1.1",
+				Header: http.Header{
+					"X-Foo": []string{"baz"},
+				},
+			},
+		}
+
+		// Store fixtures.
+		for _, senderReq := range fixtures {
+			err = database.StoreSenderRequest(context.Background(), senderReq)
+			if err != nil {
+				t.Fatalf("unexpected error creating request log fixture: %v", err)
+			}
+
+			if senderReq.Response != nil {
+				err = database.StoreResponseLog(context.Background(), senderReq.ID, *senderReq.Response)
+				if err != nil {
+					t.Fatalf("unexpected error creating response log fixture: %v", err)
+				}
+			}
+		}
+
+		filter := sender.FindRequestsFilter{
+			ProjectID: projectID,
+		}
+
+		got, err := database.FindSenderRequests(context.Background(), filter, nil)
+		if err != nil {
+			t.Fatalf("unexpected error finding sender requests: %v", err)
+		}
+
+		// We expect the found sender requests are *reversed*, e.g. newest first.
+		exp := make([]sender.Request, len(fixtures))
+		for i, j := 0, len(fixtures)-1; i < j; i, j = i+1, j-1 {
+			exp[i], exp[j] = fixtures[j], fixtures[i]
+		}
+
+		if diff := cmp.Diff(exp, got); diff != "" {
+			t.Fatalf("sender requests not equal (-exp, +got):\n%v", diff)
+		}
+	})
+}