From 38d78f984422cbeeab00cc7fee8818f6c14491f0 Mon Sep 17 00:00:00 2001
From: Austin Cherry <acherry@cloudflare.com>
Date: Mon, 13 May 2019 14:38:36 -0500
Subject: [PATCH] AUTH-1706: fixes and testing

---
 sshgen/sshgen.go      |  27 +++++++++-
 sshgen/sshgen_test.go | 112 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 138 insertions(+), 1 deletion(-)
 create mode 100644 sshgen/sshgen_test.go

diff --git a/sshgen/sshgen.go b/sshgen/sshgen.go
index 14d66a3f..12538c84 100644
--- a/sshgen/sshgen.go
+++ b/sshgen/sshgen.go
@@ -9,6 +9,8 @@ import (
 	"encoding/json"
 	"encoding/pem"
 	"errors"
+	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"net/url"
@@ -40,6 +42,14 @@ type signResponse struct {
 	ExpiresAt   time.Time `json:"expires_at"`
 }
 
+// ErrorResponse struct stores error information after any error-prone function
+type errorResponse struct {
+	Status  int    `json:"status"`
+	Message string `json:"message"`
+}
+
+var mockRequest func(url, contentType string, body io.Reader) (*http.Response, error) = nil
+
 // GenerateShortLivedCertificate generates and stores a keypair for short lived certs
 func GenerateShortLivedCertificate(appURL *url.URL, token string) error {
 	fullName, err := cfpath.GenerateFilePathFromURL(appURL, keyName)
@@ -96,13 +106,28 @@ func handleCertificateGeneration(token, fullName string) (string, error) {
 		return "", err
 	}
 
-	res, err := http.Post(issuer+signEndpoint, "application/json", bytes.NewBuffer(buf))
+	var res *http.Response
+	if mockRequest != nil {
+		res, err = mockRequest(issuer+signEndpoint, "application/json", bytes.NewBuffer(buf))
+	} else {
+		res, err = http.Post(issuer+signEndpoint, "application/json", bytes.NewBuffer(buf))
+	}
+
 	if err != nil {
 		return "", err
 	}
 	defer res.Body.Close()
 
 	decoder := json.NewDecoder(res.Body)
+
+	if res.StatusCode != 200 {
+		var errResponse errorResponse
+		if err := decoder.Decode(&errResponse); err != nil {
+			return "", err
+		}
+		return "", fmt.Errorf("%d: %s", errResponse.Status, errResponse.Message)
+	}
+
 	var signRes signResponse
 	if err := decoder.Decode(&signRes); err != nil {
 		return "", err
diff --git a/sshgen/sshgen_test.go b/sshgen/sshgen_test.go
new file mode 100644
index 00000000..d7094d14
--- /dev/null
+++ b/sshgen/sshgen_test.go
@@ -0,0 +1,112 @@
+package sshgen
+
+import (
+	"crypto/rand"
+	"crypto/rsa"
+	"encoding/json"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/cloudflare/cloudflared/cmd/cloudflared/config"
+	cfpath "github.com/cloudflare/cloudflared/cmd/cloudflared/path"
+	"github.com/coreos/go-oidc/jose"
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	audTest   = "cf-test-aud"
+	nonceTest = "asfd"
+)
+
+type signingArguments struct {
+	Principals   []string `json:"principals"`
+	ClientPubKey string   `json:"public_key"`
+	Duration     string   `json:"duration"`
+}
+
+func TestCertGenSuccess(t *testing.T) {
+	url, _ := url.Parse("https://cf-test-access.com/testpath")
+	token := tokenGenerator()
+
+	fullName, err := cfpath.GenerateFilePathFromURL(url, keyName)
+	assert.NoError(t, err)
+
+	pubKeyName := fullName + ".pub"
+	certKeyName := fullName + "-cert.pub"
+
+	defer func() {
+		os.Remove(fullName)
+		os.Remove(pubKeyName)
+		os.Remove(certKeyName)
+	}()
+
+	resp := signingArguments{
+		Principals:   []string{"dalton"},
+		ClientPubKey: "ecdsa-sha2-nistp256-cert-v01@openssh.com AAAAKGVjZHNhLXNoYTItbmlzdHAyNTYtY2VydC12MDFAb3BlbnNzaC5jb20AAAAg+0rYq4mNGIAHiH1xPOJXfmOpTEwFIcyXzGJieTOhRs8AAAAIbmlzdHAyNTYAAABBBJIcsq02e8ZaofJXOZKp7yQdKW/JIouJ90lybr76hHIRrZBL1t4JEimfLvNDphPrTW9VDQaIcBSKNaxRqHOS8jezoJbhFGWhqQAAAAEAAAAgZWU5OTliNGRkZmFmNjgxNDEwMTVhMDJiY2ZhMTdiN2UAAAAKAAAABmF1c3RpbgAAAABc1KFoAAAAAFzUohwAAAAAAAAARwAAABdwZXJtaXQtYWdlbnQtZm9yd2FyZGluZwAAAAAAAAAKcGVybWl0LXB0eQAAAAAAAAAOcGVybWl0LXVzZXItcmMAAAAAAAAAAAAAAGgAAAATZWNkc2Etc2hhMi1uaXN0cDI1NgAAAAhuaXN0cDI1NgAAAEEEeAuYR56XaxvH5Z1p0hDCTQ7wC4dbj0Gc+LOKu1f94og2ilZTv9tutg8cZrqAT97REmGH6j9KIOVLGsPVjajSKAAAAGQAAAATZWNkc2Etc2hhMi1uaXN0cDI1NgAAAEkAAAAhAORY9ZO3TQsrUm6ajnVW+arbnVfTkxYBYFlVoeOEXKZuAAAAIG96A8nQnTuprWXLSemWL68RXC1NVKnBOIPD2Z7UIOB1",
+		Duration:     "3m",
+	}
+	w := httptest.NewRecorder()
+	respJson, err := json.Marshal(resp)
+	assert.NoError(t, err)
+	w.Write(respJson)
+	mockRequest = func(url, contentType string, body io.Reader) (*http.Response, error) {
+		assert.Contains(t, "/cdn-cgi/access/cert_sign", url)
+		assert.Equal(t, "application/json", contentType)
+		buf, err := ioutil.ReadAll(body)
+		assert.NoError(t, err)
+		assert.NotEmpty(t, buf)
+		return w.Result(), nil
+	}
+
+	err = GenerateShortLivedCertificate(url, token)
+	assert.NoError(t, err)
+
+	exist, err := config.FileExists(fullName)
+	assert.NoError(t, err)
+	if !exist {
+		assert.FailNow(t, fmt.Sprintf("key should exist at: %s", fullName), fullName)
+		return
+	}
+
+	exist, err = config.FileExists(pubKeyName)
+	assert.NoError(t, err)
+	if !exist {
+		assert.FailNow(t, fmt.Sprintf("key should exist at: %s", pubKeyName), pubKeyName)
+		return
+	}
+
+	exist, err = config.FileExists(certKeyName)
+	assert.NoError(t, err)
+	if !exist {
+		assert.FailNow(t, fmt.Sprintf("key should exist at: %s", certKeyName), certKeyName)
+		return
+	}
+}
+
+func tokenGenerator() string {
+	iat := time.Now().Unix()
+	exp := time.Now().Add(time.Minute * 5).Unix()
+	claims := jose.Claims{}
+	claims.Add("aud", audTest)
+	claims.Add("iat", iat)
+	claims.Add("nonce", nonceTest)
+	claims.Add("exp", exp)
+
+	k, err := rsa.GenerateKey(rand.Reader, 512)
+	if err != nil {
+		return ""
+	}
+	signer := jose.NewSignerRSA("asdf", *k)
+	token, terr := jose.NewSignedJWT(claims, signer)
+	if terr != nil {
+		return ""
+	}
+	return token.Encode()
+}