Implemented --proxy-header Flag

cmd/serve.go

@@ -65,6 +65,7 @@ var serveCmd = &cobra.Command{
 
 		// other
 		enableEnc := viper.GetBool("enable-encryption")
+		proxyHeader := viper.GetString("proxy-header")
 
 		// create server & start
 		s := server.NewServer(&server.YAxCConfig{
@@ -81,6 +82,7 @@ var serveCmd = &cobra.Command{
 			MaxBodyLength: maxBodyLen,
 			// Other
 			EnableEncryption: enableEnc,
+			ProxyHeader:      proxyHeader,
 		})
 		s.Start()
 	},
@@ -105,6 +107,7 @@ func init() {
 	// other
 	regIntP(serveCmd, "max-body-length", "x", 1024, "Max Body Length")
 	regBoolP(serveCmd, "enable-encryption", "e", true, "Enable Encryption")
+	regStr(serveCmd, "proxy-header", "", "Proxy Header")
 }
 
 func regStrP(cmd *cobra.Command, name, shorthand, def, usage string) {

internal/server/server.go

@@ -1,10 +1,38 @@
 package server
 
-import "github.com/gofiber/fiber/v2"
+import (
+	"github.com/gofiber/fiber/v2"
+	"github.com/gofiber/fiber/v2/middleware/limiter"
+	"strings"
+	"time"
+)
 
 func (s *yAxCServer) Start() {
 	log.Info("Starting YAxC server on", s.BindAddress)
-	s.App = fiber.New()
+
+	cfg := &fiber.Config{}
+	if s.ProxyHeader != "" {
+		if s.ProxyHeader == "$proxy" {
+			s.ProxyHeader = "X-Forwarded-For"
+		}
+		cfg.ProxyHeader = s.ProxyHeader
+	}
+	s.App = fiber.New(*cfg)
+
+	// limiter middleware
+	s.App.Use(limiter.New(limiter.Config{
+		Next: func(c *fiber.Ctx) bool {
+			if c.IP() == "127.0.0.1" {
+				return true
+			}
+			if strings.HasPrefix(c.Path(), "/hash/") {
+				return true
+			}
+			return false
+		},
+		Max:        30,
+		Expiration: 60 * time.Second,
+	}))
 
 	// register routes
 	s.App.Get("/", func(ctx *fiber.Ctx) error {

internal/server/type.go

@@ -39,6 +39,7 @@ type YAxCConfig struct {
 	// Other
 	MaxBodyLength    int
 	EnableEncryption bool
+	ProxyHeader      string
 }
 
 type yAxCServer struct {