From 7c6c6e4e2d7ed75ba64240aa071c51951f48fff3 Mon Sep 17 00:00:00 2001 From: DarrylWong Date: Fri, 5 Jan 2024 16:36:24 -0500 Subject: [PATCH] roachtest: add helpers to extract sessionID and create default HTTPClient These helpers are for roachtests to easily create a HTTPClient with appropriate cookies. It also sets InsecureSkipVerify to true as any certs we create won't be valid anyway. This change is so that we can create http requests to secure clusters. Release note: None --- pkg/cmd/roachtest/clusterstats/BUILD.bazel | 1 + pkg/cmd/roachtest/roachtestutil/BUILD.bazel | 2 + pkg/cmd/roachtest/roachtestutil/httpclient.go | 213 ++++++++++++++++++ pkg/roachprod/install/cockroach.go | 2 - 4 files changed, 216 insertions(+), 2 deletions(-) create mode 100644 pkg/cmd/roachtest/roachtestutil/httpclient.go diff --git a/pkg/cmd/roachtest/clusterstats/BUILD.bazel b/pkg/cmd/roachtest/clusterstats/BUILD.bazel index 4af2286ab1e1..7b678a564ce6 100644 --- a/pkg/cmd/roachtest/clusterstats/BUILD.bazel +++ b/pkg/cmd/roachtest/clusterstats/BUILD.bazel @@ -15,6 +15,7 @@ go_library( deps = [ "//pkg/cmd/roachtest/cluster", "//pkg/cmd/roachtest/option", + "//pkg/cmd/roachtest/roachtestutil", "//pkg/cmd/roachtest/test", "//pkg/roachprod/logger", "//pkg/roachprod/prometheus", diff --git a/pkg/cmd/roachtest/roachtestutil/BUILD.bazel b/pkg/cmd/roachtest/roachtestutil/BUILD.bazel index c8e7dfa53955..732bea5b198d 100644 --- a/pkg/cmd/roachtest/roachtestutil/BUILD.bazel +++ b/pkg/cmd/roachtest/roachtestutil/BUILD.bazel @@ -22,10 +22,12 @@ go_library( "//pkg/roachprod/install", "//pkg/roachprod/logger", "//pkg/testutils/sqlutils", + "//pkg/util/httputil", "//pkg/util/humanizeutil", "//pkg/util/log", "//pkg/util/timeutil", "@com_github_cockroachdb_errors//:errors", + "@com_github_pkg_errors//:errors", ], ) diff --git a/pkg/cmd/roachtest/roachtestutil/httpclient.go b/pkg/cmd/roachtest/roachtestutil/httpclient.go new file mode 100644 index 000000000000..38a0bcca7fa2 --- /dev/null +++ b/pkg/cmd/roachtest/roachtestutil/httpclient.go @@ -0,0 +1,213 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package roachtestutil + +import ( + "context" + "crypto/tls" + "fmt" + "net/http" + "net/http/cookiejar" + "net/url" + "strings" + "time" + + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster" + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option" + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test" + "github.com/cockroachdb/cockroach/pkg/roachprod/logger" + "github.com/cockroachdb/cockroach/pkg/util/httputil" + "github.com/cockroachdb/cockroach/pkg/util/protoutil" + "github.com/cockroachdb/cockroach/pkg/util/syncutil" + "github.com/cockroachdb/errors" +) + +// RoachtestHTTPClient is a wrapper over httputil.client that: +// - Adds cookies for every request, allowing the client to +// authenticate with secure clusters. +// - Sets InsecureSkipVerify to true. Roachprod certs are not +// signed by a verified root authority. +// - Unsets DialContext and DisableKeepAlives set in httputil. +// These two settings are not needed for testing purposes. +// +// This hides the details of creating authenticated https requests +// from the roachtest user, as most of it is just boilerplate and +// the same across all tests. +// +// Note that this currently only supports requests to CRDB clusters. +type RoachtestHTTPClient struct { + client *httputil.Client + sessionID string + cluster cluster.Cluster + l *logger.Logger + // Used for safely adding to the cookie jar. + mu syncutil.Mutex +} + +type RoachtestHTTPOptions struct { + Timeout time.Duration +} + +func HTTPTimeout(timeout time.Duration) func(options *RoachtestHTTPOptions) { + return func(options *RoachtestHTTPOptions) { + options.Timeout = timeout + } +} + +// DefaultHTTPClient returns a default RoachtestHTTPClient. +// This should be the default HTTP client used if you are +// trying to make a request to a secure cluster. +func DefaultHTTPClient( + c cluster.Cluster, l *logger.Logger, opts ...func(options *RoachtestHTTPOptions), +) *RoachtestHTTPClient { + roachtestHTTP := RoachtestHTTPClient{ + client: httputil.NewClientWithTimeout(httputil.StandardHTTPTimeout), + cluster: c, + l: l, + } + + // Certificates created by roachprod start are not signed by + // a verified root authority. + roachtestHTTP.client.Transport = &http.Transport{ + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, + } + + httpOptions := &RoachtestHTTPOptions{} + for _, opt := range opts { + opt(httpOptions) + } + if httpOptions.Timeout != 0 { + roachtestHTTP.client.Timeout = httpOptions.Timeout + } + + return &roachtestHTTP +} + +func (r *RoachtestHTTPClient) Get(ctx context.Context, url string) (*http.Response, error) { + if err := r.addCookie(ctx, url); err != nil { + return nil, err + } + return r.client.Get(ctx, url) +} + +func (r *RoachtestHTTPClient) GetJSON( + ctx context.Context, path string, response protoutil.Message, +) error { + if err := r.addCookie(ctx, path); err != nil { + return err + } + return httputil.GetJSON(*r.client.Client, path, response) +} + +func (r *RoachtestHTTPClient) PostProtobuf( + ctx context.Context, path string, request, response protoutil.Message, +) error { + if err := r.addCookie(ctx, path); err != nil { + return err + } + return httputil.PostProtobuf(ctx, *r.client.Client, path, request, response) +} + +// ResetSession removes the saved sessionID so that it is retrieved +// again the next time a request is made. This is usually not required +// as the sessionID does not change unless the cluster is restarted. +func (r *RoachtestHTTPClient) ResetSession() { + r.sessionID = "" +} + +func (r *RoachtestHTTPClient) addCookie(ctx context.Context, cookieUrl string) error { + // If we haven't extracted the sessionID yet, do so. + if r.sessionID == "" { + id, err := getSessionID(ctx, r.cluster, r.l, r.cluster.All()) + if err != nil { + return errors.Wrapf(err, "roachtestutil.addCookie: unable to extract sessionID") + } + r.sessionID = id + } + + name, value, found := strings.Cut(r.sessionID, "=") + if !found { + return errors.New("Cookie not formatted correctly") + } + + url, err := url.Parse(cookieUrl) + if err != nil { + return errors.Wrapf(err, "roachtestutil.addCookie: unable to parse cookieUrl") + } + err = r.SetCookies(url, []*http.Cookie{{Name: name, Value: value}}) + if err != nil { + return errors.Wrapf(err, "roachtestutil.addCookie: unable to set cookie") + } + + return nil +} + +// SetCookies is a helper that checks if a client.CookieJar exists and creates +// one if it doesn't. It then sets the provided cookies through CookieJar.SetCookies. +func (r *RoachtestHTTPClient) SetCookies(u *url.URL, cookies []*http.Cookie) error { + r.mu.Lock() + defer r.mu.Unlock() + if r.client.Jar == nil { + jar, err := cookiejar.New(nil) + if err != nil { + return err + } + r.client.Jar = jar + } + r.client.Jar.SetCookies(u, cookies) + return nil +} + +func getSessionIDOnSingleNode( + ctx context.Context, c cluster.Cluster, l *logger.Logger, node option.NodeListOption, +) (string, error) { + loginCmd := fmt.Sprintf( + "%s auth-session login root --port={pgport%s} --certs-dir ./certs --format raw", + test.DefaultCockroachPath, node, + ) + res, err := c.RunWithDetailsSingleNode(ctx, l, option.WithNodes(node), loginCmd) + if err != nil { + return "", errors.Wrap(err, "failed to authenticate") + } + + var sessionCookie string + for _, line := range strings.Split(res.Stdout, "\n") { + if strings.HasPrefix(line, "session=") { + sessionCookie = line + } + } + if sessionCookie == "" { + return "", fmt.Errorf("failed to find session cookie in `login` output") + } + + return sessionCookie, nil +} + +func getSessionID( + ctx context.Context, c cluster.Cluster, l *logger.Logger, nodes option.NodeListOption, +) (string, error) { + var err error + var cookie string + // The session ID should be the same for all nodes so stop after we successfully + // get it from one node. + for _, node := range nodes { + cookie, err = getSessionIDOnSingleNode(ctx, c, l, c.Node(node)) + if err == nil { + break + } + l.Printf("%s auth session login failed on node %d: %v", test.DefaultCockroachPath, node, err) + } + if err != nil { + return "", errors.Wrapf(err, "roachtestutil.GetSessionID") + } + sessionID := strings.Split(cookie, ";")[0] + return sessionID, nil +} diff --git a/pkg/roachprod/install/cockroach.go b/pkg/roachprod/install/cockroach.go index d64e1b9b67bb..c43bec2a8f6f 100644 --- a/pkg/roachprod/install/cockroach.go +++ b/pkg/roachprod/install/cockroach.go @@ -174,8 +174,6 @@ const ( NoSQLTimeout = 0 defaultInitTarget = Node(1) - - Username = "roach" ) func (st StartTarget) String() string {