golang.org/x/build@v0.0.0-20240506185731-218518f32b70/cmd/buildlet/reverse_test.go (about)

     1  // Copyright 2021 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  // Same requirements as internal/coordinator/pool/reverse.go.
     6  //go:build linux || darwin
     7  
     8  package main
     9  
    10  import (
    11  	"crypto/tls"
    12  	"fmt"
    13  	"net"
    14  	"net/http"
    15  	"testing"
    16  	"time"
    17  
    18  	"golang.org/x/build"
    19  	"golang.org/x/build/internal/coordinator/pool"
    20  	"golang.org/x/build/revdial/v2"
    21  )
    22  
    23  // coordinatorServer creates a server and listener for the coordinator side of
    24  // revdial. They should be closed when done.
    25  func coordinatorServer() (*http.Server, net.Listener, error) {
    26  	mux := http.NewServeMux()
    27  	mux.HandleFunc("/reverse", pool.HandleReverse)
    28  	mux.Handle("/revdial", revdial.ConnHandler())
    29  
    30  	ln, err := net.Listen("tcp", "")
    31  	if err != nil {
    32  		return nil, nil, fmt.Errorf(`net.Listen(":"): %v`, err)
    33  	}
    34  
    35  	cert, err := tls.X509KeyPair([]byte(build.DevCoordinatorCA), []byte(build.DevCoordinatorKey))
    36  	if err != nil {
    37  		return nil, nil, fmt.Errorf("error creating TLS cert: %v", err)
    38  	}
    39  
    40  	ln = tls.NewListener(ln, &tls.Config{
    41  		Certificates: []tls.Certificate{cert},
    42  	})
    43  
    44  	addr := ln.Addr().String()
    45  	srv := &http.Server{
    46  		Addr:    addr,
    47  		Handler: mux,
    48  	}
    49  	return srv, ln, nil
    50  }
    51  
    52  // testReverseDial verifies that a revdial connection can be established and
    53  // registered in the coordinator reverse pool at coordAddr.
    54  func testReverseDial(t *testing.T, coordAddr, hostType string) {
    55  	t.Helper()
    56  
    57  	oldCoordinator := *coordinator
    58  	defer func() {
    59  		*coordinator = oldCoordinator
    60  	}()
    61  	*coordinator = coordAddr
    62  
    63  	// N.B. We don't need to set *hostname to anything in particular as it
    64  	// is only advisory in the coordinator. It is not used to connect back
    65  	// to reverse buildlets.
    66  
    67  	oldReverseType := *reverseType
    68  	defer func() {
    69  		*reverseType = oldReverseType
    70  	}()
    71  	*reverseType = hostType
    72  
    73  	ln, err := dialCoordinator()
    74  	if err != nil {
    75  		t.Fatalf("dialCoordinator got err %v want nil", err)
    76  	}
    77  
    78  	mux := http.NewServeMux()
    79  	mux.HandleFunc("/status", handleStatus)
    80  	srv := &http.Server{
    81  		Handler: mux,
    82  	}
    83  	c := make(chan error, 1)
    84  	go func() {
    85  		c <- srv.Serve(ln)
    86  	}()
    87  	defer func() {
    88  		srv.Close()
    89  		err := <-c
    90  		if err != http.ErrServerClosed {
    91  			t.Errorf("Server shutdown got err %v want ErrServerClosed", err)
    92  		}
    93  	}()
    94  
    95  	// Verify that we eventually get the "buildlet" registered with the pool.
    96  	tick := time.NewTicker(10 * time.Millisecond)
    97  	defer tick.Stop()
    98  	start := time.Now()
    99  	for range tick.C {
   100  		if time.Since(start) > 1*time.Second {
   101  			t.Fatalf("Buildlet failed to register within 1s.")
   102  		}
   103  
   104  		types := pool.ReversePool().HostTypes()
   105  		for _, typ := range types {
   106  			if typ == hostType {
   107  				// Success!
   108  				return
   109  			}
   110  		}
   111  	}
   112  }
   113  
   114  // TestReverseDial verifies that a revdial connection can be established and
   115  // registered in the coordinator reverse pool.
   116  func TestReverseDial(t *testing.T) {
   117  	pool.SetBuilderMasterKey([]byte(devMasterKey))
   118  
   119  	srv, ln, err := coordinatorServer()
   120  	if err != nil {
   121  		t.Fatalf("serveCoordinator got err %v want nil", err)
   122  	}
   123  	go srv.Serve(ln)
   124  	defer srv.Close()
   125  
   126  	const hostType = "test-reverse-dial"
   127  	testReverseDial(t, srv.Addr, hostType)
   128  }
   129  
   130  // TestReverseDialRedirect verifies that a revdial connection works with a 307
   131  // redirect to the endpoints. The coordinator will do this in dev mode.
   132  func TestReverseDialRedirect(t *testing.T) {
   133  	pool.SetBuilderMasterKey([]byte(devMasterKey))
   134  
   135  	srv, ln, err := coordinatorServer()
   136  	if err != nil {
   137  		t.Fatalf("serveCoordinator got err %v want nil", err)
   138  	}
   139  
   140  	mux := http.NewServeMux()
   141  	mux.HandleFunc("/redirected/reverse", pool.HandleReverse)
   142  	mux.Handle("/redirected/revdial", revdial.ConnHandler())
   143  
   144  	redirect := func(w http.ResponseWriter, r *http.Request) {
   145  		u := *r.URL
   146  		u.Path = "/redirected/" + u.Path
   147  		http.Redirect(w, r, u.String(), http.StatusTemporaryRedirect)
   148  	}
   149  	mux.HandleFunc("/reverse", redirect)
   150  	mux.HandleFunc("/revdial", redirect)
   151  	srv.Handler = mux
   152  
   153  	go srv.Serve(ln)
   154  	defer srv.Close()
   155  
   156  	const hostType = "test-reverse-dial-redirect"
   157  	testReverseDial(t, srv.Addr, hostType)
   158  }