chore: tunnel recorder & slight refactor

store/context.go

@@ -0,0 +1,18 @@
+package store
+
+import (
+	"context"
+)
+
+type contextKey struct{}
+
+var recordIDKey = contextKey{}
+
+func withRecord(ctx context.Context, rec *TunnelRecord) context.Context {
+	return context.WithValue(ctx, recordIDKey, rec)
+}
+
+func getRecord(ctx context.Context) (*TunnelRecord, bool) {
+	id, ok := ctx.Value(recordIDKey).(*TunnelRecord)
+	return id, ok
+}

store/store.go

@@ -0,0 +1,196 @@
+package store
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"mime"
+	"net/http"
+	"net/url"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/google/uuid"
+)
+
+const (
+	defaultQueueSize = 100
+	maxQueueSize     = 100
+)
+
+var ErrRecordNotFound = errors.New("record not found")
+
+type TunnelStore interface {
+	Get(before time.Time, count int) (results []*TunnelRecord, more bool)
+	RecordTCP()
+	RecordRequest(req *http.Request)
+	RecordResponse(resp *http.Response) error
+}
+
+type TunnelRecord struct {
+	ID     uuid.UUID
+	Time   time.Time
+	URL    *url.URL
+	Method string
+	Status int
+
+	RequestHeaders  http.Header
+	RequestBodyType string
+	RequestBody     []byte
+
+	ResponseHeaders  http.Header
+	ResponseBodyType string
+	ResponseBody     []byte
+}
+
+func NewTunnelStore(queueSize int) TunnelStore {
+	if queueSize <= 0 {
+		queueSize = defaultQueueSize
+	} else if queueSize > maxQueueSize {
+		queueSize = maxQueueSize
+	}
+
+	return &tunnelStoreImpl{queueSize: queueSize}
+}
+
+type tunnelStoreImpl struct {
+	orderedRecords []*TunnelRecord
+	queueSize      int
+	mu             sync.Mutex
+}
+
+func (s *tunnelStoreImpl) Get(before time.Time, count int) ([]*TunnelRecord, bool) {
+	// Find First
+	start := -1
+	for i, r := range s.orderedRecords {
+		if r.Time.Before(before) {
+			start = i
+			break
+		}
+	}
+
+	// Not Found
+	if start == -1 {
+		return nil, false
+	}
+
+	// Subslice Records
+	end := min(start+count, len(s.orderedRecords))
+	results := s.orderedRecords[start:end]
+	more := end < len(s.orderedRecords)
+
+	return results, more
+}
+
+func (s *tunnelStoreImpl) RecordRequest(req *http.Request) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	url := *req.URL
+	rec := &TunnelRecord{
+		ID:              uuid.New(),
+		Time:            time.Now(),
+		URL:             &url,
+		Method:          req.Method,
+		RequestHeaders:  req.Header,
+		RequestBodyType: req.Header.Get("Content-Type"),
+	}
+
+	if bodyData, err := getRequestBody(req); err == nil {
+		rec.RequestBody = bodyData
+	}
+
+	// Add Record & Truncate
+	s.orderedRecords = append(s.orderedRecords, rec)
+	if len(s.orderedRecords) > s.queueSize {
+		s.orderedRecords = s.orderedRecords[len(s.orderedRecords)-s.queueSize:]
+	}
+
+	*req = *req.WithContext(withRecord(req.Context(), rec))
+}
+
+func (s *tunnelStoreImpl) RecordResponse(resp *http.Response) error {
+	rec, found := getRecord(resp.Request.Context())
+	if !found {
+		return ErrRecordNotFound
+	}
+
+	rec.ResponseHeaders = resp.Header
+	rec.ResponseBodyType = resp.Header.Get("Content-Type")
+
+	if bodyData, err := getResponseBody(resp); err == nil {
+		rec.ResponseBody = bodyData
+	}
+
+	return nil
+}
+
+func (s *tunnelStoreImpl) RecordTCP() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// TODO
+}
+
+func getRequestBody(req *http.Request) ([]byte, error) {
+	if req.ContentLength == 0 || req.Body == nil || req.Body == http.NoBody {
+		return nil, nil
+	}
+
+	if !isTextContentType(req.Header.Get("Content-Type")) {
+		return nil, nil
+	}
+
+	// Read Body
+	bodyBytes, err := io.ReadAll(req.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	// Restore Body
+	req.Body = io.NopCloser(bytes.NewReader(bodyBytes))
+	return bodyBytes, nil
+}
+
+func getResponseBody(resp *http.Response) ([]byte, error) {
+	if resp.ContentLength == 0 || resp.Body == nil || resp.Body == http.NoBody {
+		return nil, nil
+	}
+
+	if !isTextContentType(resp.Header.Get("Content-Type")) {
+		return nil, nil
+	}
+
+	// Read Body
+	bodyBytes, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	// Restore Body
+	resp.Body = io.NopCloser(bytes.NewReader(bodyBytes))
+	return bodyBytes, nil
+}
+
+func isTextContentType(contentType string) bool {
+	mediaType, _, err := mime.ParseMediaType(contentType)
+	if err != nil {
+		return false
+	}
+
+	if strings.HasPrefix(mediaType, "text/") {
+		return true
+	}
+
+	switch mediaType {
+	case "application/json":
+		return true
+	case "application/xml":
+		return true
+	case "application/x-www-form-urlencoded":
+		return true
+	default:
+		return false
+	}
+}