intercept.go

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