package server import ( "fmt" cache "github.com/pmylund/go-cache" "math/rand" "net" "strings" "sync" "time" ) const ( cacheDuration time.Duration = 24 * time.Hour cacheCleanupInterval time.Duration = time.Minute ) /** * TunnelManager: Manages a set of tunnels */ type TunnelManager struct { domain string tunnels map[string]*Tunnel idDomainAffinity *cache.Cache ipDomainAffinity *cache.Cache sync.RWMutex } func NewTunnelManager(domain string) *TunnelManager { return &TunnelManager{ domain: domain, tunnels: make(map[string]*Tunnel), idDomainAffinity: cache.New(cacheDuration, cacheCleanupInterval), ipDomainAffinity: cache.New(cacheDuration, cacheCleanupInterval), } } func (m *TunnelManager) Add(t *Tunnel) error { assignTunnel := func(url string) bool { m.Lock() defer m.Unlock() if m.tunnels[url] == nil { m.tunnels[url] = t return true } return false } url := "" switch t.regMsg.Protocol { case "tcp": addr := t.listener.Addr().(*net.TCPAddr) url = fmt.Sprintf("tcp://%s:%d", m.domain, addr.Port) if !assignTunnel(url) { return t.Error("TCP at %s already registered!", url) } metrics.tcpTunnelMeter.Mark(1) case "http": if strings.TrimSpace(t.regMsg.Hostname) != "" { url = fmt.Sprintf("http://%s", t.regMsg.Hostname) } else if strings.TrimSpace(t.regMsg.Subdomain) != "" { url = fmt.Sprintf("http://%s.%s", t.regMsg.Subdomain, m.domain) } if url != "" { if !assignTunnel(url) { return t.Warn("The tunnel address %s is already registered!", url) } } else { clientIp := t.ctl.conn.RemoteAddr().(*net.TCPAddr).IP.String() clientId := t.regMsg.ClientId // try to give the same subdomain back if it's available subdomain := fmt.Sprintf("%x", rand.Int31()) if lastDomain, ok := m.idDomainAffinity.Get(clientId); ok { t.Debug("Found affinity for subdomain %s with client id %s", subdomain, clientId) subdomain = lastDomain.(string) } else if lastDomain, ok = m.ipDomainAffinity.Get(clientIp); ok { t.Debug("Found affinity for subdomain %s with client ip %s", subdomain, clientIp) subdomain = lastDomain.(string) } // pick one randomly for { url = fmt.Sprintf("http://%s.%s", subdomain, m.domain) if assignTunnel(url) { break } else { subdomain = fmt.Sprintf("%x", rand.Int31()) } } // save our choice for later // XXX: this is going to leak memory m.idDomainAffinity.Set(clientId, subdomain, 0) m.ipDomainAffinity.Set(clientIp, subdomain, 0) } default: return t.Error("Unrecognized protocol type %s", t.regMsg.Protocol) } t.url = url metrics.tunnelMeter.Mark(1) //metrics.tunnelGauge.Update(int64(len(m.tunnels))) switch t.regMsg.OS { case "windows": metrics.windowsCounter.Inc(1) case "linux": metrics.linuxCounter.Inc(1) case "darwin": metrics.osxCounter.Inc(1) default: metrics.otherCounter.Inc(1) } return nil } func (m *TunnelManager) Del(url string) { m.Lock() defer m.Unlock() delete(m.tunnels, url) } func (m *TunnelManager) Get(url string) *Tunnel { m.RLock() defer m.RUnlock() return m.tunnels[url] }