1 files changed, 22 insertions, 31 deletions

diff --git a/device/peer.go b/device/peer.go
index 5bd52df..79feae7 100644
--- a/device/peer.go
+++ b/device/peer.go
@@ -16,24 +16,16 @@ import (
 )
 
 type Peer struct {
-	isRunning    AtomicBool
-	sync.RWMutex // Mostly protects endpoint, but is generally taken whenever we modify peer
-	keypairs     Keypairs
-	handshake    Handshake
-	device       *Device
-	endpoint     conn.Endpoint
-	stopping     sync.WaitGroup // routines pending stop
-
-	// These fields are accessed with atomic operations, which must be
-	// 64-bit aligned even on 32-bit platforms. Go guarantees that an
-	// allocated struct will be 64-bit aligned. So we place
-	// atomically-accessed fields up front, so that they can share in
-	// this alignment before smaller fields throw it off.
-	stats struct {
-		txBytes           uint64 // bytes send to peer (endpoint)
-		rxBytes           uint64 // bytes received from peer
-		lastHandshakeNano int64  // nano seconds since epoch
-	}
+	isRunning         atomic.Bool
+	sync.RWMutex      // Mostly protects endpoint, but is generally taken whenever we modify peer
+	keypairs          Keypairs
+	handshake         Handshake
+	device            *Device
+	endpoint          conn.Endpoint
+	stopping          sync.WaitGroup // routines pending stop
+	txBytes           atomic.Uint64  // bytes send to peer (endpoint)
+	rxBytes           atomic.Uint64  // bytes received from peer
+	lastHandshakeNano atomic.Int64   // nano seconds since epoch
 
 	disableRoaming bool
 
@@ -43,9 +35,9 @@ type Peer struct {
 		newHandshake            *Timer
 		zeroKeyMaterial         *Timer
 		persistentKeepalive     *Timer
-		handshakeAttempts       uint32
-		needAnotherKeepalive    AtomicBool
-		sentLastMinuteHandshake AtomicBool
+		handshakeAttempts       atomic.Uint32
+		needAnotherKeepalive    atomic.Bool
+		sentLastMinuteHandshake atomic.Bool
 	}
 
 	state struct {
@@ -60,7 +52,7 @@ type Peer struct {
 
 	cookieGenerator             CookieGenerator
 	trieEntries                 list.List
-	persistentKeepaliveInterval uint32 // accessed atomically
+	persistentKeepaliveInterval atomic.Uint32
 }
 
 func (device *Device) NewPeer(pk NoisePublicKey) (*Peer, error) {
@@ -133,7 +125,7 @@ func (peer *Peer) SendBuffer(buffer []byte) error {
 
 	err := peer.device.net.bind.Send(buffer, peer.endpoint)
 	if err == nil {
-		atomic.AddUint64(&peer.stats.txBytes, uint64(len(buffer)))
+		peer.txBytes.Add(uint64(len(buffer)))
 	}
 	return err
 }
@@ -174,7 +166,7 @@ func (peer *Peer) Start() {
 	peer.state.Lock()
 	defer peer.state.Unlock()
 
-	if peer.isRunning.Get() {
+	if peer.isRunning.Load() {
 		return
 	}
 
@@ -198,7 +190,7 @@ func (peer *Peer) Start() {
 	go peer.RoutineSequentialSender()
 	go peer.RoutineSequentialReceiver()
 
-	peer.isRunning.Set(true)
+	peer.isRunning.Store(true)
 }
 
 func (peer *Peer) ZeroAndFlushAll() {
@@ -210,10 +202,10 @@ func (peer *Peer) ZeroAndFlushAll() {
 	keypairs.Lock()
 	device.DeleteKeypair(keypairs.previous)
 	device.DeleteKeypair(keypairs.current)
-	device.DeleteKeypair(keypairs.loadNext())
+	device.DeleteKeypair(keypairs.next.Load())
 	keypairs.previous = nil
 	keypairs.current = nil
-	keypairs.storeNext(nil)
+	keypairs.next.Store(nil)
 	keypairs.Unlock()
 
 	// clear handshake state
@@ -238,11 +230,10 @@ func (peer *Peer) ExpireCurrentKeypairs() {
 	keypairs := &peer.keypairs
 	keypairs.Lock()
 	if keypairs.current != nil {
-		atomic.StoreUint64(&keypairs.current.sendNonce, RejectAfterMessages)
+		keypairs.current.sendNonce.Store(RejectAfterMessages)
 	}
-	if keypairs.next != nil {
-		next := keypairs.loadNext()
-		atomic.StoreUint64(&next.sendNonce, RejectAfterMessages)
+	if next := keypairs.next.Load(); next != nil {
+		next.sendNonce.Store(RejectAfterMessages)
 	}
 	keypairs.Unlock()
 }