From 6a43060cce7a9294ee52534ddede2bdf7de9d3f9 Mon Sep 17 00:00:00 2001
From: lordwelch <timmy@narnian.us>
Date: Fri, 7 Aug 2020 14:43:44 -0700
Subject: [PATCH] Improve GCP error handling

---
 gcp/gcp.go             | 40 +++++++++++++++++++++++-----------------
 internal/signer/gcp.go |  2 +-
 2 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/gcp/gcp.go b/gcp/gcp.go
index f8141b7..fcda895 100644
--- a/gcp/gcp.go
+++ b/gcp/gcp.go
@@ -1,6 +1,7 @@
 package gcp
 
 import (
+	"bytes"
 	"crypto/rand"
 	"encoding/json"
 	"fmt"
@@ -8,35 +9,36 @@ import (
 	"os"
 
 	"github.com/stoggi/sshrimp/internal/config"
+	"github.com/stoggi/sshrimp/internal/identity"
 	"github.com/stoggi/sshrimp/internal/signer"
 	"golang.org/x/crypto/ssh"
 )
 
 func httpError(w http.ResponseWriter, v interface{}, statusCode int) {
-	e := json.NewEncoder(w)
-	err := e.Encode(v)
-	http.Error(w, err.Error(), statusCode)
+	var b bytes.Buffer
+	e := json.NewEncoder(&b)
+	_ = e.Encode(v)
+	http.Error(w, b.String(), statusCode)
 }
 
-// HandleRequest handles a request to sign an SSH public key verified by an OpenIDConnect id_token
+// SSHrimp handles a request to sign an SSH public key verified by an OpenIDConnect id_token
 func SSHrimp(w http.ResponseWriter, r *http.Request) {
-
 	// Load the configuration file, if not exsits, exit.
 	c := config.NewSSHrimp()
-	if err := c.Read(config.GetPath()); err != nil {
-		httpError(w, signer.SSHrimpResult{"", err.Error(), http.StatusText(http.StatusInternalServerError)}, http.StatusInternalServerError)
+	if err := c.Read("./serverless_function_source_code/sshrimp.toml"); err != nil {
+		httpError(w, signer.SSHrimpResult{Certificate: "", ErrorMessage: err.Error(), ErrorType: http.StatusText(http.StatusInternalServerError)}, http.StatusInternalServerError)
 		return
 	}
 
 	var event signer.SSHrimpEvent
 	if err := json.NewDecoder(r.Body).Decode(&event); err != nil {
-		httpError(w, signer.SSHrimpResult{"", err.Error(), http.StatusText(http.StatusBadRequest)}, http.StatusBadRequest)
+		httpError(w, signer.SSHrimpResult{Certificate: "", ErrorMessage: err.Error(), ErrorType: http.StatusText(http.StatusBadRequest)}, http.StatusBadRequest)
 		return
 	}
 
 	certificate, err := signer.ValidateRequest(event, c, r.Header.Get("Function-Execution-Id"), fmt.Sprintf("%s/%s/%s", os.Getenv("GCP_PROJECT"), os.Getenv("FUNCTION_REGION"), os.Getenv("FUNCTION_NAME")))
 	if err != nil {
-		httpError(w, signer.SSHrimpResult{"", err.Error(), http.StatusText(http.StatusBadRequest)}, http.StatusBadRequest)
+		httpError(w, signer.SSHrimpResult{Certificate: "", ErrorMessage: err.Error(), ErrorType: http.StatusText(http.StatusBadRequest)}, http.StatusBadRequest)
 		return
 	}
 
@@ -45,20 +47,28 @@ func SSHrimp(w http.ResponseWriter, r *http.Request) {
 
 	sshAlgorithmSigner, err := signer.NewAlgorithmSignerFromSigner(kmsSigner, ssh.SigAlgoRSASHA2256)
 	if err != nil {
-		httpError(w, signer.SSHrimpResult{"", err.Error(), http.StatusText(http.StatusBadRequest)}, http.StatusBadRequest)
+		httpError(w, signer.SSHrimpResult{Certificate: "", ErrorMessage: err.Error(), ErrorType: http.StatusText(http.StatusBadRequest)}, http.StatusBadRequest)
 		return
 	}
 
 	// Sign the certificate!!
 	if err := certificate.SignCert(rand.Reader, sshAlgorithmSigner); err != nil {
-		httpError(w, signer.SSHrimpResult{"", err.Error(), http.StatusText(http.StatusBadRequest)}, http.StatusBadRequest)
+		httpError(w, signer.SSHrimpResult{Certificate: "", ErrorMessage: err.Error(), ErrorType: http.StatusText(http.StatusInternalServerError)}, http.StatusInternalServerError)
+		return
+	}
+	i, _ := identity.NewIdentity(c)
+	username, _ := i.Validate(event.Token)
+	cc := ssh.CertChecker{}
+	err = cc.CheckCert(username, &certificate)
+	if err != nil {
+		httpError(w, signer.SSHrimpResult{Certificate: "", ErrorMessage: err.Error(), ErrorType: http.StatusText(http.StatusInternalServerError)}, http.StatusBadRequest)
 		return
 	}
 
 	// Extract the public key (certificate) to return to the user
 	pubkey, err := ssh.ParsePublicKey(certificate.Marshal())
 	if err != nil {
-		httpError(w, signer.SSHrimpResult{"", err.Error(), http.StatusText(http.StatusBadRequest)}, http.StatusBadRequest)
+		httpError(w, signer.SSHrimpResult{Certificate: "", ErrorMessage: err.Error(), ErrorType: http.StatusText(http.StatusInternalServerError)}, http.StatusInternalServerError)
 		return
 	}
 
@@ -69,9 +79,5 @@ func SSHrimp(w http.ResponseWriter, r *http.Request) {
 		ErrorType:    "",
 	}
 	e := json.NewEncoder(w)
-	err = e.Encode(res)
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
+	_ = e.Encode(res)
 }
diff --git a/internal/signer/gcp.go b/internal/signer/gcp.go
index 6945f0f..0230c62 100644
--- a/internal/signer/gcp.go
+++ b/internal/signer/gcp.go
@@ -52,7 +52,7 @@ func (s *GCPSigner) Public() crypto.PublicKey {
 
 	publicKey, err := x509.ParsePKIXPublicKey(pemBlock.Bytes)
 	if err != nil {
-		fmt.Printf(err.Error())
+		fmt.Println(err.Error())
 		return nil
 	}