ngrok/server/manager.go

package server

import (
	"fmt"
	cache "github.com/pmylund/go-cache"
	"math/rand"
	"net"
	"strings"
	"sync"
	"time"
)

/**
 * TunnelManager: Manages a set of tunnels
 */
type TunnelManager struct {
	domain         string
	tunnels        map[string]*Tunnel
	domainAffinity *cache.Cache
	sync.RWMutex
}

func NewTunnelManager(domain string) *TunnelManager {
	return &TunnelManager{
		domain:         domain,
		tunnels:        make(map[string]*Tunnel),
		domainAffinity: cache.New(24*time.Hour, time.Minute),
	}
}

func (m *TunnelManager) Add(t *Tunnel) {
	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) {
			panic("TCP at %s already registered!")
		}
		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) {
				panic(fmt.Sprintf("The tunnel address %s is already registered!", url))
			}
		} else {
			// try to give the same subdomain back if it's available
			subdomain, ok := m.domainAffinity.Get(t.regMsg.ClientId)
			if !ok {
				subdomain = fmt.Sprintf("%x", rand.Int31())
			}

			// 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.domainAffinity.Set(t.regMsg.ClientId, subdomain, 0)
		}

	default:
		panic(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)
	}
}

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]
}