diff --git a/.gitignore b/.gitignore index af85f2a..cd36f86 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ tmp *.o *.sh *.txt +/config.yaml diff --git a/Dockerfile b/Dockerfile index d841404..d943ec6 100644 --- a/Dockerfile +++ b/Dockerfile @@ -16,5 +16,6 @@ RUN apt-get update \ && rm -rf /var/lib/apt/lists/* COPY --from=corebuilder /work/hyperproxy /usr/local/bin +COPY ./config.yaml.sample /etc/hyperproxy/config.yaml CMD ["hyperproxy"] diff --git a/config.yaml.sample b/config.yaml.sample new file mode 100644 index 0000000..cba199f --- /dev/null +++ b/config.yaml.sample @@ -0,0 +1 @@ +whitelist: \ No newline at end of file diff --git a/go.mod b/go.mod index 75f91fb..086c088 100644 --- a/go.mod +++ b/go.mod @@ -44,6 +44,7 @@ require ( google.golang.org/genproto/googleapis/rpc v0.0.0-20241104194629-dd2ea8efbc28 // indirect google.golang.org/grpc v1.67.1 // indirect google.golang.org/protobuf v1.35.1 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect ) replace github.com/chai2010/webp => github.com/totegamma/webp v0.0.0 diff --git a/go.sum b/go.sum index 0b287fe..39eecb4 100644 --- a/go.sum +++ b/go.sum @@ -88,5 +88,6 @@ google.golang.org/grpc v1.67.1 h1:zWnc1Vrcno+lHZCOofnIMvycFcc0QRGIzm9dhnDX68E= google.golang.org/grpc v1.67.1/go.mod h1:1gLDyUQU7CTLJI90u3nXZ9ekeghjeM7pTDZlqFNg2AA= google.golang.org/protobuf v1.35.1 h1:m3LfL6/Ca+fqnjnlqQXNpFPABW1UD7mjh8KO2mKFytA= google.golang.org/protobuf v1.35.1/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/image.go b/image.go index 8a952ac..fd5cdda 100644 --- a/image.go +++ b/image.go @@ -205,15 +205,40 @@ func ImageHandler(c echo.Context) error { return c.String(400, err.Error()) } - for _, denyIP := range denyIps { - _, ipnet, err := net.ParseCIDR(denyIP) - if err != nil { - fmt.Println("Error parsing CIDR: ", err) - span.RecordError(err) + whiteListMap := make(map[string]bool) + var whiteListCIDRs []*net.IPNet + for _, entry := range IpsWhiteList { + if _, ipNet, err := net.ParseCIDR(entry); err == nil { + whiteListCIDRs = append(whiteListCIDRs, ipNet) + } else { + whiteListMap[entry] = true + } + } + + for _, targetIP := range targetIPs { + + if whiteListMap[targetIP.String()] { continue } + inWhiteList := false + for _, cidr := range whiteListCIDRs { + if cidr.Contains(targetIP) { + inWhiteList = true + break + } + } + if inWhiteList { + continue + } + + for _, denyIP := range denyIps { + _, ipnet, err := net.ParseCIDR(denyIP) + if err != nil { + fmt.Println("Error parsing CIDR: ", err) + span.RecordError(err) + continue + } - for _, targetIP := range targetIPs { if ipnet.Contains(targetIP) { err := errors.New("IP is in deny list") span.RecordError(err) diff --git a/main.go b/main.go index 368f4aa..65a1077 100644 --- a/main.go +++ b/main.go @@ -3,6 +3,8 @@ package main import ( "context" "fmt" + "gopkg.in/yaml.v3" + "io/ioutil" "net/http" "os" "time" @@ -30,6 +32,27 @@ var denyIps = []string{ "fc00::/7", } +var IpsWhiteList []string + +type Config struct { + Whitelist []string `yaml:"whitelist"` +} + +func LoadWhitelist(filePath string) ([]string, error) { + data, err := ioutil.ReadFile(filePath) + if err != nil { + return nil, err + } + + var config Config + err = yaml.Unmarshal(data, &config) + if err != nil { + return nil, err + } + + return config.Whitelist, nil +} + var ( mc *memcache.Client client = &http.Client{ @@ -44,6 +67,16 @@ const ( func main() { + whitelistFile := "/etc/hyperproxy/config.yaml" + loadedWhitelist, err := LoadWhitelist(whitelistFile) + if err != nil { + fmt.Println("Error loading whitelist:", err) + os.Exit(1) + } + + IpsWhiteList = loadedWhitelist + fmt.Println("Loaded whitelist:", IpsWhiteList) + mc = memcache.New(os.Getenv("MEMCACHED_HOST")) defer mc.Close()