github.com/grafana/pyroscope@v1.18.0/pkg/testhelper/http.go (about)

     1  // Copyright 2021-2022 Buf Technologies, Inc.
     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  // This has been adapted from https://connectrpc.com/connect/blob/cce7065d23ae00021eb4b31284361a2d8525df21/example_init_test.go#L44-L144
    16  
    17  package testhelper
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"net"
    23  	"net/http"
    24  	"net/http/httptest"
    25  	"sync"
    26  )
    27  
    28  // InMemoryServer is an HTTP server that uses in-memory pipes instead of TCP.
    29  // It supports HTTP/2 and has TLS enabled.
    30  //
    31  // The Go Playground panics if we try to start a TCP-backed server. If you're
    32  // not familiar with the Playground's behavior, it looks like our examples are
    33  // broken. This server lets us write examples that work in the playground
    34  // without abstracting over HTTP.
    35  type InMemoryServer struct {
    36  	server   *httptest.Server
    37  	listener *memoryListener
    38  }
    39  
    40  // NewInMemoryServer constructs and starts an inMemoryServer.
    41  func NewInMemoryServer(handler http.Handler) *InMemoryServer {
    42  	lis := &memoryListener{
    43  		conns:  make(chan net.Conn),
    44  		closed: make(chan struct{}),
    45  	}
    46  	server := httptest.NewUnstartedServer(handler)
    47  	server.Listener = lis
    48  	server.EnableHTTP2 = true
    49  	server.StartTLS()
    50  	return &InMemoryServer{
    51  		server:   server,
    52  		listener: lis,
    53  	}
    54  }
    55  
    56  // Client returns an HTTP client configured to trust the server's TLS
    57  // certificate and use HTTP/2 over an in-memory pipe. Automatic HTTP-level gzip
    58  // compression is disabled. It closes its idle connections when the server is
    59  // closed.
    60  func (s *InMemoryServer) Client() *http.Client {
    61  	client := s.server.Client()
    62  	if transport, ok := client.Transport.(*http.Transport); ok {
    63  		transport.DialContext = s.listener.DialContext
    64  		transport.DisableCompression = true
    65  	}
    66  	return client
    67  }
    68  
    69  // URL is the server's URL.
    70  func (s *InMemoryServer) URL() string {
    71  	return s.server.URL
    72  }
    73  
    74  // Close shuts down the server, blocking until all outstanding requests have
    75  // completed.
    76  func (s *InMemoryServer) Close() {
    77  	s.server.Close()
    78  }
    79  
    80  type memoryListener struct {
    81  	conns  chan net.Conn
    82  	once   sync.Once
    83  	closed chan struct{}
    84  }
    85  
    86  // Accept implements net.Listener.
    87  func (l *memoryListener) Accept() (net.Conn, error) {
    88  	select {
    89  	case conn := <-l.conns:
    90  		return conn, nil
    91  	case <-l.closed:
    92  		return nil, errors.New("listener closed")
    93  	}
    94  }
    95  
    96  // Close implements net.Listener.
    97  func (l *memoryListener) Close() error {
    98  	l.once.Do(func() {
    99  		close(l.closed)
   100  	})
   101  	return nil
   102  }
   103  
   104  // Addr implements net.Listener.
   105  func (l *memoryListener) Addr() net.Addr {
   106  	return &memoryAddr{}
   107  }
   108  
   109  // DialContext is the type expected by http.Transport.DialContext.
   110  func (l *memoryListener) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
   111  	select {
   112  	case <-l.closed:
   113  		return nil, errors.New("listener closed")
   114  	default:
   115  	}
   116  	server, client := net.Pipe()
   117  	l.conns <- server
   118  	return client, nil
   119  }
   120  
   121  type memoryAddr struct{}
   122  
   123  // Network implements net.Addr.
   124  func (*memoryAddr) Network() string { return "memory" }
   125  
   126  // String implements io.Stringer, returning a value that matches the
   127  // certificates used by net/http/httptest.
   128  func (*memoryAddr) String() string { return "example.com" }