http_recorder.go

Documentation: github.com/GoogleCloudPlatform/k8s-config-connector/pkg/test

     1  // Copyright 2022 Google LLC
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package test
    16  
    17  import (
    18  	"bytes"
    19  	"fmt"
    20  	"io/ioutil"
    21  	"net/http"
    22  	"os"
    23  	"path/filepath"
    24  	"strings"
    25  	"sync"
    26  	"time"
    27  	"unicode"
    28  
    29  	"gopkg.in/yaml.v2"
    30  	"k8s.io/klog/v2"
    31  )
    32  
    33  type LogEntry struct {
    34  	Timestamp time.Time `json:"timestamp,omitempty"`
    35  	Request   Request   `json:"request,omitempty"`
    36  	Response  Response  `json:"response,omitempty"`
    37  	Error     string    `json:"error,omitempty"`
    38  }
    39  
    40  type Request struct {
    41  	Method string      `json:"method,omitempty"`
    42  	URL    string      `json:"url,omitempty"`
    43  	Header http.Header `json:"header,omitempty"`
    44  	Body   string      `json:"body,omitempty"`
    45  }
    46  
    47  type Response struct {
    48  	Status     string      `json:"status,omitempty"`
    49  	StatusCode int         `json:"statusCode,omitempty"`
    50  	Header     http.Header `json:"header,omitempty"`
    51  	Body       string      `json:"body,omitempty"`
    52  }
    53  
    54  type HTTPRecorder struct {
    55  	outputDir string
    56  	inner     http.RoundTripper
    57  
    58  	// mutex to avoid concurrent writes to the same file
    59  	mutex sync.Mutex
    60  }
    61  
    62  func NewHTTPRecorder(inner http.RoundTripper, outputDir string) *HTTPRecorder {
    63  	rt := &HTTPRecorder{outputDir: outputDir, inner: inner}
    64  	return rt
    65  }
    66  
    67  func (r *HTTPRecorder) RoundTrip(req *http.Request) (*http.Response, error) {
    68  	var entry LogEntry
    69  	entry.Timestamp = time.Now()
    70  	entry.Request.Method = req.Method
    71  	entry.Request.URL = req.URL.String()
    72  
    73  	entry.Request.Header = make(http.Header)
    74  	for k, values := range req.Header {
    75  		switch strings.ToLower(k) {
    76  		case "authorization":
    77  			entry.Request.Header[k] = []string{"(removed)"}
    78  		default:
    79  			entry.Request.Header[k] = values
    80  		}
    81  	}
    82  
    83  	if req.Body != nil {
    84  		requestBody, err := ioutil.ReadAll(req.Body)
    85  		if err != nil {
    86  			panic("failed to read request body")
    87  		}
    88  		entry.Request.Body = string(requestBody)
    89  		req.Body = ioutil.NopCloser(bytes.NewReader(requestBody))
    90  	}
    91  
    92  	response, err := r.inner.RoundTrip(req)
    93  
    94  	if err != nil {
    95  		entry.Error = fmt.Sprintf("%v", err)
    96  	}
    97  
    98  	if recordErr := r.record(&entry, req, response); recordErr != nil {
    99  		klog.Warningf("failed to record HTTP request: %v", err)
   100  	}
   101  
   102  	return response, err
   103  }
   104  
   105  func (r *HTTPRecorder) record(entry *LogEntry, req *http.Request, resp *http.Response) error {
   106  	if resp != nil {
   107  		entry.Response.Status = resp.Status
   108  		entry.Response.StatusCode = resp.StatusCode
   109  
   110  		entry.Response.Header = make(http.Header)
   111  		for k, values := range resp.Header {
   112  			switch strings.ToLower(k) {
   113  			case "authorization":
   114  				entry.Response.Header[k] = []string{"(removed)"}
   115  			default:
   116  				entry.Response.Header[k] = values
   117  			}
   118  		}
   119  
   120  		if resp.Body != nil {
   121  			requestBody, err := ioutil.ReadAll(resp.Body)
   122  			if err != nil {
   123  				panic("failed to read response body")
   124  			}
   125  			entry.Response.Body = string(requestBody)
   126  			resp.Body = ioutil.NopCloser(bytes.NewReader(requestBody))
   127  		}
   128  	}
   129  
   130  	ctx := req.Context()
   131  	t := TestFromContext(ctx)
   132  	testName := "unknown"
   133  	if t != nil {
   134  		testName = t.Name()
   135  	}
   136  	dirName := sanitizePath(testName)
   137  	p := filepath.Join(r.outputDir, dirName, "requests.log")
   138  
   139  	b, err := yaml.Marshal(entry)
   140  	if err != nil {
   141  		return fmt.Errorf("failed to marshal data: %w", err)
   142  	}
   143  
   144  	// Just in case we are writing to the same file concurrently
   145  	r.mutex.Lock()
   146  	defer r.mutex.Unlock()
   147  
   148  	if err := os.MkdirAll(filepath.Dir(p), 0755); err != nil {
   149  		return fmt.Errorf("failed to create directory %q: %w", filepath.Dir(p), err)
   150  	}
   151  	f, err := os.OpenFile(p, os.O_CREATE|os.O_APPEND|os.O_RDWR, 0644)
   152  	if err != nil {
   153  		return fmt.Errorf("failed to open file %q: %w", p, err)
   154  	}
   155  	defer f.Close()
   156  
   157  	if _, err := f.Write(b); err != nil {
   158  		return fmt.Errorf("failed to write to file %q: %w", p, err)
   159  	}
   160  	delimeter := "\n\n---\n\n"
   161  	if _, err := f.Write([]byte(delimeter)); err != nil {
   162  		return fmt.Errorf("failed to write to file %q: %w", p, err)
   163  	}
   164  
   165  	return nil
   166  }
   167  
   168  func sanitizePath(s string) string {
   169  	var out strings.Builder
   170  	for _, r := range s {
   171  		if unicode.IsLetter(r) || unicode.IsDigit(r) {
   172  			out.WriteRune(r)
   173  		} else {
   174  			out.WriteRune('_')
   175  		}
   176  	}
   177  	return out.String()
   178  }
   179
View as plain text