conntracker.go 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677
  1. package tunnelstate
  2. import (
  3. "sync"
  4. "github.com/rs/zerolog"
  5. "github.com/cloudflare/cloudflared/connection"
  6. )
  7. type ConnTracker struct {
  8. mutex sync.RWMutex
  9. // int is the connection Index
  10. connectionInfo map[uint8]ConnectionInfo
  11. log *zerolog.Logger
  12. }
  13. type ConnectionInfo struct {
  14. IsConnected bool
  15. Protocol connection.Protocol
  16. }
  17. func NewConnTracker(
  18. log *zerolog.Logger,
  19. ) *ConnTracker {
  20. return &ConnTracker{
  21. connectionInfo: make(map[uint8]ConnectionInfo, 0),
  22. log: log,
  23. }
  24. }
  25. func (ct *ConnTracker) OnTunnelEvent(c connection.Event) {
  26. switch c.EventType {
  27. case connection.Connected:
  28. ct.mutex.Lock()
  29. ci := ConnectionInfo{
  30. IsConnected: true,
  31. Protocol: c.Protocol,
  32. }
  33. ct.connectionInfo[c.Index] = ci
  34. ct.mutex.Unlock()
  35. case connection.Disconnected, connection.Reconnecting, connection.RegisteringTunnel, connection.Unregistering:
  36. ct.mutex.Lock()
  37. ci := ct.connectionInfo[c.Index]
  38. ci.IsConnected = false
  39. ct.connectionInfo[c.Index] = ci
  40. ct.mutex.Unlock()
  41. default:
  42. ct.log.Error().Msgf("Unknown connection event case %v", c)
  43. }
  44. }
  45. func (ct *ConnTracker) CountActiveConns() uint {
  46. ct.mutex.RLock()
  47. defer ct.mutex.RUnlock()
  48. active := uint(0)
  49. for _, ci := range ct.connectionInfo {
  50. if ci.IsConnected {
  51. active++
  52. }
  53. }
  54. return active
  55. }
  56. // HasConnectedWith checks if we've ever had a successful connection to the edge
  57. // with said protocol.
  58. func (ct *ConnTracker) HasConnectedWith(protocol connection.Protocol) bool {
  59. ct.mutex.RLock()
  60. defer ct.mutex.RUnlock()
  61. for _, ci := range ct.connectionInfo {
  62. if ci.Protocol == protocol {
  63. return true
  64. }
  65. }
  66. return false
  67. }