github.com/google/go-safeweb@v0.0.0-20231219055052-64d8cfc90fbb/safehttp/safehttptest/recorder.go (about)

     1  // Copyright 2020 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  //	https://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 safehttptest
    16  
    17  import (
    18  	"net/http"
    19  	"net/http/httptest"
    20  
    21  	"github.com/google/go-safeweb/safehttp"
    22  )
    23  
    24  // FakeResponseWriter creates a fake safehttp.ResponseWriter implementation.
    25  //
    26  // It performs no error checking nor runs interceptors.
    27  type FakeResponseWriter struct {
    28  	// The Dispatcher implementation
    29  	Dispatcher *FakeDispatcher
    30  
    31  	// ResponseWriter is only used for calls to Dispatcher. Calls to AddCookie()
    32  	// do not affect it.
    33  	ResponseWriter http.ResponseWriter
    34  
    35  	// Cookies coming from AddCookie() calls. Use the ResponseWriter to see what
    36  	// cookies have been set by the Dispatcher.
    37  	Cookies []*safehttp.Cookie
    38  
    39  	// Response headers.
    40  	Headers safehttp.Header
    41  }
    42  
    43  // FakeDispatcher provides a minimal implementation of the Dispatcher to be used for testing Interceptors.
    44  type FakeDispatcher struct {
    45  	Written    safehttp.Response
    46  	Dispatcher safehttp.Dispatcher
    47  }
    48  
    49  // Write records the written responses and calls Dispatcher.Write.
    50  // If Dispatcher is nil, the default one is used.
    51  func (d *FakeDispatcher) Write(rw http.ResponseWriter, resp safehttp.Response) error {
    52  	d.Written = resp
    53  	if d.Dispatcher != nil {
    54  		return d.Dispatcher.Write(rw, resp)
    55  	}
    56  	return safehttp.DefaultDispatcher{}.Write(rw, resp)
    57  }
    58  
    59  // Error writes just the status code.
    60  func (d *FakeDispatcher) Error(rw http.ResponseWriter, resp safehttp.ErrorResponse) error {
    61  	rw.WriteHeader(int(resp.Code()))
    62  	return nil
    63  }
    64  
    65  // NewFakeResponseWriter creates a new safehttp.ResponseWriter implementation
    66  // and a httptest.ResponseRecorder, for testing purposes only.
    67  func NewFakeResponseWriter() (*FakeResponseWriter, *httptest.ResponseRecorder) {
    68  	recorder := httptest.NewRecorder()
    69  	return &FakeResponseWriter{
    70  		Dispatcher:     &FakeDispatcher{},
    71  		ResponseWriter: recorder,
    72  		Headers:        safehttp.NewHeader(recorder.HeaderMap),
    73  	}, recorder
    74  }
    75  
    76  var _ safehttp.ResponseWriter = (*FakeResponseWriter)(nil)
    77  
    78  // Header returns the Header.
    79  func (frw *FakeResponseWriter) Header() safehttp.Header {
    80  	return frw.Headers
    81  }
    82  
    83  // AddCookie appends the given cookie to the Cookies field.
    84  func (frw *FakeResponseWriter) AddCookie(c *safehttp.Cookie) error {
    85  	if len(c.Name()) == 0 {
    86  		panic("empty cookie name")
    87  	}
    88  
    89  	frw.Cookies = append(frw.Cookies, c)
    90  	return nil
    91  }
    92  
    93  // Write forwards the response to Dispatcher.Write.
    94  func (frw *FakeResponseWriter) Write(resp safehttp.Response) safehttp.Result {
    95  	if err := frw.Dispatcher.Write(frw.ResponseWriter, resp); err != nil {
    96  		panic(err)
    97  	}
    98  	return safehttp.Result{}
    99  }
   100  
   101  // NoContent writes just the NoContent status code.
   102  func (frw *FakeResponseWriter) NoContent() safehttp.Result {
   103  	frw.ResponseWriter.WriteHeader(int(safehttp.StatusNoContent))
   104  	return safehttp.Result{}
   105  }
   106  
   107  // WriteError forwards the error response to Dispatcher.WriteError.
   108  func (frw *FakeResponseWriter) WriteError(resp safehttp.ErrorResponse) safehttp.Result {
   109  	if err := frw.Dispatcher.Error(frw.ResponseWriter, resp); err != nil {
   110  		panic(err)
   111  	}
   112  	return safehttp.Result{}
   113  }