WIP - TrackFilters

TrackFilters allow modifying what is stored in tracks.
An example use case is filtering Authorization headers and cookies.

Author: Michael Weber

cassette.go

@@ -511,3 +511,135 @@ func decompress(data []byte) ([]byte, error) {
 
 	return data, nil
 }
+
+// Track Filters
+
+// TrackFilter allows modifying information that is being stored in a track
+//
+// The filters receive a deep copy of the request and response, that can be modified.
+type TrackFilter func(*http.Request, *http.Response, error) (*http.Request, *http.Response)
+
+// TrackFilters is a slice of TrackFilter
+type TrackFilters []TrackFilter
+
+// TrackRequestDeleteHeaderKeys will delete one or more header keys on the request
+// before the request is matched against the cassette.
+func TrackRequestDeleteHeaderKeys(keys ...string) TrackFilter {
+	return func(req *http.Request, resp *http.Response, err error) (*http.Request, *http.Response) {
+		for _, key := range keys {
+			req.Header.Del(key)
+		}
+		return req, resp
+	}
+}
+
+// TrackResponseDeleteHeaderKeys will delete one or more header keys on the request
+// before the request is matched against the cassette.
+func TrackResponseDeleteHeaderKeys(keys ...string) TrackFilter {
+	return func(req *http.Request, resp *http.Response, err error) (*http.Request, *http.Response) {
+		if resp == nil || resp.Header == nil {
+			return req, resp
+		}
+		for _, key := range keys {
+			resp.Header.Del(key)
+		}
+		return req, resp
+	}
+}
+
+// TrackRequestChangeBody will allows to change the body.
+// Supply a function that does input to output transformation.
+func TrackRequestChangeBody(fn func(b []byte) []byte) TrackFilter {
+	return func(req *http.Request, resp *http.Response, err error) (*http.Request, *http.Response) {
+		body, err2 := readRequestBody(req) // XXX does too much, don't recreate body
+		if err2 != nil {
+			return req, resp
+		}
+		body = fn(body)
+		req.Body = toReadCloser(body)
+		req.ContentLength = int64(len(body))
+		return req, resp
+	}
+}
+
+// TrackResponseChangeBody will allows to change the body.
+// Supply a function that does input to output transformation.
+func TrackResponseChangeBody(fn func(b []byte) []byte) TrackFilter {
+	return func(req *http.Request, resp *http.Response, err error) (*http.Request, *http.Response) {
+		if resp == nil || resp.Header == nil {
+			return req, resp
+		}
+		body, err2 := readResponseBody(resp) // XXX does too much, don't recreate body
+		if err2 != nil {
+			return req, resp
+		}
+		body = fn(body)
+		resp.Body = toReadCloser(body)
+		resp.ContentLength = int64(len(body))
+		return req, resp
+	}
+}
+
+// OnMethod will return a new filter that will only apply 'r'
+// if the method of the request matches.
+// Original filter is unmodified.
+func (r TrackFilter) OnMethod(method string) TrackFilter {
+	return func(req *http.Request, resp *http.Response, err error) (*http.Request, *http.Response) {
+		if req.Method != method {
+			return req, resp
+		}
+		return r(req, resp, err)
+	}
+}
+
+// OnPath will return a track filter that will only apply 'r'
+// if the url string of the request matches the supplied regex.
+// Original filter is unmodified.
+func (r TrackFilter) OnPath(pathRegEx string) TrackFilter {
+	if pathRegEx == "" {
+		pathRegEx = ".*"
+	}
+	re := regexp.MustCompile(pathRegEx)
+	return func(req *http.Request, resp *http.Response, err error) (*http.Request, *http.Response) {
+		if !re.MatchString(req.URL.String()) {
+			return req, resp
+		}
+		return r(req, resp, err)
+	}
+}
+
+// OnStatus will return a Track filter that will only apply 'r'  if the response status matches.
+// Original filter is unmodified.
+func (r TrackFilter) OnStatus(status int) TrackFilter {
+	return func(req *http.Request, resp *http.Response, err error) (*http.Request, *http.Response) {
+		if resp == nil || resp.StatusCode != status {
+			return req, resp
+		}
+		return r(req, resp, err)
+	}
+}
+
+// Add one or more filters at the end of the filter chain.
+func (r *TrackFilters) Add(filters ...TrackFilter) {
+	v := *r
+	v = append(v, filters...)
+	*r = v
+}
+
+// Prepend one or more filters before the current ones.
+func (r *TrackFilters) Prepend(filters ...TrackFilter) {
+	src := *r
+	dst := make(TrackFilters, 0, len(filters)+len(src))
+	dst = append(dst, filters...)
+	*r = append(dst, src...)
+}
+
+// combined returns the filters as a single filter.
+func (r TrackFilters) combined() TrackFilter {
+	return func(req *http.Request, resp *http.Response, err error) (*http.Request, *http.Response) {
+		for _, filter := range r {
+			req, resp = filter(req, resp, err)
+		}
+		return req, resp
+	}
+}

govcr.go

@@ -35,6 +35,9 @@ type VCRConfig struct {
 	// Filter to run before a response is returned.
 	ResponseFilters ResponseFilters
 
+	// Filter to run before storing a request/response pair in a track
+	TrackFilters TrackFilters
+
 	// LongPlay will compress data on cassettes.
 	LongPlay         bool
 	DisableRecording bool
@@ -85,6 +88,7 @@ func NewVCR(cassetteName string, vcrConfig *VCRConfig) *VCRControlPanel {
 		Transport:        vcrConfig.Client.Transport,
 		RequestFilter:    vcrConfig.RequestFilters.combined(),
 		ResponseFilter:   vcrConfig.ResponseFilters.combined(),
+		TrackFilter:      vcrConfig.TrackFilters.combined(),
 		Logger:           logger,
 		CassettePath:     vcrConfig.CassettePath,
 	}
@@ -242,6 +246,33 @@ func readRequestBody(req *http.Request) ([]byte, error) {
 	return bodyData, nil
 }
 
+// copyResponse makes a copy an HTTP response.
+// It ensures that the original response Body stream is restored to its original state
+// and can be read from again.
+// TODO: should perform a deep copy of the TLS property as with URL
+func copyResponse(resp *http.Response) (*http.Response, error) {
+	if resp == nil {
+		return nil, nil
+	}
+
+	// get a shallow copy
+	copiedResp := *resp
+
+	copiedResp.Header = cloneHeader(resp.Header)
+
+	// deal with the Body
+	bodyCopy, err := readResponseBody(resp)
+	if err != nil {
+		return nil, err
+	}
+
+	// restore Body stream state
+	resp.Body = toReadCloser(bodyCopy)
+	copiedResp.Body = toReadCloser(bodyCopy)
+
+	return &copiedResp, nil
+}
+
 // readResponseBody reads the Body data stream and restores its states.
 // It ensures the stream is restored to its original state and can be read from again.
 func readResponseBody(resp *http.Response) ([]byte, error) {

pcb.go

@@ -12,6 +12,7 @@ type pcb struct {
 	Transport        http.RoundTripper
 	RequestFilter    RequestFilter
 	ResponseFilter   ResponseFilter
+	TrackFilter      TrackFilter
 	Logger           *log.Logger
 	DisableRecording bool
 	CassettePath     string

vcr_transport.go

@@ -24,6 +24,7 @@ func (t *vcrTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 		t.PCB.Logger.Println(err)
 		return nil, err
 	}
+	copiedReq, _ = t.PCB.TrackFilter(copiedReq, nil, nil)
 
 	// attempt to use a track from the cassette that matches
 	// the request if one exists.
@@ -40,8 +41,14 @@ func (t *vcrTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 
 	if !t.PCB.DisableRecording {
 		// the VCR is not in read-only mode so
+		copiedResp, errResp := copyResponse(resp)
+		if errResp != nil {
+			t.PCB.Logger.Println(errResp)
+			return nil, errResp
+		}
+		copiedReq, copiedResp = t.PCB.TrackFilter(copiedReq, copiedResp, err)
 		t.PCB.Logger.Printf("INFO - Cassette '%s' - Recording new track for %s %s as %s %s\n", t.Cassette.Name, req.Method, req.URL.String(), copiedReq.Method, copiedReq.URL)
-		if err := recordNewTrackToCassette(t.Cassette, copiedReq, resp, err); err != nil {
+		if err := recordNewTrackToCassette(t.Cassette, copiedReq, copiedResp, err); err != nil {
 			t.PCB.Logger.Println(err)
 		}
 	}