Add unit test for TrustEnv

Emyrk · Emyrk · commit 8483d2152131 · 2021-09-30T23:20:32.000-04:00
diff --git a/coder-sdk/workspace.go b/coder-sdk/workspace.go
@@ -2,6 +2,7 @@ package coder
 
 import (
 	"context"
+	"crypto/x509"
 	"encoding/pem"
 	"fmt"
 	"io"
@@ -517,8 +518,9 @@ func (c *DefaultClient) SetPolicyTemplate(ctx context.Context, templateID string
 
 // TrustChallengeResponse returns the tls certs used in the response.
 type TrustChallengeResponse struct {
-	// Certificates are the returned response certificates
-	Certificates [][]byte `json:"-"`
+	// PemCertificates are the returned response certificates
+	PemCertificates [][]byte            `json:"-"`
+	Certificates    []*x509.Certificate `json:"-"`
 }
 
 // TrustEnvironment is used to make a secure handshake with a coderd. This is intended to run from within a workspace.
@@ -534,12 +536,14 @@ func (c *DefaultClient) TrustEnvironment(ctx context.Context, id string) (*Trust
 	}
 
 	if len(resp.TLS.PeerCertificates) > 0 {
-		for _, c := range resp.TLS.PeerCertificates {
+		for i := range resp.TLS.PeerCertificates {
+			c := resp.TLS.PeerCertificates[i]
 			data := pem.EncodeToMemory(&pem.Block{
 				Type:  "CERTIFICATE",
 				Bytes: c.Raw,
 			})
-			challenge.Certificates = append(challenge.Certificates, data)
+			challenge.PemCertificates = append(challenge.PemCertificates, data)
+			challenge.Certificates = append(challenge.Certificates, c)
 		}
 	}
 
diff --git a/coder-sdk/workspace_test.go b/coder-sdk/workspace_test.go
@@ -0,0 +1,71 @@
+package coder
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_TrustEnvironment(t *testing.T) {
+	ctx := context.Background()
+
+	const version = "test"
+	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("coder-version", version)
+		_, _ = w.Write([]byte("{}"))
+		w.WriteHeader(http.StatusNoContent)
+	}))
+
+	u, err := url.Parse(srv.URL)
+	require.NoError(t, err)
+
+	c, err := NewClient(ClientOptions{
+		BaseURL: u,
+		Token:   "random",
+	})
+	require.NoError(t, err)
+
+	_, err = c.APIVersion(ctx)
+	require.Error(t, err)
+	require.Regexp(t, "x509: certificate signed by unknown authority", err.Error())
+
+	// TODO: @emyrk check proper handshake
+	c.httpClient = insecureHTTPClient()
+	challenge, err := c.TrustEnvironment(ctx, "random")
+	require.NoError(t, err)
+	require.Len(t, challenge.PemCertificates, 1)
+
+	// Add the cert to the trusted pool, try the api call again
+	pool := x509.NewCertPool()
+	for i := range challenge.Certificates {
+		pool.AddCert(challenge.Certificates[i])
+	}
+	conf := &tls.Config{RootCAs: pool}
+	c.httpClient = &http.Client{
+		Timeout: time.Second * 3,
+		Transport: &http.Transport{
+			TLSClientConfig: conf,
+		},
+	}
+
+	v, err := c.APIVersion(ctx)
+	require.NoError(t, err)
+	require.Equal(t, version, v)
+}
+
+func insecureHTTPClient() *http.Client {
+	conf := &tls.Config{InsecureSkipVerify: true}
+	return &http.Client{
+		Timeout: time.Second * 3,
+		Transport: &http.Transport{
+			TLSClientConfig: conf,
+		},
+	}
+}
diff --git a/internal/cmd/agent.go b/internal/cmd/agent.go
@@ -214,5 +214,5 @@ func trustCertificate(ctx context.Context, sessionToken string) ([][]byte, error
 		return nil, xerrors.Errorf("challenge failed: %w", err)
 	}
 
-	return challenge.Certificates, nil
+	return challenge.PemCertificates, nil
 }

Original file line number	Diff line number	Diff line change
`@@ -214,5 +214,5 @@ func trustCertificate(ctx context.Context, sessionToken string) ([][]byte, error`
`214`	`214`	`return nil, xerrors.Errorf("challenge failed: %w", err)`
`215`	`215`	`}`
`216`	`216`
`217`		`- return challenge.Certificates, nil`
	`217`	`+ return challenge.PemCertificates, nil`
`218`	`218`	`}`