export_test.go

     1  // Copyright 2011 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  // Bridge package to expose http internals to tests in the http_test
     6  // package.
     7  
     8  package http
     9  
    10  import (
    11  	"context"
    12  	"fmt"
    13  	"net"
    14  	"net/url"
    15  	"slices"
    16  	"sync"
    17  	"testing"
    18  	"time"
    19  )
    20  
    21  var (
    22  	DefaultUserAgent                  = defaultUserAgent
    23  	NewLoggingConn                    = newLoggingConn
    24  	ExportRefererForURL               = refererForURL
    25  	ExportServerNewConn               = (*Server).newConn
    26  	ExportCloseWriteAndWait           = (*conn).closeWriteAndWait
    27  	ExportErrRequestCanceled          = errRequestCanceled
    28  	ExportErrRequestCanceledConn      = errRequestCanceledConn
    29  	ExportErrServerClosedIdle         = errServerClosedIdle
    30  	ExportServeFile                   = serveFile
    31  	ExportScanETag                    = scanETag
    32  	ExportHttp2ConfigureServer        = http2ConfigureServer
    33  	Export_shouldCopyHeaderOnRedirect = shouldCopyHeaderOnRedirect
    34  	Export_writeStatusLine            = writeStatusLine
    35  	Export_is408Message               = is408Message
    36  )
    37  
    38  var MaxWriteWaitBeforeConnReuse = &maxWriteWaitBeforeConnReuse
    39  
    40  func init() {
    41  	// We only want to pay for this cost during testing.
    42  	// When not under test, these values are always nil
    43  	// and never assigned to.
    44  	testHookMu = new(sync.Mutex)
    45  
    46  	testHookClientDoResult = func(res *Response, err error) {
    47  		if err != nil {
    48  			if _, ok := err.(*url.Error); !ok {
    49  				panic(fmt.Sprintf("unexpected Client.Do error of type %T; want *url.Error", err))
    50  			}
    51  		} else {
    52  			if res == nil {
    53  				panic("Client.Do returned nil, nil")
    54  			}
    55  			if res.Body == nil {
    56  				panic("Client.Do returned nil res.Body and no error")
    57  			}
    58  		}
    59  	}
    60  }
    61  
    62  func CondSkipHTTP2(t testing.TB) {
    63  	if omitBundledHTTP2 {
    64  		t.Skip("skipping HTTP/2 test when nethttpomithttp2 build tag in use")
    65  	}
    66  }
    67  
    68  var (
    69  	SetEnterRoundTripHook = hookSetter(&testHookEnterRoundTrip)
    70  	SetRoundTripRetried   = hookSetter(&testHookRoundTripRetried)
    71  )
    72  
    73  func SetReadLoopBeforeNextReadHook(f func()) {
    74  	unnilTestHook(&f)
    75  	testHookReadLoopBeforeNextRead = f
    76  }
    77  
    78  // SetPendingDialHooks sets the hooks that run before and after handling
    79  // pending dials.
    80  func SetPendingDialHooks(before, after func()) {
    81  	unnilTestHook(&before)
    82  	unnilTestHook(&after)
    83  	testHookPrePendingDial, testHookPostPendingDial = before, after
    84  }
    85  
    86  func SetTestHookServerServe(fn func(*Server, net.Listener)) { testHookServerServe = fn }
    87  
    88  func SetTestHookProxyConnectTimeout(t *testing.T, f func(context.Context, time.Duration) (context.Context, context.CancelFunc)) {
    89  	orig := testHookProxyConnectTimeout
    90  	t.Cleanup(func() {
    91  		testHookProxyConnectTimeout = orig
    92  	})
    93  	testHookProxyConnectTimeout = f
    94  }
    95  
    96  func NewTestTimeoutHandler(handler Handler, ctx context.Context) Handler {
    97  	return &timeoutHandler{
    98  		handler:     handler,
    99  		testContext: ctx,
   100  		// (no body)
   101  	}
   102  }
   103  
   104  func ResetCachedEnvironment() {
   105  	resetProxyConfig()
   106  }
   107  
   108  func (t *Transport) NumPendingRequestsForTesting() int {
   109  	t.reqMu.Lock()
   110  	defer t.reqMu.Unlock()
   111  	return len(t.reqCanceler)
   112  }
   113  
   114  func (t *Transport) IdleConnKeysForTesting() (keys []string) {
   115  	keys = make([]string, 0)
   116  	t.idleMu.Lock()
   117  	defer t.idleMu.Unlock()
   118  	for key := range t.idleConn {
   119  		keys = append(keys, key.String())
   120  	}
   121  	slices.Sort(keys)
   122  	return
   123  }
   124  
   125  func (t *Transport) IdleConnKeyCountForTesting() int {
   126  	t.idleMu.Lock()
   127  	defer t.idleMu.Unlock()
   128  	return len(t.idleConn)
   129  }
   130  
   131  func (t *Transport) IdleConnStrsForTesting() []string {
   132  	var ret []string
   133  	t.idleMu.Lock()
   134  	defer t.idleMu.Unlock()
   135  	for _, conns := range t.idleConn {
   136  		for _, pc := range conns {
   137  			ret = append(ret, pc.conn.LocalAddr().String()+"/"+pc.conn.RemoteAddr().String())
   138  		}
   139  	}
   140  	slices.Sort(ret)
   141  	return ret
   142  }
   143  
   144  func (t *Transport) IdleConnStrsForTesting_h2() []string {
   145  	var ret []string
   146  	noDialPool := t.h2transport.(*http2Transport).ConnPool.(http2noDialClientConnPool)
   147  	pool := noDialPool.http2clientConnPool
   148  
   149  	pool.mu.Lock()
   150  	defer pool.mu.Unlock()
   151  
   152  	for k, ccs := range pool.conns {
   153  		for _, cc := range ccs {
   154  			if cc.idleState().canTakeNewRequest {
   155  				ret = append(ret, k)
   156  			}
   157  		}
   158  	}
   159  
   160  	slices.Sort(ret)
   161  	return ret
   162  }
   163  
   164  func (t *Transport) IdleConnCountForTesting(scheme, addr string) int {
   165  	t.idleMu.Lock()
   166  	defer t.idleMu.Unlock()
   167  	key := connectMethodKey{"", scheme, addr, false}
   168  	cacheKey := key.String()
   169  	for k, conns := range t.idleConn {
   170  		if k.String() == cacheKey {
   171  			return len(conns)
   172  		}
   173  	}
   174  	return 0
   175  }
   176  
   177  func (t *Transport) IdleConnWaitMapSizeForTesting() int {
   178  	t.idleMu.Lock()
   179  	defer t.idleMu.Unlock()
   180  	return len(t.idleConnWait)
   181  }
   182  
   183  func (t *Transport) IsIdleForTesting() bool {
   184  	t.idleMu.Lock()
   185  	defer t.idleMu.Unlock()
   186  	return t.closeIdle
   187  }
   188  
   189  func (t *Transport) QueueForIdleConnForTesting() {
   190  	t.queueForIdleConn(nil)
   191  }
   192  
   193  // PutIdleTestConn reports whether it was able to insert a fresh
   194  // persistConn for scheme, addr into the idle connection pool.
   195  func (t *Transport) PutIdleTestConn(scheme, addr string) bool {
   196  	c, _ := net.Pipe()
   197  	key := connectMethodKey{"", scheme, addr, false}
   198  
   199  	if t.MaxConnsPerHost > 0 {
   200  		// Transport is tracking conns-per-host.
   201  		// Increment connection count to account
   202  		// for new persistConn created below.
   203  		t.connsPerHostMu.Lock()
   204  		if t.connsPerHost == nil {
   205  			t.connsPerHost = make(map[connectMethodKey]int)
   206  		}
   207  		t.connsPerHost[key]++
   208  		t.connsPerHostMu.Unlock()
   209  	}
   210  
   211  	return t.tryPutIdleConn(&persistConn{
   212  		t:        t,
   213  		conn:     c,                   // dummy
   214  		closech:  make(chan struct{}), // so it can be closed
   215  		cacheKey: key,
   216  	}) == nil
   217  }
   218  
   219  // PutIdleTestConnH2 reports whether it was able to insert a fresh
   220  // HTTP/2 persistConn for scheme, addr into the idle connection pool.
   221  func (t *Transport) PutIdleTestConnH2(scheme, addr string, alt RoundTripper) bool {
   222  	key := connectMethodKey{"", scheme, addr, false}
   223  
   224  	if t.MaxConnsPerHost > 0 {
   225  		// Transport is tracking conns-per-host.
   226  		// Increment connection count to account
   227  		// for new persistConn created below.
   228  		t.connsPerHostMu.Lock()
   229  		if t.connsPerHost == nil {
   230  			t.connsPerHost = make(map[connectMethodKey]int)
   231  		}
   232  		t.connsPerHost[key]++
   233  		t.connsPerHostMu.Unlock()
   234  	}
   235  
   236  	return t.tryPutIdleConn(&persistConn{
   237  		t:        t,
   238  		alt:      alt,
   239  		cacheKey: key,
   240  	}) == nil
   241  }
   242  
   243  // All test hooks must be non-nil so they can be called directly,
   244  // but the tests use nil to mean hook disabled.
   245  func unnilTestHook(f *func()) {
   246  	if *f == nil {
   247  		*f = nop
   248  	}
   249  }
   250  
   251  func hookSetter(dst *func()) func(func()) {
   252  	return func(fn func()) {
   253  		unnilTestHook(&fn)
   254  		*dst = fn
   255  	}
   256  }
   257  
   258  func ExportHttp2ConfigureTransport(t *Transport) error {
   259  	t2, err := http2configureTransports(t)
   260  	if err != nil {
   261  		return err
   262  	}
   263  	t.h2transport = t2
   264  	return nil
   265  }
   266  
   267  func (s *Server) ExportAllConnsIdle() bool {
   268  	s.mu.Lock()
   269  	defer s.mu.Unlock()
   270  	for c := range s.activeConn {
   271  		st, unixSec := c.getState()
   272  		if unixSec == 0 || st != StateIdle {
   273  			return false
   274  		}
   275  	}
   276  	return true
   277  }
   278  
   279  func (s *Server) ExportAllConnsByState() map[ConnState]int {
   280  	states := map[ConnState]int{}
   281  	s.mu.Lock()
   282  	defer s.mu.Unlock()
   283  	for c := range s.activeConn {
   284  		st, _ := c.getState()
   285  		states[st] += 1
   286  	}
   287  	return states
   288  }
   289  
   290  func (r *Request) WithT(t *testing.T) *Request {
   291  	return r.WithContext(context.WithValue(r.Context(), tLogKey{}, t.Logf))
   292  }
   293  
   294  func ExportSetH2GoawayTimeout(d time.Duration) (restore func()) {
   295  	old := http2goAwayTimeout
   296  	http2goAwayTimeout = d
   297  	return func() { http2goAwayTimeout = old }
   298  }
   299  
   300  func (r *Request) ExportIsReplayable() bool { return r.isReplayable() }
   301  
   302  // ExportCloseTransportConnsAbruptly closes all idle connections from
   303  // tr in an abrupt way, just reaching into the underlying Conns and
   304  // closing them, without telling the Transport or its persistConns
   305  // that it's doing so. This is to simulate the server closing connections
   306  // on the Transport.
   307  func ExportCloseTransportConnsAbruptly(tr *Transport) {
   308  	tr.idleMu.Lock()
   309  	for _, pcs := range tr.idleConn {
   310  		for _, pc := range pcs {
   311  			pc.conn.Close()
   312  		}
   313  	}
   314  	tr.idleMu.Unlock()
   315  }
   316  
   317  // ResponseWriterConnForTesting returns w's underlying connection, if w
   318  // is a regular *response ResponseWriter.
   319  func ResponseWriterConnForTesting(w ResponseWriter) (c net.Conn, ok bool) {
   320  	if r, ok := w.(*response); ok {
   321  		return r.conn.rwc, true
   322  	}
   323  	return nil, false
   324  }
   325  
   326  func init() {
   327  	// Set the default rstAvoidanceDelay to the minimum possible value to shake
   328  	// out tests that unexpectedly depend on it. Such tests should use
   329  	// runTimeSensitiveTest and SetRSTAvoidanceDelay to explicitly raise the delay
   330  	// if needed.
   331  	rstAvoidanceDelay = 1 * time.Nanosecond
   332  }
   333  
   334  // SetRSTAvoidanceDelay sets how long we are willing to wait between calling
   335  // CloseWrite on a connection and fully closing the connection.
   336  func SetRSTAvoidanceDelay(t *testing.T, d time.Duration) {
   337  	prevDelay := rstAvoidanceDelay
   338  	t.Cleanup(func() {
   339  		rstAvoidanceDelay = prevDelay
   340  	})
   341  	rstAvoidanceDelay = d
   342  }
   343
View as plain text