Implement GetCertificate and GetClientCertificate

The goal is to close the feature parity gap with stdlib's tls package.

Co-authored-by: Rachel Chen <rachel@chens.email>

Author: jkralik

certificate.go

@@ -1,35 +1,102 @@
 package dtls
 
 import (
+	"bytes"
 	"crypto/tls"
 	"crypto/x509"
+	"fmt"
 	"strings"
 )
 
-func (c *handshakeConfig) getCertificate(serverName string) (*tls.Certificate, error) {
-	c.mu.Lock()
-	defer c.mu.Unlock()
+// ClientHelloInfo contains information from a ClientHello message in order to
+// guide application logic in the GetCertificate.
+type ClientHelloInfo struct {
+	// ServerName indicates the name of the server requested by the client
+	// in order to support virtual hosting. ServerName is only set if the
+	// client is using SNI (see RFC 4366, Section 3.1).
+	ServerName string
 
-	if c.nameToCertificate == nil {
-		nameToCertificate := make(map[string]*tls.Certificate)
-		for i := range c.localCertificates {
-			cert := &c.localCertificates[i]
-			x509Cert := cert.Leaf
-			if x509Cert == nil {
-				var parseErr error
-				x509Cert, parseErr = x509.ParseCertificate(cert.Certificate[0])
-				if parseErr != nil {
-					continue
-				}
+	// CipherSuites lists the CipherSuites supported by the client (e.g.
+	// TLS_AES_128_GCM_SHA256, TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256).
+	CipherSuites []CipherSuiteID
+}
+
+// CertificateRequestInfo contains information from a server's
+// CertificateRequest message, which is used to demand a certificate and proof
+// of control from a client.
+type CertificateRequestInfo struct {
+	// AcceptableCAs contains zero or more, DER-encoded, X.501
+	// Distinguished Names. These are the names of root or intermediate CAs
+	// that the server wishes the returned certificate to be signed by. An
+	// empty slice indicates that the server has no preference.
+	AcceptableCAs [][]byte
+}
+
+// SupportsCertificate returns nil if the provided certificate is supported by
+// the server that sent the CertificateRequest. Otherwise, it returns an error
+// describing the reason for the incompatibility.
+// NOTE: original src: https://github.com/golang/go/blob/29b9a328d268d53833d2cc063d1d8b4bf6852675/src/crypto/tls/common.go#L1273
+func (cri *CertificateRequestInfo) SupportsCertificate(c *tls.Certificate) error {
+	if len(cri.AcceptableCAs) == 0 {
+		return nil
+	}
+
+	for j, cert := range c.Certificate {
+		x509Cert := c.Leaf
+		// Parse the certificate if this isn't the leaf node, or if
+		// chain.Leaf was nil.
+		if j != 0 || x509Cert == nil {
+			var err error
+			if x509Cert, err = x509.ParseCertificate(cert); err != nil {
+				return fmt.Errorf("failed to parse certificate #%d in the chain: %w", j, err)
 			}
-			if len(x509Cert.Subject.CommonName) > 0 {
-				nameToCertificate[strings.ToLower(x509Cert.Subject.CommonName)] = cert
+		}
+
+		for _, ca := range cri.AcceptableCAs {
+			if bytes.Equal(x509Cert.RawIssuer, ca) {
+				return nil
 			}
-			for _, san := range x509Cert.DNSNames {
-				nameToCertificate[strings.ToLower(san)] = cert
+		}
+	}
+	return errNotAcceptableCertificateChain
+}
+
+func (c *handshakeConfig) setNameToCertificateLocked() {
+	nameToCertificate := make(map[string]*tls.Certificate)
+	for i := range c.localCertificates {
+		cert := &c.localCertificates[i]
+		x509Cert := cert.Leaf
+		if x509Cert == nil {
+			var parseErr error
+			x509Cert, parseErr = x509.ParseCertificate(cert.Certificate[0])
+			if parseErr != nil {
+				continue
 			}
 		}
-		c.nameToCertificate = nameToCertificate
+		if len(x509Cert.Subject.CommonName) > 0 {
+			nameToCertificate[strings.ToLower(x509Cert.Subject.CommonName)] = cert
+		}
+		for _, san := range x509Cert.DNSNames {
+			nameToCertificate[strings.ToLower(san)] = cert
+		}
+	}
+	c.nameToCertificate = nameToCertificate
+}
+
+func (c *handshakeConfig) getCertificate(clientHelloInfo *ClientHelloInfo) (*tls.Certificate, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.localGetCertificate != nil &&
+		(len(c.localCertificates) == 0 || len(clientHelloInfo.ServerName) > 0) {
+		cert, err := c.localGetCertificate(clientHelloInfo)
+		if cert != nil || err != nil {
+			return cert, err
+		}
+	}
+
+	if c.nameToCertificate == nil {
+		c.setNameToCertificateLocked()
 	}
 
 	if len(c.localCertificates) == 0 {
@@ -41,11 +108,11 @@ func (c *handshakeConfig) getCertificate(serverName string) (*tls.Certificate, e
 		return &c.localCertificates[0], nil
 	}
 
-	if len(serverName) == 0 {
+	if len(clientHelloInfo.ServerName) == 0 {
 		return &c.localCertificates[0], nil
 	}
 
-	name := strings.TrimRight(strings.ToLower(serverName), ".")
+	name := strings.TrimRight(strings.ToLower(clientHelloInfo.ServerName), ".")
 
 	if cert, ok := c.nameToCertificate[name]; ok {
 		return cert, nil
@@ -65,3 +132,23 @@ func (c *handshakeConfig) getCertificate(serverName string) (*tls.Certificate, e
 	// If nothing matches, return the first certificate.
 	return &c.localCertificates[0], nil
 }
+
+// NOTE: original src: https://github.com/golang/go/blob/29b9a328d268d53833d2cc063d1d8b4bf6852675/src/crypto/tls/handshake_client.go#L974
+func (c *handshakeConfig) getClientCertificate(cri *CertificateRequestInfo) (*tls.Certificate, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.localGetClientCertificate != nil {
+		return c.localGetClientCertificate(cri)
+	}
+
+	for i := range c.localCertificates {
+		chain := c.localCertificates[i]
+		if err := cri.SupportsCertificate(&chain); err != nil {
+			continue
+		}
+		return &chain, nil
+	}
+
+	// No acceptable certificate found. Don't send a certificate.
+	return new(tls.Certificate), nil
+}

certificate_test.go

@@ -24,47 +24,71 @@ func TestGetCertificate(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	cfg := &handshakeConfig{
-		localCertificates: []tls.Certificate{
-			certificateRandom,
-			certificateTest,
-			certificateWildcard,
-		},
-	}
-
 	testCases := []struct {
+		localCertificates   []tls.Certificate
 		desc                string
 		serverName          string
 		expectedCertificate tls.Certificate
+		getCertificate      func(info *ClientHelloInfo) (*tls.Certificate, error)
 	}{
 		{
-			desc:                "Simple match in CN",
+			desc: "Simple match in CN",
+			localCertificates: []tls.Certificate{
+				certificateRandom,
+				certificateTest,
+				certificateWildcard,
+			},
 			serverName:          "test.test",
 			expectedCertificate: certificateTest,
 		},
 		{
-			desc:                "Simple match in SANs",
+			desc: "Simple match in SANs",
+			localCertificates: []tls.Certificate{
+				certificateRandom,
+				certificateTest,
+				certificateWildcard,
+			},
 			serverName:          "www.test.test",
 			expectedCertificate: certificateTest,
 		},
 
 		{
-			desc:                "Wildcard match",
+			desc: "Wildcard match",
+			localCertificates: []tls.Certificate{
+				certificateRandom,
+				certificateTest,
+				certificateWildcard,
+			},
 			serverName:          "foo.test.test",
 			expectedCertificate: certificateWildcard,
 		},
 		{
-			desc:                "No match return first",
+			desc: "No match return first",
+			localCertificates: []tls.Certificate{
+				certificateRandom,
+				certificateTest,
+				certificateWildcard,
+			},
 			serverName:          "foo.bar",
 			expectedCertificate: certificateRandom,
 		},
+		{
+			desc: "Get certificate from callback",
+			getCertificate: func(info *ClientHelloInfo) (*tls.Certificate, error) {
+				return &certificateTest, nil
+			},
+			expectedCertificate: certificateTest,
+		},
 	}
 
 	for _, test := range testCases {
 		test := test
-
 		t.Run(test.desc, func(t *testing.T) {
-			cert, err := cfg.getCertificate(test.serverName)
+			cfg := &handshakeConfig{
+				localCertificates:   test.localCertificates,
+				localGetCertificate: test.getCertificate,
+			}
+			cert, err := cfg.getCertificate(&ClientHelloInfo{ServerName: test.serverName})
 			if err != nil {
 				t.Fatal(err)
 			}

config.go

@@ -147,6 +147,27 @@ type Config struct {
 	// If an ECC ciphersuite is configured and EllipticCurves is empty
 	// it will default to X25519, P-256, P-384 in this specific order.
 	EllipticCurves []elliptic.Curve
+
+	// GetCertificate returns a Certificate based on the given
+	// ClientHelloInfo. It will only be called if the client supplies SNI
+	// information or if Certificates is empty.
+	//
+	// If GetCertificate is nil or returns nil, then the certificate is
+	// retrieved from NameToCertificate. If NameToCertificate is nil, the
+	// best element of Certificates will be used.
+	GetCertificate func(*ClientHelloInfo) (*tls.Certificate, error)
+
+	// GetClientCertificate, if not nil, is called when a server requests a
+	// certificate from a client. If set, the contents of Certificates will
+	// be ignored.
+	//
+	// If GetClientCertificate returns an error, the handshake will be
+	// aborted and that error will be returned. Otherwise
+	// GetClientCertificate must return a non-nil Certificate. If
+	// Certificate.Certificate is empty then no certificate will be sent to
+	// the server. If this is unacceptable to the server then it may abort
+	// the handshake.
+	GetClientCertificate func(*CertificateRequestInfo) (*tls.Certificate, error)
 }
 
 func defaultConnectContextMaker() (context.Context, func()) {
@@ -160,6 +181,10 @@ func (c *Config) connectContextMaker() (context.Context, func()) {
 	return c.ConnectContextMaker()
 }
 
+func (c *Config) includeCertificateSuites() bool {
+	return c.PSK == nil || len(c.Certificates) > 0 || c.GetCertificate != nil || c.GetClientCertificate != nil
+}
+
 const defaultMTU = 1200 // bytes
 
 var defaultCurves = []elliptic.Curve{elliptic.X25519, elliptic.P256, elliptic.P384} //nolint:gochecknoglobals
@@ -215,6 +240,6 @@ func validateConfig(config *Config) error {
 		}
 	}
 
-	_, err := parseCipherSuites(config.CipherSuites, config.CustomCipherSuites, config.PSK == nil || len(config.Certificates) > 0, config.PSK != nil)
+	_, err := parseCipherSuites(config.CipherSuites, config.CustomCipherSuites, config.includeCertificateSuites(), config.PSK != nil)
 	return err
 }