go_subnetter/netsplit/funcs_cache.go

package netsplit

import (
	`encoding/json`
	`errors`
	`io/fs`
	`net/http`
	`net/netip`
	`os`
	`path/filepath`
	`sync`

	`github.com/go-resty/resty/v2`
	`r00t2.io/goutils/multierr`
	`r00t2.io/sysutils/envs`
	`r00t2.io/sysutils/paths`
)

/*
	CacheReserved caches the IANA address/network reservations to disk
	(and updates ianaReserved4, ianaReserved6, and reservedNets).

	It is up to the caller to schedule periodic CacheReserved calls
	according to their needs for long-lived processes.
	It *shouldn't* cause any memory leaks, but this has not been tested/confirmed.

	If caching is not enabled, CacheReserved exits withouth any retrieval, parsing, etc.
*/
func CacheReserved() (err error) {

	var dat4 []byte
	var dat6 []byte

	if !isCaching {
		return
	}

	if cacheDir == "" {
		if cacheDir, err = getDefCachePath(); err != nil {
			return
		}
	}

	if dat4, dat6, err = getLive(); err != nil {
		return
	}
	if err = writeCache(dat4, dat6); err != nil {
		return
	}

	return
}

// CleanCache clears the cache completely.
func CleanCache() (err error) {

	if cacheDir == "" {
		if cacheDir, err = getDefCachePath(); err != nil {
			return
		}
	}

	if err = os.RemoveAll(cacheDir); err != nil {
		return
	}

	return
}

/*
	RetrieveReserved returns the current reservations and (re-)populates reservedNets.

	It returns a copy of reservedNets and the current reservations that are safe to use concurrently/separately
	from subnetter.

	If caching is enabled, then:

		1.) First the local cache will be checked.

		2.)	If no cache exists, subnetter will attempt to populate it (CacheReserved());
			otherwise the data there will be used.

			2.b.) If an error occurs while parsing the cached data, the cache will be invalidated
					and attempt to be updated.

		3.) If no cache exists and the live resource is unavailable, an error will be returned.

	If not:

		1.) The live resource will be fetched. If it is unavailable, an error will be returned.
*/
func RetrieveReserved() (ipv4, ipv6 IANARegistry, reserved map[netip.Prefix]*IANAAddrNetResRecord, err error) {

	var b []byte
	var dat4 []byte
	var dat6 []byte
	var hasCache bool

	if isCaching {
		if hasCache, err = checkCache(); err != nil {
			return
		}
		if !hasCache {
			if err = CacheReserved(); err != nil {
				return
			}
		}
		if dat4, dat6, err = readCache(); err != nil {
			return
		}
	} else {
		if dat4, dat6, err = getLive(); err != nil {
			return
		}
	}

	if err = json.Unmarshal(dat4, &ipv4); err != nil {
		return
	}
	if err = json.Unmarshal(dat6, &ipv6); err != nil {
		return
	}
	ianaReserved4 = new(IANARegistry)
	ianaReserved4 = new(IANARegistry)
	if err = json.Unmarshal(dat4, ianaReserved4); err != nil {
		return
	}
	if err = json.Unmarshal(dat6, ianaReserved6); err != nil {
		return
	}

	reservedNets = make(map[netip.Prefix]*IANAAddrNetResRecord)
	for _, reg := range []*IANARegistry{
		ianaReserved4, ianaReserved6,
	} {
		for _, rec := range reg.Notice.Records {
			for _, n := range rec.Networks.Prefixes {
				reservedNets[*n] = rec
			}
		}
	}
	if b, err = json.Marshal(reservedNets); err != nil {
		return
	}
	reserved = make(map[netip.Prefix]*IANAAddrNetResRecord)
	if err = json.Unmarshal(b, &reserved); err != nil {
		return
	}

	return
}

// GetCacheConfig returns the current state and path of subnetter's cache.
func GetCacheConfig() (enabled bool, cacheDirPath string) {

	enabled = isCaching
	cacheDirPath = cacheDir

	return
}

// EnableCache enables or disables subnetter's caching.
func EnableCache(enable bool) (err error) {

	var oldVal bool = isCaching

	isCaching = enable

	if isCaching && (oldVal != isCaching) {
		if err = os.MkdirAll(cacheDir, 0o0640); err != nil {
			return
		}
	}

	return
}

/*
	SetCachePath sets the cache path. Use an empty cacheDirPath to use the default path.

	If the cache dir was changed from its previous value, subnetter will attempt to create it.
*/
func SetCachePath(cacheDirPath string) (err error) {

	var oldPath string = cacheDir

	if cacheDirPath == "" {
		if cacheDirPath, err = getDefCachePath(); err != nil {
			return
		}
	} else {
		if err = paths.RealPath(&cacheDirPath); err != nil {
			return
		}
	}

	if cacheDirPath != oldPath {
		if err = os.MkdirAll(cacheDir, cacheDirPerms); err != nil {
			return
		}
	}

	return
}

func checkCache() (hasCache bool, err error) {

	var numCached uint8
	var cacheDirEntries []fs.DirEntry

	if cacheDir == "" {
		if cacheDir, err = getDefCachePath(); err != nil {
			return
		}
	}

	if cacheDirEntries, err = os.ReadDir(cacheDir); err != nil {
		return
	}

	for _, entry := range cacheDirEntries {
		if entry.IsDir() {
			continue
		}
		switch entry.Name() {
		case ianaSpecial4Cache, ianaSpecial6Cache:
			numCached++
		}
		if numCached >= 2 {
			break
		}
	}

	hasCache = numCached > 2 && ianaReserved4 != nil && ianaReserved6 != nil

	return
}

func getDefCachePath() (val string, err error) {

	if envs.HasEnv(cachedirEnvName) {
		val = os.Getenv(cachedirEnvName)
	} else {
		if val, err = os.UserCacheDir(); err != nil {
			return
		}
	}
	if err = paths.RealPath(&val); err != nil {
		return
	}

	return
}

func getLive() (dat4, dat6 []byte, err error) {

	var wg sync.WaitGroup
	var errChan chan error
	var doneChan chan bool
	var mErr *multierr.MultiError
	var numJobs int = 2

	doneChan = make(chan bool, 1)
	mErr = multierr.NewMultiError(nil)
	wg.Add(numJobs)
	errChan = make(chan error, numJobs)

	if cacheClient == nil {
		cacheClient = resty.New()
	}

	// IPv4
	go func() {
		var rErr error
		var req *resty.Request
		var resp *resty.Response
		var dat *IANARegistry = new(IANARegistry)

		defer wg.Done()

		req = cacheClient.R()
		req.SetResult(dat)

		if resp, rErr = req.Get(ianaSpecial4); rErr != nil {
			errChan <- rErr
			return
		}
		if resp.StatusCode() != http.StatusOK {
			errChan <- errors.New(resp.Status())
			return
		}

		ianaReserved4 = new(IANARegistry)
		*ianaReserved4 = *dat

		if dat4, rErr = json.Marshal(dat); rErr != nil {
			errChan <- rErr
			return
		}

	}()

	// IPv6
	go func() {
		var rErr error
		var req *resty.Request
		var resp *resty.Response
		var dat *IANARegistry = new(IANARegistry)

		defer wg.Done()

		req = cacheClient.R()
		req.SetResult(dat)

		if resp, rErr = req.Get(ianaSpecial6); rErr != nil {
			errChan <- rErr
			return
		}
		if resp.StatusCode() != http.StatusOK {
			errChan <- errors.New(resp.Status())
			return
		}

		if dat6, rErr = json.Marshal(dat); rErr != nil {
			errChan <- rErr
			return
		}
	}()

	go func() {
		wg.Wait()
		close(errChan)
		doneChan <- true
	}()

	<-doneChan

	for i := 0; i < numJobs; i++ {
		if err = <-errChan; err != nil {
			mErr.AddError(err)
			err = nil
		}
	}

	if !mErr.IsEmpty() {
		err = mErr
		return
	}

	return
}

func readCache() (dat4, dat6 []byte, err error) {

	if dat4, err = os.ReadFile(filepath.Join(cacheDir, ianaSpecial4Cache)); err != nil {
		return
	}
	if dat6, err = os.ReadFile(filepath.Join(cacheDir, ianaSpecial6Cache)); err != nil {
		return
	}

	return
}

func writeCache(dat4, dat6 []byte) (err error) {

	if err = os.WriteFile(
		filepath.Join(cacheDir, ianaSpecial4Cache),
		dat4,
		cacheFilePerms,
	); err != nil {
		return
	}
	if err = os.WriteFile(
		filepath.Join(cacheDir, ianaSpecial6Cache),
		dat6,
		cacheFilePerms,
	); err != nil {
		return
	}

	return
}