package nsqd
import (
"bufio"
"compress/flate"
"crypto/tls"
"fmt"
"net"
"sync"
"sync/atomic"
"time"
"github.com/mreiferson/go-snappystream"
"github.com/nsqio/nsq/internal/auth"
)
const defaultBufferSize = 16 * 1024
const (
stateInit = iota
stateDisconnected
stateConnected
stateSubscribed
stateClosing
)
type identifyDataV2 struct { ShortID string `json:"short_id"` // TODO: deprecated, remove in 1.0
LongID string `json:"long_id"` // TODO: deprecated, remove in 1.0
ClientID string `json:"client_id"`
Hostname string `json:"hostname"`
HeartbeatInterval int `json:"heartbeat_interval"`
OutputBufferSize int `json:"output_buffer_size"`
OutputBufferTimeout int `json:"output_buffer_timeout"`
FeatureNegotiation bool `json:"feature_negotiation"`
TLSv1 bool `json:"tls_v1"`
Deflate bool `json:"deflate"`
DeflateLevel int `json:"deflate_level"`
Snappy bool `json:"snappy"`
SampleRate int32 `json:"sample_rate"`
UserAgent string `json:"user_agent"`
MsgTimeout int `json:"msg_timeout"`
}
type identifyEvent struct { OutputBufferTimeout time.Duration
HeartbeatInterval time.Duration
SampleRate int32
MsgTimeout time.Duration
}
type clientV2 struct { // 64bit atomic vars need to be first for proper alignment on 32bit platforms
ReadyCount int64
InFlightCount int64
MessageCount uint64
FinishCount uint64
RequeueCount uint64
writeLock sync.RWMutex
metaLock sync.RWMutex
ID int64
ctx *context
UserAgent string
// original connection
net.Conn
// connections based on negotiated features
tlsConn *tls.Conn
flateWriter *flate.Writer
// reading/writing interfaces
Reader *bufio.Reader
Writer *bufio.Writer
OutputBufferSize int
OutputBufferTimeout time.Duration
HeartbeatInterval time.Duration
MsgTimeout time.Duration
State int32
ConnectTime time.Time
Channel *Channel
ReadyStateChan chan int
ExitChan chan int
ClientID string
Hostname string
SampleRate int32
IdentifyEventChan chan identifyEvent
SubEventChan chan *Channel
TLS int32
Snappy int32
Deflate int32
// re-usable buffer for reading the 4-byte lengths off the wire
lenBuf [4]byte
lenSlice []byte
AuthSecret string
AuthState *auth.State
}
func newClientV2(id int64, conn net.Conn, ctx *context) *clientV2 { var identifier string
if conn != nil { identifier, _, _ = net.SplitHostPort(conn.RemoteAddr().String())
}
c := &clientV2{ ID: id,
ctx: ctx,
Conn: conn,
Reader: bufio.NewReaderSize(conn, defaultBufferSize),
Writer: bufio.NewWriterSize(conn, defaultBufferSize),
OutputBufferSize: defaultBufferSize,
OutputBufferTimeout: 250 * time.Millisecond,
MsgTimeout: ctx.nsqd.getOpts().MsgTimeout,
// ReadyStateChan has a buffer of 1 to guarantee that in the event
// there is a race the state update is not lost
ReadyStateChan: make(chan int, 1),
ExitChan: make(chan int),
ConnectTime: time.Now(),
State: stateInit,
ClientID: identifier,
Hostname: identifier,
SubEventChan: make(chan *Channel, 1),
IdentifyEventChan: make(chan identifyEvent, 1),
// heartbeats are client configurable but default to 30s
HeartbeatInterval: ctx.nsqd.getOpts().ClientTimeout / 2,
}
c.lenSlice = c.lenBuf[:]
return c
}
func (c *clientV2) String() string { return c.RemoteAddr().String()
}
func (c *clientV2) Identify(data identifyDataV2) error { c.ctx.nsqd.logf("[%s] IDENTIFY: %+v", c, data)
// TODO: for backwards compatibility, remove in 1.0
hostname := data.Hostname
if hostname == "" { hostname = data.LongID
}
// TODO: for backwards compatibility, remove in 1.0
clientID := data.ClientID
if clientID == "" { clientID = data.ShortID
}
c.metaLock.Lock()
c.ClientID = clientID
c.Hostname = hostname
c.UserAgent = data.UserAgent
c.metaLock.Unlock()
err := c.SetHeartbeatInterval(data.HeartbeatInterval)
if err != nil { return err
}
err = c.SetOutputBufferSize(data.OutputBufferSize)
if err != nil { return err
}
err = c.SetOutputBufferTimeout(data.OutputBufferTimeout)
if err != nil { return err
}
err = c.SetSampleRate(data.SampleRate)
if err != nil { return err
}
err = c.SetMsgTimeout(data.MsgTimeout)
if err != nil { return err
}
ie := identifyEvent{ OutputBufferTimeout: c.OutputBufferTimeout,
HeartbeatInterval: c.HeartbeatInterval,
SampleRate: c.SampleRate,
MsgTimeout: c.MsgTimeout,
}
// update the client's message pump
select { case c.IdentifyEventChan <- ie:
default:
}
return nil
}
func (c *clientV2) Stats() ClientStats { c.metaLock.RLock()
// TODO: deprecated, remove in 1.0
name := c.ClientID
clientID := c.ClientID
hostname := c.Hostname
userAgent := c.UserAgent
var identity string
var identityURL string
if c.AuthState != nil { identity = c.AuthState.Identity
identityURL = c.AuthState.IdentityURL
}
c.metaLock.RUnlock()
stats := ClientStats{ // TODO: deprecated, remove in 1.0
Name: name,
Version: "V2",
RemoteAddress: c.RemoteAddr().String(),
ClientID: clientID,
Hostname: hostname,
UserAgent: userAgent,
State: atomic.LoadInt32(&c.State),
ReadyCount: atomic.LoadInt64(&c.ReadyCount),
InFlightCount: atomic.LoadInt64(&c.InFlightCount),
MessageCount: atomic.LoadUint64(&c.MessageCount),
FinishCount: atomic.LoadUint64(&c.FinishCount),
RequeueCount: atomic.LoadUint64(&c.RequeueCount),
ConnectTime: c.ConnectTime.Unix(),
SampleRate: atomic.LoadInt32(&c.SampleRate),
TLS: atomic.LoadInt32(&c.TLS) == 1,
Deflate: atomic.LoadInt32(&c.Deflate) == 1,
Snappy: atomic.LoadInt32(&c.Snappy) == 1,
Authed: c.HasAuthorizations(),
AuthIdentity: identity,
AuthIdentityURL: identityURL,
}
if stats.TLS { p := prettyConnectionState{c.tlsConn.ConnectionState()} stats.CipherSuite = p.GetCipherSuite()
stats.TLSVersion = p.GetVersion()
stats.TLSNegotiatedProtocol = p.NegotiatedProtocol
stats.TLSNegotiatedProtocolIsMutual = p.NegotiatedProtocolIsMutual
}
return stats
}
// struct to convert from integers to the human readable strings
type prettyConnectionState struct { tls.ConnectionState
}
func (p *prettyConnectionState) GetCipherSuite() string { switch p.CipherSuite { case tls.TLS_RSA_WITH_RC4_128_SHA:
return "TLS_RSA_WITH_RC4_128_SHA"
case tls.TLS_RSA_WITH_3DES_EDE_CBC_SHA:
return "TLS_RSA_WITH_3DES_EDE_CBC_SHA"
case tls.TLS_RSA_WITH_AES_128_CBC_SHA:
return "TLS_RSA_WITH_AES_128_CBC_SHA"
case tls.TLS_RSA_WITH_AES_256_CBC_SHA:
return "TLS_RSA_WITH_AES_256_CBC_SHA"
case tls.TLS_ECDHE_ECDSA_WITH_RC4_128_SHA:
return "TLS_ECDHE_ECDSA_WITH_RC4_128_SHA"
case tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA:
return "TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA"
case tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA:
return "TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA"
case tls.TLS_ECDHE_RSA_WITH_RC4_128_SHA:
return "TLS_ECDHE_RSA_WITH_RC4_128_SHA"
case tls.TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA:
return "TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA"
case tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA:
return "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA"
case tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA:
return "TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA"
case tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256:
return "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256"
case tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256:
return "TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256"
}
return fmt.Sprintf("Unknown %d", p.CipherSuite)}
func (p *prettyConnectionState) GetVersion() string { switch p.Version { case tls.VersionSSL30:
return "SSL30"
case tls.VersionTLS10:
return "TLS1.0"
case tls.VersionTLS11:
return "TLS1.1"
case tls.VersionTLS12:
return "TLS1.2"
default:
return fmt.Sprintf("Unknown %d", p.Version) }
}
func (c *clientV2) IsReadyForMessages() bool { if c.Channel.IsPaused() { return false
}
readyCount := atomic.LoadInt64(&c.ReadyCount)
inFlightCount := atomic.LoadInt64(&c.InFlightCount)
if c.ctx.nsqd.getOpts().Verbose { c.ctx.nsqd.logf("[%s] state rdy: %4d inflt: %4d", c, readyCount, inFlightCount)
}
if inFlightCount >= readyCount || readyCount <= 0 { return false
}
return true
}
func (c *clientV2) SetReadyCount(count int64) { atomic.StoreInt64(&c.ReadyCount, count)
c.tryUpdateReadyState()
}
func (c *clientV2) tryUpdateReadyState() { // you can always *try* to write to ReadyStateChan because in the cases
// where you cannot the message pump loop would have iterated anyway.
// the atomic integer operations guarantee correctness of the value.
select { case c.ReadyStateChan <- 1:
default:
}
}
func (c *clientV2) FinishedMessage() { atomic.AddUint64(&c.FinishCount, 1)
atomic.AddInt64(&c.InFlightCount, -1)
c.tryUpdateReadyState()
}
func (c *clientV2) Empty() { atomic.StoreInt64(&c.InFlightCount, 0)
c.tryUpdateReadyState()
}
func (c *clientV2) SendingMessage() { atomic.AddInt64(&c.InFlightCount, 1)
atomic.AddUint64(&c.MessageCount, 1)
}
func (c *clientV2) TimedOutMessage() { atomic.AddInt64(&c.InFlightCount, -1)
c.tryUpdateReadyState()
}
func (c *clientV2) RequeuedMessage() { atomic.AddUint64(&c.RequeueCount, 1)
atomic.AddInt64(&c.InFlightCount, -1)
c.tryUpdateReadyState()
}
func (c *clientV2) StartClose() { // Force the client into ready 0
c.SetReadyCount(0)
// mark this client as closing
atomic.StoreInt32(&c.State, stateClosing)
}
func (c *clientV2) Pause() { c.tryUpdateReadyState()
}
func (c *clientV2) UnPause() { c.tryUpdateReadyState()
}
func (c *clientV2) SetHeartbeatInterval(desiredInterval int) error { c.writeLock.Lock()
defer c.writeLock.Unlock()
switch { case desiredInterval == -1:
c.HeartbeatInterval = 0
case desiredInterval == 0:
// do nothing (use default)
case desiredInterval >= 1000 &&
desiredInterval <= int(c.ctx.nsqd.getOpts().MaxHeartbeatInterval/time.Millisecond):
c.HeartbeatInterval = time.Duration(desiredInterval) * time.Millisecond
default:
return fmt.Errorf("heartbeat interval (%d) is invalid", desiredInterval) }
return nil
}
func (c *clientV2) SetOutputBufferSize(desiredSize int) error { var size int
switch { case desiredSize == -1:
// effectively no buffer (every write will go directly to the wrapped net.Conn)
size = 1
case desiredSize == 0:
// do nothing (use default)
case desiredSize >= 64 && desiredSize <= int(c.ctx.nsqd.getOpts().MaxOutputBufferSize):
size = desiredSize
default:
return fmt.Errorf("output buffer size (%d) is invalid", desiredSize) }
if size > 0 { c.writeLock.Lock()
defer c.writeLock.Unlock()
c.OutputBufferSize = size
err := c.Writer.Flush()
if err != nil { return err
}
c.Writer = bufio.NewWriterSize(c.Conn, size)
}
return nil
}
func (c *clientV2) SetOutputBufferTimeout(desiredTimeout int) error { c.writeLock.Lock()
defer c.writeLock.Unlock()
switch { case desiredTimeout == -1:
c.OutputBufferTimeout = 0
case desiredTimeout == 0:
// do nothing (use default)
case desiredTimeout >= 1 &&
desiredTimeout <= int(c.ctx.nsqd.getOpts().MaxOutputBufferTimeout/time.Millisecond):
c.OutputBufferTimeout = time.Duration(desiredTimeout) * time.Millisecond
default:
return fmt.Errorf("output buffer timeout (%d) is invalid", desiredTimeout) }
return nil
}
func (c *clientV2) SetSampleRate(sampleRate int32) error { if sampleRate < 0 || sampleRate > 99 { return fmt.Errorf("sample rate (%d) is invalid", sampleRate) }
atomic.StoreInt32(&c.SampleRate, sampleRate)
return nil
}
func (c *clientV2) SetMsgTimeout(msgTimeout int) error { c.writeLock.Lock()
defer c.writeLock.Unlock()
switch { case msgTimeout == 0:
// do nothing (use default)
case msgTimeout >= 1000 &&
msgTimeout <= int(c.ctx.nsqd.getOpts().MaxMsgTimeout/time.Millisecond):
c.MsgTimeout = time.Duration(msgTimeout) * time.Millisecond
default:
return fmt.Errorf("msg timeout (%d) is invalid", msgTimeout) }
return nil
}
func (c *clientV2) UpgradeTLS() error { c.writeLock.Lock()
defer c.writeLock.Unlock()
tlsConn := tls.Server(c.Conn, c.ctx.nsqd.tlsConfig)
tlsConn.SetDeadline(time.Now().Add(5 * time.Second))
err := tlsConn.Handshake()
if err != nil { return err
}
c.tlsConn = tlsConn
c.Reader = bufio.NewReaderSize(c.tlsConn, defaultBufferSize)
c.Writer = bufio.NewWriterSize(c.tlsConn, c.OutputBufferSize)
atomic.StoreInt32(&c.TLS, 1)
return nil
}
func (c *clientV2) UpgradeDeflate(level int) error { c.writeLock.Lock()
defer c.writeLock.Unlock()
conn := c.Conn
if c.tlsConn != nil { conn = c.tlsConn
}
c.Reader = bufio.NewReaderSize(flate.NewReader(conn), defaultBufferSize)
fw, _ := flate.NewWriter(conn, level)
c.flateWriter = fw
c.Writer = bufio.NewWriterSize(fw, c.OutputBufferSize)
atomic.StoreInt32(&c.Deflate, 1)
return nil
}
func (c *clientV2) UpgradeSnappy() error { c.writeLock.Lock()
defer c.writeLock.Unlock()
conn := c.Conn
if c.tlsConn != nil { conn = c.tlsConn
}
c.Reader = bufio.NewReaderSize(snappystream.NewReader(conn, snappystream.SkipVerifyChecksum), defaultBufferSize)
c.Writer = bufio.NewWriterSize(snappystream.NewWriter(conn), c.OutputBufferSize)
atomic.StoreInt32(&c.Snappy, 1)
return nil
}
func (c *clientV2) Flush() error { var zeroTime time.Time
if c.HeartbeatInterval > 0 { c.SetWriteDeadline(time.Now().Add(c.HeartbeatInterval))
} else { c.SetWriteDeadline(zeroTime)
}
err := c.Writer.Flush()
if err != nil { return err
}
if c.flateWriter != nil { return c.flateWriter.Flush()
}
return nil
}
func (c *clientV2) QueryAuthd() error { remoteIP, _, err := net.SplitHostPort(c.String())
if err != nil { return err
}
tls := atomic.LoadInt32(&c.TLS) == 1
tlsEnabled := "false"
if tls { tlsEnabled = "true"
}
authState, err := auth.QueryAnyAuthd(c.ctx.nsqd.getOpts().AuthHTTPAddresses,
remoteIP, tlsEnabled, c.AuthSecret, c.ctx.nsqd.getOpts().HTTPClientConnectTimeout,
c.ctx.nsqd.getOpts().HTTPClientRequestTimeout)
if err != nil { return err
}
c.AuthState = authState
return nil
}
func (c *clientV2) Auth(secret string) error { c.AuthSecret = secret
return c.QueryAuthd()
}
func (c *clientV2) IsAuthorized(topic, channel string) (bool, error) { if c.AuthState == nil { return false, nil
}
if c.AuthState.IsExpired() { err := c.QueryAuthd()
if err != nil { return false, err
}
}
if c.AuthState.IsAllowed(topic, channel) { return true, nil
}
return false, nil
}
func (c *clientV2) HasAuthorizations() bool { if c.AuthState != nil { return len(c.AuthState.Authorizations) != 0
}
return false
}
