Implement VerifyConnection as is in tls.Config

VerifyConnection, if not nil, is called after normal certificate
verification/PSK and after VerifyPeerCertificate by either a TLS client
or server. If it returns a non-nil error, the handshake is aborted
and that error results.

If normal verification fails then the handshake will abort before
considering this callback. This callback will run for all connections
regardless of InsecureSkipVerify or ClientAuth settings.

Author: jkralik

config.go

@@ -83,6 +83,16 @@ type Config struct {
 	// be considered but the verifiedChains will always be nil.
 	VerifyPeerCertificate func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error
 
+	// VerifyConnection, if not nil, is called after normal certificate
+	// verification/PSK and after VerifyPeerCertificate by either a TLS client
+	// or server. If it returns a non-nil error, the handshake is aborted
+	// and that error results.
+	//
+	// If normal verification fails then the handshake will abort before
+	// considering this callback. This callback will run for all connections
+	// regardless of InsecureSkipVerify or ClientAuth settings.
+	VerifyConnection func(*State) error
+
 	// RootCAs defines the set of root certificate authorities
 	// that one peer uses when verifying the other peer's certificates.
 	// If RootCAs is nil, TLS uses the host's root CA set.

conn.go

@@ -172,6 +172,7 @@ func createConn(ctx context.Context, nextConn net.Conn, config *Config, isClient
 		localCertificates:           config.Certificates,
 		insecureSkipVerify:          config.InsecureSkipVerify,
 		verifyPeerCertificate:       config.VerifyPeerCertificate,
+		verifyConnection:            config.VerifyConnection,
 		rootCAs:                     config.RootCAs,
 		clientCAs:                   config.ClientCAs,
 		customCipherSuites:          config.CustomCipherSuites,

conn_test.go

@@ -15,6 +15,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -466,15 +467,42 @@ func TestPSK(t *testing.T) {
 	defer report()
 
 	for _, test := range []struct {
-		Name           string
-		ServerIdentity []byte
-		CipherSuites   []CipherSuiteID
+		Name                   string
+		ServerIdentity         []byte
+		CipherSuites           []CipherSuiteID
+		ClientVerifyConnection func(*State) error
+		ServerVerifyConnection func(*State) error
+		WantFail               bool
+		ExpectedServerErr      string
+		ExpectedClientErr      string
 	}{
 		{
 			Name:           "Server identity specified",
 			ServerIdentity: []byte("Test Identity"),
 			CipherSuites:   []CipherSuiteID{TLS_PSK_WITH_AES_128_CCM_8},
 		},
+		{
+			Name:           "Server identity specified - Server verify connection fails",
+			ServerIdentity: []byte("Test Identity"),
+			CipherSuites:   []CipherSuiteID{TLS_PSK_WITH_AES_128_CCM_8},
+			ServerVerifyConnection: func(s *State) error {
+				return errExample
+			},
+			WantFail:          true,
+			ExpectedServerErr: errExample.Error(),
+			ExpectedClientErr: alert.BadCertificate.String(),
+		},
+		{
+			Name:           "Server identity specified - Client verify connection fails",
+			ServerIdentity: []byte("Test Identity"),
+			CipherSuites:   []CipherSuiteID{TLS_PSK_WITH_AES_128_CCM_8},
+			ClientVerifyConnection: func(s *State) error {
+				return errExample
+			},
+			WantFail:          true,
+			ExpectedServerErr: alert.BadCertificate.String(),
+			ExpectedClientErr: errExample.Error(),
+		},
 		{
 			Name:           "Server identity nil",
 			ServerIdentity: nil,
@@ -513,8 +541,9 @@ func TestPSK(t *testing.T) {
 
 						return []byte{0xAB, 0xC1, 0x23}, nil
 					},
-					PSKIdentityHint: clientIdentity,
-					CipherSuites:    test.CipherSuites,
+					PSKIdentityHint:  clientIdentity,
+					CipherSuites:     test.CipherSuites,
+					VerifyConnection: test.ClientVerifyConnection,
 				}
 
 				c, err := testClient(ctx, ca, conf, false)
@@ -528,11 +557,22 @@ func TestPSK(t *testing.T) {
 					}
 					return []byte{0xAB, 0xC1, 0x23}, nil
 				},
-				PSKIdentityHint: test.ServerIdentity,
-				CipherSuites:    test.CipherSuites,
+				PSKIdentityHint:  test.ServerIdentity,
+				CipherSuites:     test.CipherSuites,
+				VerifyConnection: test.ServerVerifyConnection,
 			}
 
 			server, err := testServer(ctx, cb, config, false)
+			if test.WantFail {
+				res := <-clientRes
+				if err == nil || !strings.Contains(err.Error(), test.ExpectedServerErr) {
+					t.Fatalf("TestPSK: Server expected(%v) actual(%v)", test.ExpectedServerErr, err)
+				}
+				if res.err == nil || !strings.Contains(res.err.Error(), test.ExpectedClientErr) {
+					t.Fatalf("TestPSK: Client expected(%v) actual(%v)", test.ExpectedClientErr, res.err)
+				}
+				return
+			}
 			if err != nil {
 				t.Fatalf("TestPSK: Server failed(%v)", err)
 			}
@@ -788,6 +828,29 @@ func TestClientCertificate(t *testing.T) {
 					ClientCAs:    caPool,
 				},
 			},
+			"NoClientCert_ServerVerifyConnectionFails": {
+				clientCfg: &Config{RootCAs: srvCAPool},
+				serverCfg: &Config{
+					Certificates: []tls.Certificate{srvCert},
+					ClientAuth:   NoClientCert,
+					ClientCAs:    caPool,
+					VerifyConnection: func(s *State) error {
+						return errExample
+					},
+				},
+				wantErr: true,
+			},
+			"NoClientCert_ClientVerifyConnectionFails": {
+				clientCfg: &Config{RootCAs: srvCAPool, VerifyConnection: func(s *State) error {
+					return errExample
+				}},
+				serverCfg: &Config{
+					Certificates: []tls.Certificate{srvCert},
+					ClientAuth:   NoClientCert,
+					ClientCAs:    caPool,
+				},
+				wantErr: true,
+			},
 			"NoClientCert_cert": {
 				clientCfg: &Config{RootCAs: srvCAPool, Certificates: []tls.Certificate{cert}},
 				serverCfg: &Config{
@@ -850,11 +913,22 @@ func TestClientCertificate(t *testing.T) {
 				wantErr: true,
 			},
 			"RequireAndVerifyClientCert": {
-				clientCfg: &Config{RootCAs: srvCAPool, Certificates: []tls.Certificate{cert}},
+				clientCfg: &Config{RootCAs: srvCAPool, Certificates: []tls.Certificate{cert}, VerifyConnection: func(s *State) error {
+					if ok := bytes.Equal(s.PeerCertificates[0], srvCertificate.Raw); !ok {
+						return errExample
+					}
+					return nil
+				}},
 				serverCfg: &Config{
 					Certificates: []tls.Certificate{srvCert},
 					ClientAuth:   RequireAndVerifyClientCert,
 					ClientCAs:    caPool,
+					VerifyConnection: func(s *State) error {
+						if ok := bytes.Equal(s.PeerCertificates[0], certificate.Raw); !ok {
+							return errExample
+						}
+						return nil
+					},
 				},
 			},
 		}

flight4handler.go

@@ -91,7 +91,7 @@ func flight4Parse(ctx context.Context, c flightConn, state *State, cache *handsh
 		state.peerCertificatesVerified = verified
 	} else if state.PeerCertificates != nil {
 		// A certificate was received, but we haven't seen a CertificateVerify
-		// keep reading until we receieve one
+		// keep reading until we receive one
 		return 0, nil, nil
 	}
 
@@ -178,6 +178,11 @@ func flight4Parse(ctx context.Context, c flightConn, state *State, cache *handsh
 	}
 
 	if state.cipherSuite.AuthenticationType() == CipherSuiteAuthenticationTypeAnonymous {
+		if cfg.verifyConnection != nil {
+			if err := cfg.verifyConnection(state.clone()); err != nil {
+				return 0, &alert.Alert{Level: alert.Fatal, Description: alert.BadCertificate}, err
+			}
+		}
 		return flight6, nil, nil
 	}
 
@@ -198,7 +203,12 @@ func flight4Parse(ctx context.Context, c flightConn, state *State, cache *handsh
 			return 0, &alert.Alert{Level: alert.Fatal, Description: alert.BadCertificate}, errClientCertificateNotVerified
 		}
 	case NoClientCert, RequestClientCert:
-		return flight6, nil, nil
+		// go to flight6
+	}
+	if cfg.verifyConnection != nil {
+		if err := cfg.verifyConnection(state.clone()); err != nil {
+			return 0, &alert.Alert{Level: alert.Fatal, Description: alert.BadCertificate}, err
+		}
 	}
 
 	return flight6, nil, nil

flight5handler.go

@@ -328,6 +328,11 @@ func initalizeCipherSuite(state *State, cache *handshakeCache, cfg *handshakeCon
 			}
 		}
 	}
+	if cfg.verifyConnection != nil {
+		if err = cfg.verifyConnection(state.clone()); err != nil {
+			return &alert.Alert{Level: alert.Fatal, Description: alert.BadCertificate}, err
+		}
+	}
 
 	if err = state.cipherSuite.Init(state.masterSecret, clientRandom[:], serverRandom[:], true); err != nil {
 		return &alert.Alert{Level: alert.Fatal, Description: alert.InternalError}, err

handshaker.go

@@ -102,6 +102,7 @@ type handshakeConfig struct {
 	nameToCertificate           map[string]*tls.Certificate
 	insecureSkipVerify          bool
 	verifyPeerCertificate       func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error
+	verifyConnection            func(*State) error
 	sessionStore                SessionStore
 	rootCAs                     *x509.CertPool
 	clientCAs                   *x509.CertPool