94 lines
2.7 KiB
Go
94 lines
2.7 KiB
Go
package tunnelhostnamemapper
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/cloudflare/cloudflared/h2mux"
|
|
"github.com/cloudflare/cloudflared/originservice"
|
|
"github.com/cloudflare/cloudflared/tunnelrpc/pogs"
|
|
)
|
|
|
|
// TunnelHostnameMapper maps TunnelHostname to an OriginService
|
|
type TunnelHostnameMapper struct {
|
|
sync.RWMutex
|
|
tunnelHostnameToOrigin map[h2mux.TunnelHostname]originservice.OriginService
|
|
}
|
|
|
|
func NewTunnelHostnameMapper() *TunnelHostnameMapper {
|
|
return &TunnelHostnameMapper{
|
|
tunnelHostnameToOrigin: make(map[h2mux.TunnelHostname]originservice.OriginService),
|
|
}
|
|
}
|
|
|
|
// Get an OriginService given a TunnelHostname
|
|
func (om *TunnelHostnameMapper) Get(key h2mux.TunnelHostname) (originservice.OriginService, bool) {
|
|
om.RLock()
|
|
defer om.RUnlock()
|
|
originService, ok := om.tunnelHostnameToOrigin[key]
|
|
return originService, ok
|
|
}
|
|
|
|
// Add a mapping. If there is already an OriginService with this key, shutdown the old origin service and replace it
|
|
// with the new one
|
|
func (om *TunnelHostnameMapper) Add(key h2mux.TunnelHostname, os originservice.OriginService) {
|
|
om.Lock()
|
|
defer om.Unlock()
|
|
if oldOS, ok := om.tunnelHostnameToOrigin[key]; ok {
|
|
oldOS.Shutdown()
|
|
}
|
|
om.tunnelHostnameToOrigin[key] = os
|
|
}
|
|
|
|
// Delete a mapping, and shutdown its OriginService
|
|
func (om *TunnelHostnameMapper) Delete(key h2mux.TunnelHostname) (keyFound bool) {
|
|
om.Lock()
|
|
defer om.Unlock()
|
|
if os, ok := om.tunnelHostnameToOrigin[key]; ok {
|
|
os.Shutdown()
|
|
delete(om.tunnelHostnameToOrigin, key)
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
// ToRemove finds all keys that should be removed from the TunnelHostnameMapper.
|
|
func (om *TunnelHostnameMapper) ToRemove(newConfigs []*pogs.ReverseProxyConfig) (toRemove []h2mux.TunnelHostname) {
|
|
om.Lock()
|
|
defer om.Unlock()
|
|
|
|
// Convert into a set, for O(1) lookups instead of O(n)
|
|
newConfigSet := toSet(newConfigs)
|
|
|
|
// If a config in `om` isn't in `newConfigs`, it must be removed.
|
|
for hostname := range om.tunnelHostnameToOrigin {
|
|
if _, ok := newConfigSet[hostname]; !ok {
|
|
toRemove = append(toRemove, hostname)
|
|
}
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
// ToAdd filters the given configs, keeping those that should be added to the TunnelHostnameMapper.
|
|
func (om *TunnelHostnameMapper) ToAdd(newConfigs []*pogs.ReverseProxyConfig) (toAdd []*pogs.ReverseProxyConfig) {
|
|
om.Lock()
|
|
defer om.Unlock()
|
|
|
|
// If a config in `newConfigs` isn't in `om`, it must be added.
|
|
for _, config := range newConfigs {
|
|
if _, ok := om.tunnelHostnameToOrigin[config.TunnelHostname]; !ok {
|
|
toAdd = append(toAdd, config)
|
|
}
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
func toSet(configs []*pogs.ReverseProxyConfig) map[h2mux.TunnelHostname]*pogs.ReverseProxyConfig {
|
|
m := make(map[h2mux.TunnelHostname]*pogs.ReverseProxyConfig)
|
|
for _, config := range configs {
|
|
m[config.TunnelHostname] = config
|
|
}
|
|
return m
|
|
}
|