From 88c6311ab801d181d5a346ac9c0be3a8623b3db4 Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Sun, 8 May 2016 22:57:53 +0200
Subject: [PATCH] read tls.Config data properly in RSA proof implementation

ref #48
---
 crypto/proof_rsa.go      | 54 ++++++++++++++++++++++------------------
 crypto/proof_rsa_test.go | 12 ++++++---
 2 files changed, 39 insertions(+), 27 deletions(-)

diff --git a/crypto/proof_rsa.go b/crypto/proof_rsa.go
index 386477a5..758cbb6e 100644
--- a/crypto/proof_rsa.go
+++ b/crypto/proof_rsa.go
@@ -9,7 +9,6 @@ import (
 	"crypto/rsa"
 	"crypto/sha256"
 	"crypto/tls"
-	"crypto/x509"
 	"errors"
 	"strings"
 
@@ -18,33 +17,25 @@ import (
 
 // rsaSigner stores a key and a certificate for the server proof
 type rsaSigner struct {
-	key    *rsa.PrivateKey
-	cert   *x509.Certificate
 	config *tls.Config
 }
 
 // NewRSASigner loads the key and cert from files
 func NewRSASigner(tlsConfig *tls.Config) (Signer, error) {
-	if len(tlsConfig.Certificates) == 0 {
-		return nil, errors.New("Expected at least one certificate in TLS config")
-	}
-	cert := tlsConfig.Certificates[0]
-
-	x509Cert, err := x509.ParseCertificate(cert.Certificate[0])
-	if err != nil {
-		return nil, err
-	}
-
-	rsaKey, ok := cert.PrivateKey.(*rsa.PrivateKey)
-	if !ok {
-		return nil, errors.New("Only RSA private keys are supported for now")
-	}
-
-	return &rsaSigner{key: rsaKey, cert: x509Cert, config: tlsConfig}, nil
+	return &rsaSigner{config: tlsConfig}, nil
 }
 
 // SignServerProof signs CHLO and server config for use in the server proof
 func (kd *rsaSigner) SignServerProof(sni string, chlo []byte, serverConfigData []byte) ([]byte, error) {
+	cert, err := kd.getCertForSNI(sni)
+	if err != nil {
+		return nil, err
+	}
+	key, ok := cert.PrivateKey.(*rsa.PrivateKey)
+	if !ok {
+		return nil, errors.New("only RSA keys are supported for now")
+	}
+
 	hash := sha256.New()
 	if len(chlo) > 0 {
 		// Version >= 31
@@ -56,34 +47,49 @@ func (kd *rsaSigner) SignServerProof(sni string, chlo []byte, serverConfigData [
 		hash.Write([]byte("QUIC server config signature\x00"))
 	}
 	hash.Write(serverConfigData)
-	return rsa.SignPSS(rand.Reader, kd.key, crypto.SHA256, hash.Sum(nil), &rsa.PSSOptions{SaltLength: 32})
+	return rsa.SignPSS(
+		rand.Reader,
+		key,
+		crypto.SHA256,
+		hash.Sum(nil),
+		&rsa.PSSOptions{SaltLength: 32},
+	)
 }
 
 // GetCertCompressed gets the certificate in the format described by the QUIC crypto doc
 func (kd *rsaSigner) GetCertCompressed(sni string) ([]byte, error) {
+	cert, err := kd.getCertForSNI(sni)
+	if err != nil {
+		return nil, err
+	}
+
 	b := &bytes.Buffer{}
 	b.WriteByte(1) // Entry type compressed
 	b.WriteByte(0) // Entry type end_of_list
-	utils.WriteUint32(b, uint32(len(kd.cert.Raw)+4))
+	utils.WriteUint32(b, uint32(len(cert.Certificate[0])+4))
 	gz, err := zlib.NewWriterLevelDict(b, flate.BestCompression, certDictZlib)
 	if err != nil {
 		panic(err)
 	}
-	lenCert := len(kd.cert.Raw)
+	lenCert := len(cert.Certificate[0])
 	gz.Write([]byte{
 		byte(lenCert & 0xff),
 		byte((lenCert >> 8) & 0xff),
 		byte((lenCert >> 16) & 0xff),
 		byte((lenCert >> 24) & 0xff),
 	})
-	gz.Write(kd.cert.Raw)
+	gz.Write(cert.Certificate[0])
 	gz.Close()
 	return b.Bytes(), nil
 }
 
 // GetCertUncompressed gets the certificate in DER
 func (kd *rsaSigner) GetCertUncompressed(sni string) ([]byte, error) {
-	return kd.cert.Raw, nil
+	cert, err := kd.getCertForSNI(sni)
+	if err != nil {
+		return nil, err
+	}
+	return cert.Certificate[0], nil
 }
 
 func (kd *rsaSigner) getCertForSNI(sni string) (*tls.Certificate, error) {
diff --git a/crypto/proof_rsa_test.go b/crypto/proof_rsa_test.go
index 84fa0dd3..d4340452 100644
--- a/crypto/proof_rsa_test.go
+++ b/crypto/proof_rsa_test.go
@@ -7,7 +7,6 @@ import (
 	"crypto"
 	"crypto/rsa"
 	"crypto/tls"
-	"crypto/x509"
 
 	"github.com/lucas-clemente/quic-go/testdata"
 
@@ -24,7 +23,13 @@ var _ = Describe("ProofRsa", func() {
 		z.Write([]byte{0x04, 0x00, 0x00, 0x00})
 		z.Write(cert)
 		z.Close()
-		kd := &rsaSigner{cert: &x509.Certificate{Raw: cert}}
+		kd := &rsaSigner{
+			config: &tls.Config{
+				Certificates: []tls.Certificate{
+					tls.Certificate{Certificate: [][]byte{cert}},
+				},
+			},
+		}
 		certCompressed, err := kd.GetCertCompressed("")
 		Expect(err).ToNot(HaveOccurred())
 		Expect(certCompressed).To(Equal(append([]byte{
@@ -34,6 +39,7 @@ var _ = Describe("ProofRsa", func() {
 	})
 
 	It("gives valid signatures", func() {
+		key := testdata.GetTLSConfig().Certificates[0].PrivateKey.(*rsa.PrivateKey).Public().(*rsa.PublicKey)
 		kd, err := NewRSASigner(testdata.GetTLSConfig())
 		Expect(err).ToNot(HaveOccurred())
 		signature, err := kd.SignServerProof("", []byte{'C', 'H', 'L', 'O'}, []byte{'S', 'C', 'F', 'G'})
@@ -41,7 +47,7 @@ var _ = Describe("ProofRsa", func() {
 		// Generated with:
 		// ruby -e 'require "digest"; p Digest::SHA256.digest("QUIC CHLO and server config signature\x00" + "\x20\x00\x00\x00" + Digest::SHA256.digest("CHLO") + "SCFG")'
 		data := []byte("W\xA6\xFC\xDE\xC7\xD2>c\xE6\xB5\xF6\tq\x9E|<~1\xA33\x01\xCA=\x19\xBD\xC1\xE4\xB0\xBA\x9B\x16%")
-		err = rsa.VerifyPSS(kd.(*rsaSigner).cert.PublicKey.(*rsa.PublicKey), crypto.SHA256, data, signature, &rsa.PSSOptions{SaltLength: 32})
+		err = rsa.VerifyPSS(key, crypto.SHA256, data, signature, &rsa.PSSOptions{SaltLength: 32})
 		Expect(err).ToNot(HaveOccurred())
 	})