ecdh: allow for setting of handshake message contents

nikileshsa · Feb 19, 2019 · c6cd19e · c6cd19e
1 parent 00ff2fc
commit c6cd19e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/handshake/ecdh/mod.go b/handshake/ecdh/mod.go
@@ -56,7 +56,7 @@ func (b *block) OnBegin(p *protocol.Protocol, peer *noise.Peer) error {
 	}
 
 	req := Handshake{publicKey: ephemeralPublicKey}
-	req.signature, err = ephemeralPrivateKey.Sign(nil, []byte(msgEphemeralHandshake), crypto.Hash(0))
+	req.signature, err = ephemeralPrivateKey.Sign(nil, []byte(b.handshakeMessage), crypto.Hash(0))
 	if err != nil {
 		return errors.Wrap(errors.Wrap(protocol.DisconnectPeer, err.Error()), "failed to sign handshake message using ECDSA")
 	}
@@ -86,7 +86,7 @@ func (b *block) OnBegin(p *protocol.Protocol, peer *noise.Peer) error {
 		return errors.Wrap(protocol.DisconnectPeer, "failed to unmarshal our peers ephemeral public key")
 	}
 
-	if !edwards25519.Verify(peersPublicKey, []byte(msgEphemeralHandshake), res.signature) {
+	if !edwards25519.Verify(peersPublicKey, []byte(b.handshakeMessage), res.signature) {
 		return errors.Wrap(protocol.DisconnectPeer, "failed to verify signature in handshake request")
 	}