intercept.go

     1  package intercept
     2  
     3  import (
     4  	"errors"
     5  	"net/http"
     6  	"net/url"
     7  )
     8  
     9  // Interceptor is used to change the host, and maybe the client,
    10  // for a request to point to a test host.
    11  type Interceptor struct {
    12  	Scheme   string
    13  	FromHost string
    14  	ToHost   string
    15  	Client   *http.Client
    16  }
    17  
    18  // EnableTestHooks installs the given interceptors to be used by URL and Request.
    19  func EnableTestHooks(interceptors []Interceptor) error {
    20  	if TestHooksEnabled {
    21  		return errors.New("web: test hooks already enabled")
    22  	}
    23  
    24  	for _, t := range interceptors {
    25  		if t.FromHost == "" {
    26  			panic("EnableTestHooks: missing FromHost")
    27  		}
    28  		if t.ToHost == "" {
    29  			panic("EnableTestHooks: missing ToHost")
    30  		}
    31  	}
    32  
    33  	testInterceptors = interceptors
    34  	TestHooksEnabled = true
    35  	return nil
    36  }
    37  
    38  // DisableTestHooks disables the installed interceptors.
    39  func DisableTestHooks() {
    40  	if !TestHooksEnabled {
    41  		panic("web: test hooks not enabled")
    42  	}
    43  	TestHooksEnabled = false
    44  	testInterceptors = nil
    45  }
    46  
    47  var (
    48  	// TestHooksEnabled is true if interceptors are installed
    49  	TestHooksEnabled = false
    50  	testInterceptors []Interceptor
    51  )
    52  
    53  // URL returns the Interceptor to be used for a given URL.
    54  func URL(u *url.URL) (*Interceptor, bool) {
    55  	if !TestHooksEnabled {
    56  		return nil, false
    57  	}
    58  	for i, t := range testInterceptors {
    59  		if u.Host == t.FromHost && (u.Scheme == "" || u.Scheme == t.Scheme) {
    60  			return &testInterceptors[i], true
    61  		}
    62  	}
    63  	return nil, false
    64  }
    65  
    66  // Request updates the host to actually use for the request, if it is to be intercepted.
    67  func Request(req *http.Request) {
    68  	if t, ok := URL(req.URL); ok {
    69  		req.Host = req.URL.Host
    70  		req.URL.Host = t.ToHost
    71  	}
    72  }
    73
View as plain text