Zoraxy/mod/access/access.go

package access

import (
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"sync"

	"imuslab.com/zoraxy/mod/database"
	"imuslab.com/zoraxy/mod/geodb"
	"imuslab.com/zoraxy/mod/info/logger"
	"imuslab.com/zoraxy/mod/utils"
)

/*
	Access.go

	This module is the new version of access control system
	where now the blacklist / whitelist are seperated from
	geodb module
*/

type Options struct {
	Logger       logger.Logger
	ConfigFolder string             //Path for storing config files
	GeoDB        *geodb.Store       //For resolving country code
	Database     *database.Database //System key-value database
}

type AccessRule struct {
	ID               string
	Name             string
	Desc             string
	BlacklistEnabled bool
	WhitelistEnabled bool

	/* Whitelist Blacklist Table, value is comment if supported */
	WhiteListCountryCode *map[string]string
	WhiteListIP          *map[string]string
	BlackListContryCode  *map[string]string
	BlackListIP          *map[string]string

	parent *Controller
}

type Controller struct {
	GlobalAccessRule *AccessRule
	ProxyAccessRule  *sync.Map
	Options          *Options
}

// Create a new access controller to handle blacklist / whitelist
func NewAccessController(options *Options) (*Controller, error) {
	sysdb := options.Database
	if sysdb == nil {
		return nil, errors.New("missing database access")
	}

	//Create the config folder if not exists
	confFolder := options.ConfigFolder
	if !utils.FileExists(confFolder) {
		err := os.MkdirAll(confFolder, 0775)
		if err != nil {
			return nil, err
		}
	}

	// Create the global access rule if not exists
	globalAccessRule := AccessRule{
		ID:               "default",
		Name:             "Default",
		Desc:             "Default access rule for all HTTP proxy hosts",
		BlacklistEnabled: false,
		WhitelistEnabled: false,
	}
	defaultRuleSettingFile := filepath.Join(confFolder, "default.json")
	if utils.FileExists(defaultRuleSettingFile) {
		//Load from file
		defaultRuleBytes, err := os.ReadFile(defaultRuleSettingFile)
		if err == nil {
			err = json.Unmarshal(defaultRuleBytes, &globalAccessRule)
			if err != nil {
				options.Logger.PrintAndLog("Access", "Unable to parse default routing rule config file. Using default", err)
			}
		}
	} else {
		//Create one
		js, _ := json.MarshalIndent(globalAccessRule, "", " ")
		os.WriteFile(defaultRuleSettingFile, js, 0775)
	}

	//Generate a controller object
	thisController := Controller{
		GlobalAccessRule: &globalAccessRule,
		Options:          options,
	}

	//Load all acccess rules from file
	configFiles, err := filepath.Glob(options.ConfigFolder + "/*.json")
	if err != nil {
		return nil, err
	}
	ProxyAccessRules := sync.Map{}
	for _, configFile := range configFiles {
		if filepath.Base(configFile) == "default.json" {
			//Skip this, as this was already loaded as default
			continue
		}

		configContent, err := os.ReadFile(configFile)
		if err != nil {
			options.Logger.PrintAndLog("Access", "Unable to load config "+filepath.Base(configFile), err)
			continue
		}

		//Parse the config file into AccessRule
		thisAccessRule := AccessRule{}
		err = json.Unmarshal(configContent, &thisAccessRule)
		if err != nil {
			options.Logger.PrintAndLog("Access", "Unable to parse config "+filepath.Base(configFile), err)
			continue
		}
		thisAccessRule.parent = &thisController
		ProxyAccessRules.Store(thisAccessRule.ID, &thisAccessRule)
	}
	thisController.ProxyAccessRule = &ProxyAccessRules

	return &thisController, nil
}

// Get the global access rule
func (c *Controller) GetGlobalAccessRule() (*AccessRule, error) {
	if c.GlobalAccessRule == nil {
		return nil, errors.New("global access rule is not set")
	}
	return c.GlobalAccessRule, nil
}

// Load access rules to runtime, require rule ID
func (c *Controller) GetAccessRuleByID(accessRuleID string) (*AccessRule, error) {
	//Load from sync.Map, should be O(1)
	targetRule, ok := c.ProxyAccessRule.Load(accessRuleID)
	if !ok {
		return nil, errors.New("target access rule not exists")
	}

	ar, ok := targetRule.(*AccessRule)
	if !ok {
		return nil, errors.New("assertion of access rule failed, version too old?")
	}
	return ar, nil
}