nonce should be tracked by conn

Author: jhawk28

conn.go

@@ -36,6 +36,9 @@ type Conn struct {
 
 	closed         int32
 	onCloseErrorCB func(c *Conn)
+
+	nonceIdx     uint64
+	peerNonceIdx uint64
 }
 
 func (c *Conn) Close() error {
@@ -518,3 +521,27 @@ func (conn *Conn) notifyOnCloseError() {
 	}
 	conn.onCloseErrorCB(conn)
 }
+
+func (conn *Conn) SetNonce(i uint64) {
+	conn.nonceIdx = i
+}
+
+func (conn *Conn) Nonce() uint64 {
+	return conn.nonceIdx
+}
+
+func (conn *Conn) SetPeerNonce(i uint64) {
+	conn.peerNonceIdx = i
+}
+
+func (conn *Conn) PeerNonce() uint64 {
+	return conn.peerNonceIdx
+}
+
+func (conn *Conn) IncrNonce() {
+	conn.nonceIdx++
+}
+
+func (conn *Conn) IncrPeerNonce() {
+	conn.peerNonceIdx++
+}

security.go

@@ -24,10 +24,10 @@ type Security interface {
 
 type SecurityEncryption interface {
 	// Encrypt writes the encrypted form of data to w.
-	Encrypt(data []byte, more bool) ([]byte, error)
+	Encrypt(conn *Conn, data []byte, more bool) ([]byte, error)
 
 	// Decrypt writes the decrypted form of data to w.
-	Decrypt(data []byte) ([]byte, bool, error)
+	Decrypt(conn *Conn, data []byte) ([]byte, bool, error)
 }
 
 // SecurityType denotes types of ZMTP security mechanisms

security/curve/curve.go

@@ -50,8 +50,6 @@ type security struct {
 	secretKey    [keySize]byte // Derived shared secret key
 	asServer     bool          // True if this is a server
 	sharedKey    [keySize]byte // Pre-computed shared key (for optimization)
-	nonceIdx     uint64
-	peerNonceIdx uint64
 }
 
 // SecurityForClient returns a CURVE security mechanism for a client.
@@ -123,8 +121,8 @@ func (sec *security) clientHandshake(conn *zmq4.Conn) error {
 		return fmt.Errorf("security/curve: failed READY: %w", err)
 	}
 
-	sec.nonceIdx = 3
-	sec.peerNonceIdx = 1
+	conn.SetNonce(3)
+	conn.SetPeerNonce(1)
 
 	box.Precompute(&sec.sharedKey, &sec.secretKey, &sec.ephemeral.Private)
 
@@ -168,26 +166,27 @@ func (sec *security) serverHandshake(conn *zmq4.Conn) error {
 		return fmt.Errorf("security/curve: Server ready failed: %w", err)
 	}
 
-	sec.nonceIdx = 2
-	sec.peerNonceIdx = 2
+	conn.SetNonce(2)
+	conn.SetPeerNonce(2)
+
 	box.Precompute(&sec.sharedKey, &clientTransPubKey, &kp.Private)
 	return nil
 }
 
 // Encrypt writes the encrypted form of data to w.
-func (sec *security) Encrypt(data []byte, more bool) ([]byte, error) {
-	defer func() { sec.nonceIdx++ }()
+func (sec *security) Encrypt(conn *zmq4.Conn, data []byte, more bool) ([]byte, error) {
+	defer func() { conn.IncrNonce() }()
 	out := make([]byte, 8+8+17+len(data))
 	out[0] = uint8(7)
 	copy(out[1:], "MESSAGE")
 
 	var nonce Nonce
 	if sec.asServer {
-		nonce.Short("CurveZMQMESSAGES", sec.nonceIdx) // From server
+		nonce.Short("CurveZMQMESSAGES", conn.Nonce()) // From server
 	} else {
-		nonce.Short("CurveZMQMESSAGEC", sec.nonceIdx) // From client
+		nonce.Short("CurveZMQMESSAGEC", conn.Nonce()) // From client
 	}
-	binary.BigEndian.AppendUint64(out[8:8], sec.nonceIdx)
+	binary.BigEndian.AppendUint64(out[8:8], conn.Nonce())
 	toSeal := make([]byte, 1+len(data))
 	if more {
 		toSeal[0] = 0x1
@@ -198,7 +197,7 @@ func (sec *security) Encrypt(data []byte, more bool) ([]byte, error) {
 }
 
 // Decrypt writes the decrypted form of data to w.
-func (sec *security) Decrypt(body []byte) ([]byte, bool, error) {
+func (sec *security) Decrypt(conn *zmq4.Conn, body []byte) ([]byte, bool, error) {
 	if len(body) < 33 {
 		return nil, false, fmt.Errorf("security/curve: invalid message: too short")
 	}
@@ -217,10 +216,10 @@ func (sec *security) Decrypt(body []byte) ([]byte, bool, error) {
 	} else {
 		nonce.Short("CurveZMQMESSAGES", shortNonce) // From server
 	}
-	if shortNonce != sec.peerNonceIdx+1 {
-		return nil, false, fmt.Errorf("Peer used invalid nonce (expected %d, got %d)", sec.peerNonceIdx+1, shortNonce)
+	if shortNonce != conn.PeerNonce()+1 {
+		return nil, false, fmt.Errorf("Peer used invalid nonce (expected %d, got %d)", conn.PeerNonce()+1, shortNonce)
 	}
-	sec.peerNonceIdx++
+	conn.IncrPeerNonce()
 	copy(nonce[16:], body[8:])
 	out := make([]byte, len(body)-32)
 	out, ok := box.OpenAfterPrecomputation(out[0:0], body[16:], nonce.N(), &sec.sharedKey)

security_test.go

@@ -5,7 +5,6 @@
 package zmq4
 
 import (
-	"bytes"
 	"context"
 	"fmt"
 	"os"
@@ -22,25 +21,6 @@ func TestNullSecurity(t *testing.T) {
 	if got, want := sec.Type(), NullSecurity; got != want {
 		t.Fatalf("got=%v, want=%v", got, want)
 	}
-
-	data := []byte("hello world")
-	wenc := new(bytes.Buffer)
-	if _, err := sec.Encrypt(wenc, data); err != nil {
-		t.Fatalf("error encrypting data: %+v", err)
-	}
-
-	if !bytes.Equal(wenc.Bytes(), data) {
-		t.Fatalf("error encrypted data.\ngot = %q\nwant= %q\n", wenc.Bytes(), data)
-	}
-
-	wdec := new(bytes.Buffer)
-	if _, err := sec.Decrypt(wdec, wenc.Bytes()); err != nil {
-		t.Fatalf("error decrypting data: %+v", err)
-	}
-
-	if !bytes.Equal(wdec.Bytes(), data) {
-		t.Fatalf("error decrypted data.\ngot = %q\nwant= %q\n", wdec.Bytes(), data)
-	}
 }
 
 func TestNullHandshakeReqRep(t *testing.T) {