github.com/Tyktechnologies/tyk@v2.9.5+incompatible/dnscache/manager.go (about)

     1  package dnscache
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"math/rand"
     7  	"net"
     8  	"time"
     9  
    10  	"github.com/TykTechnologies/tyk/config"
    11  
    12  	"github.com/sirupsen/logrus"
    13  
    14  	"github.com/TykTechnologies/tyk/log"
    15  )
    16  
    17  var (
    18  	logger = log.Get().WithField("prefix", "dnscache")
    19  )
    20  
    21  type DialContextFunc func(ctx context.Context, network, address string) (net.Conn, error)
    22  
    23  // IDnsCacheManager is an interface for abstracting interaction with dns cache. Implemented by DnsCacheManager
    24  type IDnsCacheManager interface {
    25  	InitDNSCaching(ttl, checkInterval time.Duration)
    26  	WrapDialer(dialer *net.Dialer) DialContextFunc
    27  	SetCacheStorage(cache IDnsCacheStorage)
    28  	CacheStorage() IDnsCacheStorage
    29  	IsCacheEnabled() bool
    30  	DisposeCache()
    31  }
    32  
    33  // IDnsCacheStorage is an interface for working with cached storage of dns record.
    34  // Wrapped by IDnsCacheManager/DnsCacheManager. Implemented by DnsCacheStorage
    35  type IDnsCacheStorage interface {
    36  	FetchItem(key string) ([]string, error)
    37  	Get(key string) (DnsCacheItem, bool)
    38  	Set(key string, addrs []string)
    39  	Delete(key string)
    40  	Clear()
    41  }
    42  
    43  // DnsCacheManager is responsible for in-memory dns query records cache.
    44  // It allows to init dns caching and to hook into net/http dns resolution chain in order to cache query response ip records.
    45  type DnsCacheManager struct {
    46  	cacheStorage IDnsCacheStorage
    47  	strategy     config.IPsHandleStrategy
    48  	rand         *rand.Rand
    49  }
    50  
    51  // NewDnsCacheManager returns new empty/non-initialized DnsCacheManager
    52  func NewDnsCacheManager(multipleIPsHandleStrategy config.IPsHandleStrategy) *DnsCacheManager {
    53  	manager := &DnsCacheManager{nil, multipleIPsHandleStrategy, nil}
    54  	return manager
    55  }
    56  
    57  func (m *DnsCacheManager) SetCacheStorage(cache IDnsCacheStorage) {
    58  	m.cacheStorage = cache
    59  }
    60  
    61  func (m *DnsCacheManager) CacheStorage() IDnsCacheStorage {
    62  	return m.cacheStorage
    63  }
    64  
    65  func (m *DnsCacheManager) IsCacheEnabled() bool {
    66  	return m.cacheStorage != nil
    67  }
    68  
    69  // WrapDialer returns wrapped version of net.Dialer#DialContext func with hooked up caching of dns queries.
    70  //
    71  // Actual dns server call occures in net.Resolver#LookupIPAddr method,
    72  // linked to net.Dialer instance by net.Dialer#Resolver field
    73  func (m *DnsCacheManager) WrapDialer(dialer *net.Dialer) DialContextFunc {
    74  	return func(ctx context.Context, network, address string) (net.Conn, error) {
    75  		return m.doCachedDial(dialer, ctx, network, address)
    76  	}
    77  }
    78  
    79  func (m *DnsCacheManager) doCachedDial(d *net.Dialer, ctx context.Context, network, address string) (net.Conn, error) {
    80  	safeDial := func(addr string, itemKey string) (net.Conn, error) {
    81  		conn, err := d.DialContext(ctx, network, addr)
    82  		if err != nil && itemKey != "" {
    83  			m.cacheStorage.Delete(itemKey)
    84  		}
    85  		return conn, err
    86  	}
    87  
    88  	if !m.IsCacheEnabled() {
    89  		return safeDial(address, "")
    90  	}
    91  
    92  	host, port, err := net.SplitHostPort(address)
    93  	if err != nil {
    94  		return nil, err
    95  	}
    96  
    97  	if ip := net.ParseIP(host); ip != nil {
    98  		return safeDial(address, "")
    99  	}
   100  
   101  	ips, err := m.cacheStorage.FetchItem(host)
   102  	if err != nil {
   103  		logger.WithError(err).WithFields(logrus.Fields{
   104  			"network": network,
   105  			"address": address,
   106  		}).Errorf("doCachedDial cachedStorage.FetchItem error. ips=%v", ips)
   107  
   108  		return safeDial(address, "")
   109  	}
   110  
   111  	if m.strategy == config.NoCacheStrategy {
   112  		if len(ips) > 1 {
   113  			m.cacheStorage.Delete(host)
   114  			return safeDial(ips[0]+":"+port, "")
   115  		}
   116  	}
   117  
   118  	if m.strategy == config.RandomStrategy {
   119  		if len(ips) > 1 {
   120  			ip, _ := m.getRandomIp(ips)
   121  			return safeDial(ip+":"+port, host)
   122  		}
   123  		return safeDial(ips[0]+":"+port, host)
   124  	}
   125  
   126  	return safeDial(ips[0]+":"+port, host)
   127  }
   128  
   129  func (m *DnsCacheManager) getRandomIp(ips []string) (string, error) {
   130  	if m.strategy != config.RandomStrategy {
   131  		return "", fmt.Errorf(
   132  			"getRandomIp can be called only with %v strategy. strategy=%v",
   133  			config.RandomStrategy, m.strategy)
   134  	}
   135  
   136  	if m.rand == nil {
   137  		source := rand.NewSource(time.Now().Unix())
   138  		m.rand = rand.New(source)
   139  	}
   140  
   141  	ip := ips[m.rand.Intn(len(ips))]
   142  
   143  	return ip, nil
   144  }
   145  
   146  // InitDNSCaching initializes manager's cache storage if it wasn't initialized before with provided ttl, checkinterval values
   147  // Initialized cache storage enables caching of previously hoooked net.Dialer DialContext calls
   148  //
   149  // Otherwise leave storage as is.
   150  func (m *DnsCacheManager) InitDNSCaching(ttl, checkInterval time.Duration) {
   151  	if !m.IsCacheEnabled() {
   152  		logger.Infof("Initializing dns cache with ttl=%s, duration=%s", ttl, checkInterval)
   153  		storage := NewDnsCacheStorage(ttl, checkInterval)
   154  		m.SetCacheStorage(IDnsCacheStorage(storage))
   155  	}
   156  }
   157  
   158  // DisposeCache clear all entries from cache and disposes/disables caching of dns queries
   159  func (m *DnsCacheManager) DisposeCache() {
   160  	m.cacheStorage.Clear()
   161  	m.cacheStorage = nil
   162  }