From 70a7cea10bb12286973d836b6c01ffc0af6da805 Mon Sep 17 00:00:00 2001
From: Henry Case <kris@amongbytes.com>
Date: Sun, 12 May 2019 20:22:21 +0100
Subject: [PATCH] Add negotiated group to ConnectionState

---
 13.go               | 4 ++--
 cipher_suites.go    | 3 +++
 common.go           | 3 ++-
 conn.go             | 3 +++
 handshake_client.go | 2 +-
 handshake_server.go | 2 +-
 key_agreement.go    | 8 ++++++++
 7 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/13.go b/13.go
index 167b40f..d7f2fb6 100644
--- a/13.go
+++ b/13.go
@@ -351,6 +351,7 @@ CurvePreferenceLoop:
 		c.phase = waitingClientFinished
 	}
 
+	hs.c.Group = hs.hello.keyShare.group
 	return nil
 }
 
@@ -459,6 +460,7 @@ func (hs *serverHandshakeState) readClientFinished13(hasConfirmLock bool) error
 	c.in.traceErr, c.out.traceErr = nil, nil
 	c.phase = handshakeConfirmed
 	atomic.StoreInt32(&c.handshakeConfirmed, 1)
+	hs.c.Group = hs.hello.keyShare.group
 
 	// Any read operation after handshakeRunning and before handshakeConfirmed
 	// will be holding this lock, which we release as soon as the confirmation
@@ -470,7 +472,6 @@ func (hs *serverHandshakeState) readClientFinished13(hasConfirmLock bool) error
 	if hasConfirmLock {
 		c.confirmMutex.Unlock()
 	}
-
 	return hs.sendSessionTicket13() // TODO: do in a goroutine
 }
 
@@ -532,7 +533,6 @@ func (hs *serverHandshakeState) sendCertificate13() error {
 	if _, err := c.writeRecord(recordTypeHandshake, verifyMsg.marshal()); err != nil {
 		return err
 	}
-
 	return nil
 }
 
diff --git a/cipher_suites.go b/cipher_suites.go
index 7b5ec13..9c782b6 100644
--- a/cipher_suites.go
+++ b/cipher_suites.go
@@ -35,6 +35,9 @@ type keyAgreement interface {
 	// ServerKeyExchange message.
 	processServerKeyExchange(*Config, *clientHelloMsg, *serverHelloMsg, crypto.PublicKey, *serverKeyExchangeMsg) error
 	generateClientKeyExchange(*Config, *clientHelloMsg, crypto.PublicKey) ([]byte, *clientKeyExchangeMsg, error)
+
+	// Returns ID of negotiated group
+	NegotiatedGroup() CurveID
 }
 
 const (
diff --git a/common.go b/common.go
index a04d7e0..53d6c32 100644
--- a/common.go
+++ b/common.go
@@ -241,7 +241,8 @@ type ConnectionState struct {
 	// Unique0RTTToken is only present if HandshakeConfirmed is false.
 	Unique0RTTToken []byte
 
-	ClientHello []byte // ClientHello packet
+	ClientHello []byte  // ClientHello packet
+	Group       CurveID // ECDH group used
 }
 
 // ClientAuthType declares the policy the server will follow for
diff --git a/conn.go b/conn.go
index 4b8dc55..24aee1e 100644
--- a/conn.go
+++ b/conn.go
@@ -130,6 +130,8 @@ type Conn struct {
 	// accept the 0-RTT data. Exposed as ConnectionState.Unique0RTTToken.
 	binder []byte
 
+	Group CurveID // ECDH group used
+
 	tmp [16]byte
 }
 
@@ -1698,6 +1700,7 @@ func (c *Conn) ConnectionState() ConnectionState {
 				state.TLSUnique = c.serverFinished[:]
 			}
 		}
+		state.Group = c.Group
 	}
 
 	return state
diff --git a/handshake_client.go b/handshake_client.go
index 23894e7..9cbb3ba 100644
--- a/handshake_client.go
+++ b/handshake_client.go
@@ -638,7 +638,7 @@ func (hs *clientHandshakeState) doFullHandshake() error {
 			return err
 		}
 	}
-
+	c.Group = keyAgreement.NegotiatedGroup()
 	hs.finishedHash.discardHandshakeBuffer()
 
 	return nil
diff --git a/handshake_server.go b/handshake_server.go
index f2b9600..12c237e 100644
--- a/handshake_server.go
+++ b/handshake_server.go
@@ -647,7 +647,7 @@ func (hs *serverHandshakeState) doFullHandshake() error {
 
 		hs.finishedHash.Write(certVerify.marshal())
 	}
-
+	c.Group = keyAgreement.NegotiatedGroup()
 	hs.finishedHash.discardHandshakeBuffer()
 
 	return nil
diff --git a/key_agreement.go b/key_agreement.go
index d19c0b0..f9206e9 100644
--- a/key_agreement.go
+++ b/key_agreement.go
@@ -24,6 +24,10 @@ var errServerKeyExchange = errors.New("tls: invalid ServerKeyExchange message")
 // encrypts the pre-master secret to the server's public key.
 type rsaKeyAgreement struct{}
 
+func (ka rsaKeyAgreement) NegotiatedGroup() CurveID {
+	return CurveID(0)
+}
+
 func (ka rsaKeyAgreement) generateServerKeyExchange(config *Config, sk crypto.PrivateKey, clientHello *clientHelloMsg, hello *serverHelloMsg) (*serverKeyExchangeMsg, error) {
 	return nil, nil
 }
@@ -155,6 +159,10 @@ type ecdheKeyAgreement struct {
 	x, y *big.Int
 }
 
+func (ka *ecdheKeyAgreement) NegotiatedGroup() CurveID {
+	return ka.curveid
+}
+
 func (ka *ecdheKeyAgreement) generateServerKeyExchange(config *Config, sk crypto.PrivateKey, clientHello *clientHelloMsg, hello *serverHelloMsg) (*serverKeyExchangeMsg, error) {
 	preferredCurves := config.curvePreferences()