Handshake state machine (#1656)

connection_state.go

@@ -1,15 +1,12 @@
 package nebula
 
 import (
-	"crypto/rand"
 	"encoding/json"
-	"fmt"
 	"sync"
 	"sync/atomic"
 
-	"github.com/flynn/noise"
 	"github.com/slackhq/nebula/cert"
-	"github.com/slackhq/nebula/noiseutil"
+	"github.com/slackhq/nebula/handshake"
 )
 
 const ReplayWindow = 1024
@@ -17,7 +14,6 @@ const ReplayWindow = 1024
 type ConnectionState struct {
 	eKey           *NebulaCipherState
 	dKey           *NebulaCipherState
-	H              *noise.HandshakeState
 	myCert         cert.Certificate
 	peerCert       *cert.CachedCertificate
 	initiator      bool
@@ -26,55 +22,24 @@ type ConnectionState struct {
 	writeLock      sync.Mutex
 }
 
-func NewConnectionState(cs *CertState, crt cert.Certificate, initiator bool, pattern noise.HandshakePattern) (*ConnectionState, error) {
-	var dhFunc noise.DHFunc
-	switch crt.Curve() {
-	case cert.Curve_CURVE25519:
-		dhFunc = noise.DH25519
-	case cert.Curve_P256:
-		if cs.pkcs11Backed {
-			dhFunc = noiseutil.DHP256PKCS11
-		} else {
-			dhFunc = noiseutil.DHP256
-		}
-	default:
-		return nil, fmt.Errorf("invalid curve: %s", crt.Curve())
-	}
-
-	var ncs noise.CipherSuite
-	if cs.cipher == "chachapoly" {
-		ncs = noise.NewCipherSuite(dhFunc, noise.CipherChaChaPoly, noise.HashSHA256)
-	} else {
-		ncs = noise.NewCipherSuite(dhFunc, noiseutil.CipherAESGCM, noise.HashSHA256)
-	}
-
-	static := noise.DHKey{Private: cs.privateKey, Public: crt.PublicKey()}
-	hs, err := noise.NewHandshakeState(noise.Config{
-		CipherSuite:   ncs,
-		Random:        rand.Reader,
-		Pattern:       pattern,
-		Initiator:     initiator,
-		StaticKeypair: static,
-		//NOTE: These should come from CertState (pki.go) when we finally implement it
-		PresharedKey:          []byte{},
-		PresharedKeyPlacement: 0,
-	})
-	if err != nil {
-		return nil, fmt.Errorf("NewConnectionState: %s", err)
-	}
-
-	// The queue and ready params prevent a counter race that would happen when
-	// sending stored packets and simultaneously accepting new traffic.
+// newConnectionStateFromResult builds a fully-populated ConnectionState from a
+// completed handshake.Result. It seeds messageCounter and the replay window so
+// that the post-handshake message indices already used on the wire don't count
+// as missed traffic in the data plane.
+func newConnectionStateFromResult(r *handshake.Result) *ConnectionState {
 	ci := &ConnectionState{
-		H:         hs,
-		initiator: initiator,
+		myCert:    r.MyCert,
+		initiator: r.Initiator,
+		peerCert:  r.RemoteCert,
+		eKey:      NewNebulaCipherState(r.EKey),
+		dKey:      NewNebulaCipherState(r.DKey),
 		window:    NewBits(ReplayWindow),
-		myCert:    crt,
 	}
-	// always start the counter from 2, as packet 1 and packet 2 are handshake packets.
-	ci.messageCounter.Add(2)
-
-	return ci, nil
+	ci.messageCounter.Add(r.MessageIndex)
+	for i := uint64(1); i <= r.MessageIndex; i++ {
+		ci.window.Update(nil, i)
+	}
+	return ci
 }
 
 func (cs *ConnectionState) MarshalJSON() ([]byte, error) {