notify/client_conn_attachment.go

package notify

import (
	"net"
	"time"
)

type clientConnAttachmentState struct {
	maxReadTimeout         time.Duration
	maxWriteTimeout        time.Duration
	authMode               AuthMode
	protectionMode         ProtectionMode
	msgEn                  func([]byte, []byte) []byte
	msgDe                  func([]byte, []byte) []byte
	fastStreamEncode       transportFastStreamEncoder
	fastBulkEncode         transportFastBulkEncoder
	fastPlainEncode        transportFastPlainEncoder
	modernPSKRuntime       *modernPSKCodecRuntime
	handshakeRsaKey        []byte
	secretKey              []byte
	keyMode                string
	sessionID              []byte
	forwardSecrecy         bool
	forwardSecrecyFallback bool
	peerAttached           bool
	peerAttachFallback     bool
	peerAttachAt           int64
	lastHeartBeat          int64
}

func cloneClientConnAttachmentState(src *clientConnAttachmentState) *clientConnAttachmentState {
	if src == nil {
		return &clientConnAttachmentState{}
	}
	cloned := *src
	cloned.handshakeRsaKey = cloneClientConnAttachmentBytes(src.handshakeRsaKey)
	cloned.secretKey = cloneClientConnAttachmentBytes(src.secretKey)
	cloned.sessionID = cloneClientConnAttachmentBytes(src.sessionID)
	return &cloned
}

func cloneClientConnAttachmentBytes(src []byte) []byte {
	if len(src) == 0 {
		return nil
	}
	return append([]byte(nil), src...)
}

func (c *LogicalConn) attachmentStateSnapshot() *clientConnAttachmentState {
	if c == nil {
		return &clientConnAttachmentState{}
	}
	if state := c.attachment.Load(); state != nil {
		if client := c.compatClientConn(); client != nil {
			client.attachment.Store(state)
		}
		return cloneClientConnAttachmentState(state)
	}
	client := c.compatClientConn()
	if client != nil {
		if state := client.attachment.Load(); state != nil {
			if c.attachment.CompareAndSwap(nil, state) {
				client.attachment.Store(state)
				return cloneClientConnAttachmentState(state)
			}
			return c.attachmentStateSnapshot()
		}
	}
	return &clientConnAttachmentState{}
}

func (c *LogicalConn) setAttachmentState(state *clientConnAttachmentState) {
	if c == nil {
		return
	}
	next := cloneClientConnAttachmentState(state)
	c.attachment.Store(next)
	if client := c.compatClientConn(); client != nil {
		client.attachment.Store(next)
	}
}

func (c *LogicalConn) updateAttachmentState(apply func(*clientConnAttachmentState)) {
	if c == nil || apply == nil {
		return
	}
	for {
		current := c.attachment.Load()
		if current == nil {
			if client := c.compatClientConn(); client != nil {
				current = client.attachment.Load()
			}
		}
		next := cloneClientConnAttachmentState(current)
		apply(next)
		if current == nil {
			if c.attachment.CompareAndSwap((*clientConnAttachmentState)(nil), next) {
				if client := c.compatClientConn(); client != nil {
					client.attachment.Store(next)
				}
				return
			}
			continue
		}
		if c.attachment.CompareAndSwap(current, next) {
			if client := c.compatClientConn(); client != nil {
				client.attachment.Store(next)
			}
			return
		}
	}
}

func (c *ClientConn) clientConnAttachmentStateSnapshot() *clientConnAttachmentState {
	if c == nil {
		return &clientConnAttachmentState{}
	}
	if logical := c.logicalView.Load(); logical != nil {
		return logical.attachmentStateSnapshot()
	}
	if state := c.attachment.Load(); state != nil {
		return cloneClientConnAttachmentState(state)
	}
	return &clientConnAttachmentState{}
}

func (c *ClientConn) setClientConnAttachmentState(state *clientConnAttachmentState) {
	if c == nil {
		return
	}
	if logical := c.logicalView.Load(); logical != nil {
		logical.setAttachmentState(state)
		return
	}
	c.attachment.Store(cloneClientConnAttachmentState(state))
}

func (c *ClientConn) updateClientConnAttachmentState(apply func(*clientConnAttachmentState)) {
	if c == nil || apply == nil {
		return
	}
	if logical := c.logicalView.Load(); logical != nil {
		logical.updateAttachmentState(apply)
		return
	}
	for {
		current := c.attachment.Load()
		next := cloneClientConnAttachmentState(current)
		apply(next)
		if current == nil {
			if c.attachment.CompareAndSwap((*clientConnAttachmentState)(nil), next) {
				return
			}
			continue
		}
		if c.attachment.CompareAndSwap(current, next) {
			return
		}
	}
}

func (c *ClientConn) applyClientConnAttachmentProfile(maxReadTimeout time.Duration, maxWriteTimeout time.Duration, msgEn func([]byte, []byte) []byte, msgDe func([]byte, []byte) []byte, handshakeRsaKey []byte, secretKey []byte) {
	c.updateClientConnAttachmentState(func(state *clientConnAttachmentState) {
		state.maxReadTimeout = maxReadTimeout
		state.maxWriteTimeout = maxWriteTimeout
		state.protectionMode = ProtectionManaged
		state.msgEn = msgEn
		state.msgDe = msgDe
		state.modernPSKRuntime = nil
		state.handshakeRsaKey = cloneClientConnAttachmentBytes(handshakeRsaKey)
		state.secretKey = cloneClientConnAttachmentBytes(secretKey)
	})
}

func (c *ClientConn) inheritClientConnAttachmentProfile(src *ClientConn) {
	if c == nil || src == nil {
		return
	}
	c.setClientConnAttachmentState(src.clientConnAttachmentStateSnapshot())
}

func (c *ClientConn) clientConnMaxReadTimeoutSnapshot() time.Duration {
	if c == nil {
		return 0
	}
	return c.clientConnAttachmentStateSnapshot().maxReadTimeout
}

func (c *ClientConn) setClientConnMaxWriteTimeout(timeout time.Duration) {
	if c == nil {
		return
	}
	if logical := c.logicalView.Load(); logical != nil {
		logical.updateAttachmentState(func(state *clientConnAttachmentState) {
			state.maxWriteTimeout = timeout
		})
		return
	}
	c.updateClientConnAttachmentState(func(state *clientConnAttachmentState) {
		state.maxWriteTimeout = timeout
	})
}

func (c *ClientConn) clientConnMaxWriteTimeoutSnapshot() time.Duration {
	if c == nil {
		return 0
	}
	return c.clientConnAttachmentStateSnapshot().maxWriteTimeout
}

func (c *ClientConn) clientConnMsgEnSnapshot() func([]byte, []byte) []byte {
	if c == nil {
		return nil
	}
	return c.clientConnAttachmentStateSnapshot().msgEn
}

func (c *ClientConn) setClientConnMsgEn(fn func([]byte, []byte) []byte) {
	c.updateClientConnAttachmentState(func(state *clientConnAttachmentState) {
		state.protectionMode = ProtectionManaged
		state.msgEn = fn
		state.fastStreamEncode = nil
		state.fastBulkEncode = nil
		state.fastPlainEncode = nil
		state.modernPSKRuntime = nil
	})
}

func (c *ClientConn) clientConnMsgDeSnapshot() func([]byte, []byte) []byte {
	if c == nil {
		return nil
	}
	return c.clientConnAttachmentStateSnapshot().msgDe
}

func (c *ClientConn) setClientConnMsgDe(fn func([]byte, []byte) []byte) {
	c.updateClientConnAttachmentState(func(state *clientConnAttachmentState) {
		state.protectionMode = ProtectionManaged
		state.msgDe = fn
		state.fastStreamEncode = nil
		state.fastBulkEncode = nil
		state.fastPlainEncode = nil
		state.modernPSKRuntime = nil
	})
}

func (c *ClientConn) setClientConnFastStreamEncode(fn transportFastStreamEncoder) {
	c.updateClientConnAttachmentState(func(state *clientConnAttachmentState) {
		state.fastStreamEncode = fn
	})
}

func (c *ClientConn) clientConnFastStreamEncodeSnapshot() transportFastStreamEncoder {
	if c == nil {
		return nil
	}
	return c.clientConnAttachmentStateSnapshot().fastStreamEncode
}

func (c *ClientConn) setClientConnFastBulkEncode(fn transportFastBulkEncoder) {
	c.updateClientConnAttachmentState(func(state *clientConnAttachmentState) {
		state.fastBulkEncode = fn
	})
}

func (c *ClientConn) clientConnFastBulkEncodeSnapshot() transportFastBulkEncoder {
	if c == nil {
		return nil
	}
	return c.clientConnAttachmentStateSnapshot().fastBulkEncode
}

func (c *ClientConn) setClientConnFastPlainEncode(fn transportFastPlainEncoder) {
	c.updateClientConnAttachmentState(func(state *clientConnAttachmentState) {
		state.fastPlainEncode = fn
	})
}

func (c *ClientConn) clientConnFastPlainEncodeSnapshot() transportFastPlainEncoder {
	if c == nil {
		return nil
	}
	return c.clientConnAttachmentStateSnapshot().fastPlainEncode
}

func (c *ClientConn) clientConnHandshakeRsaKeySnapshot() []byte {
	if c == nil {
		return nil
	}
	return c.clientConnAttachmentStateSnapshot().handshakeRsaKey
}

func (c *ClientConn) clientConnSecretKeySnapshot() []byte {
	if c == nil {
		return nil
	}
	return c.clientConnAttachmentStateSnapshot().secretKey
}

func (c *ClientConn) setClientConnSecretKey(key []byte) {
	c.updateClientConnAttachmentState(func(state *clientConnAttachmentState) {
		state.secretKey = cloneClientConnAttachmentBytes(key)
	})
}

func (c *LogicalConn) attachmentStateRaw() *clientConnAttachmentState {
	if c == nil {
		return nil
	}
	if state := c.attachment.Load(); state != nil {
		if client := c.compatClientConn(); client != nil {
			client.attachment.Store(state)
		}
		return state
	}
	if client := c.compatClientConn(); client != nil {
		if state := client.attachment.Load(); state != nil {
			if c.attachment.CompareAndSwap(nil, state) {
				client.attachment.Store(state)
				return state
			}
			return c.attachmentStateRaw()
		}
	}
	return nil
}

func (c *LogicalConn) modernPSKRuntimeSnapshot() *modernPSKCodecRuntime {
	if state := c.attachmentStateRaw(); state != nil {
		return state.modernPSKRuntime
	}
	return nil
}

func (c *LogicalConn) protectionModeSnapshot() ProtectionMode {
	if state := c.attachmentStateRaw(); state != nil {
		return state.protectionMode
	}
	return ProtectionManaged
}

func (c *LogicalConn) authModeSnapshot() AuthMode {
	if state := c.attachmentStateRaw(); state != nil {
		return state.authMode
	}
	return AuthNone
}

func (c *LogicalConn) peerAttachAuthenticatedSnapshot() (bool, bool, time.Time) {
	if state := c.attachmentStateRaw(); state != nil {
		if state.peerAttachAt == 0 {
			return state.peerAttached, state.peerAttachFallback, time.Time{}
		}
		return state.peerAttached, state.peerAttachFallback, time.Unix(0, state.peerAttachAt)
	}
	return false, false, time.Time{}
}

func (c *LogicalConn) markPeerAttachAuthenticated(authMode AuthMode, fallback bool, at time.Time) {
	c.updateAttachmentState(func(state *clientConnAttachmentState) {
		state.authMode = authMode
		state.peerAttached = true
		state.peerAttachFallback = fallback
		state.peerAttachAt = at.UnixNano()
	})
}

func (c *LogicalConn) setModernPSKRuntime(runtime *modernPSKCodecRuntime) {
	c.updateAttachmentState(func(state *clientConnAttachmentState) {
		state.modernPSKRuntime = runtime
	})
}

func (c *ClientConn) clientConnAttachmentStateRaw() *clientConnAttachmentState {
	if c == nil {
		return nil
	}
	if logical := c.logicalView.Load(); logical != nil {
		return logical.attachmentStateRaw()
	}
	return c.attachment.Load()
}

func (c *ClientConn) clientConnModernPSKRuntimeSnapshot() *modernPSKCodecRuntime {
	if state := c.clientConnAttachmentStateRaw(); state != nil {
		return state.modernPSKRuntime
	}
	return nil
}

func (c *ClientConn) setClientConnModernPSKRuntime(runtime *modernPSKCodecRuntime) {
	c.updateClientConnAttachmentState(func(state *clientConnAttachmentState) {
		state.modernPSKRuntime = runtime
	})
}

func (c *ClientConn) clientConnLastHeartbeatUnixSnapshot() int64 {
	if c == nil {
		return 0
	}
	return c.clientConnAttachmentStateSnapshot().lastHeartBeat
}

func (c *ClientConn) setClientConnLastHeartbeatUnix(unix int64) {
	if c == nil {
		return
	}
	if logical := c.logicalView.Load(); logical != nil {
		logical.setClientConnLastHeartbeatUnix(unix)
		return
	}
	c.updateClientConnAttachmentState(func(state *clientConnAttachmentState) {
		state.lastHeartBeat = unix
	})
}

func (c *ClientConn) markClientConnHeartbeatNow() {
	if c == nil {
		return
	}
	if logical := c.logicalView.Load(); logical != nil {
		logical.markHeartbeatNow()
		return
	}
	c.setClientConnLastHeartbeatUnix(time.Now().Unix())
}

func (c *ClientConn) setClientConnRemoteAddr(addr net.Addr) {
	if c == nil {
		return
	}
	state := c.ensureLogicalConnState()
	if state == nil {
		c.ClientAddr = addr
		return
	}
	state.updatePeer(func(peer *logicalConnPeerState) {
		peer.clientAddr = addr
	})
	c.syncLegacyLogicalFieldsFromState(state)
}