github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/testutils/sqlutils/pg_url.go (about)

     1  // Copyright 2016 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package sqlutils
    12  
    13  import (
    14  	"fmt"
    15  	"io/ioutil"
    16  	"net"
    17  	"net/url"
    18  	"os"
    19  	"path/filepath"
    20  	"testing"
    21  
    22  	"github.com/cockroachdb/cockroach/pkg/security"
    23  	"github.com/cockroachdb/cockroach/pkg/security/securitytest"
    24  	"github.com/cockroachdb/cockroach/pkg/util/fileutil"
    25  )
    26  
    27  // PGUrl returns a postgres connection url which connects to this server with the given user, and a
    28  // cleanup function which must be called after all connections created using the connection url have
    29  // been closed.
    30  //
    31  // In order to connect securely using postgres, this method will create temporary on-disk copies of
    32  // certain embedded security certificates. The certificates will be created in a new temporary
    33  // directory. The returned cleanup function will delete this temporary directory.
    34  // Note that two calls to this function for the same `user` will generate different
    35  // copies of the certificates, so the cleanup function must always be called.
    36  //
    37  // Args:
    38  //  prefix: A prefix to be prepended to the temp file names generated, for debugging.
    39  func PGUrl(t testing.TB, servingAddr, prefix string, user *url.Userinfo) (url.URL, func()) {
    40  	return PGUrlWithOptionalClientCerts(t, servingAddr, prefix, user, true /* withCerts */)
    41  }
    42  
    43  // PGUrlWithOptionalClientCerts is like PGUrl but the caller can
    44  // customize whether the client certificates are loaded on-disk and in the URL.
    45  func PGUrlWithOptionalClientCerts(
    46  	t testing.TB, servingAddr, prefix string, user *url.Userinfo, withClientCerts bool,
    47  ) (url.URL, func()) {
    48  	host, port, err := net.SplitHostPort(servingAddr)
    49  	if err != nil {
    50  		t.Fatal(err)
    51  	}
    52  
    53  	// TODO(benesch): Audit usage of prefix and replace the following line with
    54  	// `testutils.TempDir(t)` if prefix can always be `t.Name()`.
    55  	tempDir, err := ioutil.TempDir("", fileutil.EscapeFilename(prefix))
    56  	if err != nil {
    57  		t.Fatal(err)
    58  	}
    59  
    60  	caPath := filepath.Join(security.EmbeddedCertsDir, security.EmbeddedCACert)
    61  	tempCAPath := securitytest.RestrictedCopy(t, caPath, tempDir, "ca")
    62  	options := url.Values{}
    63  	options.Add("sslrootcert", tempCAPath)
    64  
    65  	if withClientCerts {
    66  		certPath := filepath.Join(security.EmbeddedCertsDir, fmt.Sprintf("client.%s.crt", user.Username()))
    67  		keyPath := filepath.Join(security.EmbeddedCertsDir, fmt.Sprintf("client.%s.key", user.Username()))
    68  
    69  		// Copy these assets to disk from embedded strings, so this test can
    70  		// run from a standalone binary.
    71  		tempCertPath := securitytest.RestrictedCopy(t, certPath, tempDir, "cert")
    72  		tempKeyPath := securitytest.RestrictedCopy(t, keyPath, tempDir, "key")
    73  		options.Add("sslcert", tempCertPath)
    74  		options.Add("sslkey", tempKeyPath)
    75  		options.Add("sslmode", "verify-full")
    76  	} else {
    77  		options.Add("sslmode", "verify-ca")
    78  	}
    79  
    80  	return url.URL{
    81  			Scheme:   "postgres",
    82  			User:     user,
    83  			Host:     net.JoinHostPort(host, port),
    84  			RawQuery: options.Encode(),
    85  		}, func() {
    86  			if err := os.RemoveAll(tempDir); err != nil {
    87  				// Not Fatal() because we might already be panicking.
    88  				t.Error(err)
    89  			}
    90  		}
    91  }