Finish first working version of reqlog

pkg/reqlog/reqlog.go

@@ -1,51 +1,150 @@
 package reqlog
 
 import (
+	"bytes"
+	"compress/gzip"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"log"
 	"net/http"
 	"sync"
+	"time"
+
+	"github.com/dstotijn/gurp/pkg/proxy"
+	"github.com/google/uuid"
 )
 
 type Request struct {
-	Request http.Request
-	Body    []byte
+	ID        uuid.UUID
+	Request   http.Request
+	Body      []byte
+	Timestamp time.Time
+	Response  *Response
 }
 
-type response struct {
-	res  http.Response
-	body []byte
+type Response struct {
+	Response  http.Response
+	Body      []byte
+	Timestamp time.Time
 }
 
-type RequestLogStore struct {
-	reqStore []Request
-	resStore []response
-	reqMu    sync.Mutex
-	resMu    sync.Mutex
+type Service struct {
+	store []Request
+	mu    sync.Mutex
 }
 
-func NewRequestLogStore() RequestLogStore {
-	return RequestLogStore{
-		reqStore: make([]Request, 0),
-		resStore: make([]response, 0),
+func NewService() Service {
+	return Service{
+		store: make([]Request, 0),
 	}
 }
 
-func (store *RequestLogStore) AddRequest(req http.Request, body []byte) {
-	store.reqMu.Lock()
-	defer store.reqMu.Unlock()
+func (svc *Service) Requests() []Request {
+	// TODO(?): Is locking necessary here?
+	svc.mu.Lock()
+	defer svc.mu.Unlock()
 
-	store.reqStore = append(store.reqStore, Request{req, body})
+	return svc.store
 }
 
-func (store *RequestLogStore) Requests() []Request {
-	store.reqMu.Lock()
-	defer store.reqMu.Unlock()
+func (svc *Service) addRequest(reqID uuid.UUID, req http.Request, body []byte) Request {
+	svc.mu.Lock()
+	defer svc.mu.Unlock()
 
-	return store.reqStore
+	reqLog := Request{
+		ID:        reqID,
+		Request:   req,
+		Body:      body,
+		Timestamp: time.Now(),
+	}
+
+	svc.store = append(svc.store, reqLog)
+
+	return reqLog
 }
 
-func (store *RequestLogStore) AddResponse(res http.Response, body []byte) {
-	store.resMu.Lock()
-	defer store.resMu.Unlock()
+func (svc *Service) addResponse(reqID uuid.UUID, res http.Response, body []byte) error {
+	svc.mu.Lock()
+	defer svc.mu.Unlock()
 
-	store.resStore = append(store.resStore, response{res, body})
+	for i := range svc.store {
+		if svc.store[i].ID == reqID {
+			svc.store[i].Response = &Response{
+				Response:  res,
+				Body:      body,
+				Timestamp: time.Now(),
+			}
+			return nil
+		}
+	}
+
+	return fmt.Errorf("no request found with ID: %s", reqID)
+}
+
+func (svc *Service) RequestModifier(next proxy.RequestModifyFunc) proxy.RequestModifyFunc {
+	return func(req *http.Request) {
+		next(req)
+
+		clone := req.Clone(req.Context())
+		var body []byte
+		if req.Body != nil {
+			// TODO: Use io.LimitReader.
+			var err error
+			body, err = ioutil.ReadAll(req.Body)
+			if err != nil {
+				log.Printf("[ERROR] Could not read request body for logging: %v", err)
+				return
+			}
+			req.Body = ioutil.NopCloser(bytes.NewBuffer(body))
+		}
+
+		reqID, _ := req.Context().Value(proxy.ReqIDKey).(uuid.UUID)
+		if reqID == uuid.Nil {
+			log.Println("[ERROR] Request is missing a related request ID")
+			return
+		}
+
+		_ = svc.addRequest(reqID, *clone, body)
+	}
+}
+
+func (svc *Service) ResponseModifier(next proxy.ResponseModifyFunc) proxy.ResponseModifyFunc {
+	return func(res *http.Response) error {
+		if err := next(res); err != nil {
+			return err
+		}
+
+		clone := *res
+
+		// TODO: Use io.LimitReader.
+		body, err := ioutil.ReadAll(res.Body)
+		if err != nil {
+			return fmt.Errorf("reqlog: could not read response body: %v", err)
+		}
+		res.Body = ioutil.NopCloser(bytes.NewBuffer(body))
+
+		if res.Header.Get("Content-Encoding") == "gzip" {
+			gzipReader, err := gzip.NewReader(bytes.NewBuffer(body))
+			if err != nil {
+				return fmt.Errorf("reqlog: could not create gzip reader: %v", err)
+			}
+			defer gzipReader.Close()
+			body, err = ioutil.ReadAll(gzipReader)
+			if err != nil {
+				return fmt.Errorf("reqlog: could not read gzipped response body: %v", err)
+			}
+		}
+
+		reqID, _ := res.Request.Context().Value(proxy.ReqIDKey).(uuid.UUID)
+		if reqID == uuid.Nil {
+			return errors.New("reqlog: request is missing ID")
+		}
+
+		if err := svc.addResponse(reqID, clone, body); err != nil {
+			return fmt.Errorf("reqlog: could not add response: %v", err)
+		}
+
+		return nil
+	}
 }