github.com/carter-ya/go-ethereum@v0.0.0-20230628080049-d2309be3983b/cmd/geth/attach_test.go (about)

     1  // Copyright 2022 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package main
    18  
    19  import (
    20  	"fmt"
    21  	"net"
    22  	"net/http"
    23  	"sync/atomic"
    24  	"testing"
    25  )
    26  
    27  type testHandler struct {
    28  	body func(http.ResponseWriter, *http.Request)
    29  }
    30  
    31  func (t *testHandler) ServeHTTP(out http.ResponseWriter, in *http.Request) {
    32  	t.body(out, in)
    33  }
    34  
    35  // TestAttachWithHeaders tests that 'geth attach' with custom headers works, i.e
    36  // that custom headers are forwarded to the target.
    37  func TestAttachWithHeaders(t *testing.T) {
    38  	t.Parallel()
    39  	ln, err := net.Listen("tcp", "localhost:0")
    40  	if err != nil {
    41  		t.Fatal(err)
    42  	}
    43  	port := ln.Addr().(*net.TCPAddr).Port
    44  	testReceiveHeaders(t, ln, "attach", "-H", "first: one", "-H", "second: two", fmt.Sprintf("http://localhost:%d", port))
    45  	// This way to do it fails due to flag ordering:
    46  	//
    47  	// testReceiveHeaders(t, ln, "-H", "first: one", "-H", "second: two", "attach", fmt.Sprintf("http://localhost:%d", port))
    48  	// This is fixed in a follow-up PR.
    49  }
    50  
    51  // TestAttachWithHeaders tests that 'geth db --remotedb' with custom headers works, i.e
    52  // that custom headers are forwarded to the target.
    53  func TestRemoteDbWithHeaders(t *testing.T) {
    54  	t.Parallel()
    55  	ln, err := net.Listen("tcp", "localhost:0")
    56  	if err != nil {
    57  		t.Fatal(err)
    58  	}
    59  	port := ln.Addr().(*net.TCPAddr).Port
    60  	testReceiveHeaders(t, ln, "db", "metadata", "--remotedb", fmt.Sprintf("http://localhost:%d", port), "-H", "first: one", "-H", "second: two")
    61  }
    62  
    63  func testReceiveHeaders(t *testing.T, ln net.Listener, gethArgs ...string) {
    64  	var ok uint32
    65  	server := &http.Server{
    66  		Addr: "localhost:0",
    67  		Handler: &testHandler{func(w http.ResponseWriter, r *http.Request) {
    68  			// We expect two headers
    69  			if have, want := r.Header.Get("first"), "one"; have != want {
    70  				t.Fatalf("missing header, have %v want %v", have, want)
    71  			}
    72  			if have, want := r.Header.Get("second"), "two"; have != want {
    73  				t.Fatalf("missing header, have %v want %v", have, want)
    74  			}
    75  			atomic.StoreUint32(&ok, 1)
    76  		}}}
    77  	go server.Serve(ln)
    78  	defer server.Close()
    79  	runGeth(t, gethArgs...).WaitExit()
    80  	if atomic.LoadUint32(&ok) != 1 {
    81  		t.Fatal("Test fail, expected invocation to succeed")
    82  	}
    83  }