Implemented Whitelist checking

cmd/jwtgen.go

@@ -1,3 +1,5 @@
+// +build dev
+
 /*
 Copyright © 2021 darmiel <hi@d2a.io>
 
@@ -17,45 +19,29 @@ package cmd
 
 import (
 	"fmt"
-	"github.com/dgrijalva/jwt-go"
+	"github.com/darmiel/yaxc/internal/whitelist"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 	"log"
-	"time"
 )
 
-type Claim struct {
-	MaxBody  int64 `json:"max_body"`
-	CountNum int   `json:"count_num_id"`
-	jwt.StandardClaims
-}
-
 // jwtgenCmd represents the jwtgen command
 var jwtgenCmd = &cobra.Command{
 	Use:   "jwtgen",
 	Short: "Generate JWT Token",
 	Run: func(cmd *cobra.Command, args []string) {
-		secret := viper.GetString("jwt")
+		secret := viper.GetString("secret")
 		maxBody := viper.GetInt64("max-body")
 		audience := viper.GetString("audience")
 		issuer := viper.GetString("issuer")
 		count := viper.GetInt("count")
 
-		claims := &Claim{
-			MaxBody: maxBody,
-			StandardClaims: jwt.StandardClaims{
-				Audience: audience,
-				IssuedAt: time.Now().Unix(),
-				Issuer:   issuer,
-			},
-		}
-
 		fmt.Println("🔨 Generating", count, "JWT-Tokens ...")
 		for i := 0; i < count; i++ {
-			claims.CountNum = i
-
-			token := jwt.NewWithClaims(jwt.SigningMethodHS512, claims)
-			signed, err := token.SignedString([]byte(secret))
+			signed, err := whitelist.GenerateToken(secret,
+				audience,
+				issuer,
+				maxBody)
 			if err != nil {
 				log.Fatalln("Error signing:", err)
 				return

cmd/serve.go

@@ -40,7 +40,7 @@ var serveCmd = &cobra.Command{
 		defTTL := viper.GetDuration("default-ttl")
 		minTTL := viper.GetDuration("min-ttl")
 		maxTTL := viper.GetDuration("max-ttl")
-		maxBodyLen := viper.GetInt("max-body-length")
+		maxBodyLen := viper.GetInt64("max-body-length")
 
 		// validate values
 		if bind == "" {
@@ -74,6 +74,7 @@ var serveCmd = &cobra.Command{
 		// other
 		enableEnc := viper.GetBool("enable-encryption")
 		proxyHeader := viper.GetString("proxy-header")
+		jwt := viper.GetString("jwt")
 
 		// create server & start
 		s := server.NewServer(&server.YAxCConfig{
@@ -92,6 +93,7 @@ var serveCmd = &cobra.Command{
 			// Other
 			EnableEncryption: enableEnc,
 			ProxyHeader:      proxyHeader,
+			JWTSign:          []byte(jwt),
 		})
 		go s.Start()
 
@@ -130,7 +132,8 @@ func init() {
 	regDurP(serveCmd, "max-ttl", "s", 60*time.Minute, "Max TTL")
 
 	// other
-	regIntP(serveCmd, "max-body-length", "x", 8192, "Max Body Length")
+	regInt64P(serveCmd, "max-body-length", "x", 8192, "Max Body Length")
 	regBoolP(serveCmd, "enable-encryption", "e", true, "Enable Encryption")
 	regStr(serveCmd, "proxy-header", "", "Proxy Header")
+	regStr(serveCmd, "jwt", "", "JWT-Token")
 }

internal/server/request_post.go

@@ -6,6 +6,7 @@ import (
 	"encoding/base64"
 	"fmt"
 	"github.com/darmiel/yaxc/internal/common"
+	"github.com/darmiel/yaxc/internal/whitelist"
 	"github.com/gofiber/fiber/v2"
 	"github.com/gofiber/fiber/v2/utils"
 	"github.com/muesli/termenv"
@@ -35,9 +36,31 @@ func (s *YAxCServer) setAnywhereWithHash(ctx *fiber.Ctx, path, hash string) (err
 		return fiber.NewError(http.StatusNotAcceptable, "invalid anywhere-path")
 	}
 
+	var maxBodyLength = s.MaxBodyLength
+
+	// check authorization
+	if auth := ctx.Get("Authorization"); auth != "" {
+		if len(s.JWTSign) == 0 {
+			return fiber.NewError(509, "whitelist not available")
+		}
+		spl := strings.Split(auth, " ")
+		if spl[0] != "JWT" {
+			return fiber.NewError(http.StatusUnauthorized, "invalid auth type - JWT required")
+		}
+		claims, err := whitelist.ValidateToken(s.JWTSign, spl[1])
+		if err != nil {
+			return fiber.NewError(510, "error validating token: "+err.Error())
+		}
+		if claims.MaxBody != 0 {
+			maxBodyLength = claims.MaxBody
+		}
+		fmt.Println(common.StyleInfo(),
+			claims.Issuer, "used whitelist key with claim: mb:", claims.MaxBody, ", ia:", claims.IssuedAt)
+	}
+
 	// Read content
 	bytes := ctx.Body()
-	if s.MaxBodyLength > 0 && len(bytes) > s.MaxBodyLength {
+	if s.MaxBodyLength > 0 && int64(len(bytes)) > maxBodyLength {
 		return fiber.NewError(http.StatusRequestEntityTooLarge, "exceeded max body length")
 	}
 	content := string(bytes)

internal/server/server.go

@@ -13,7 +13,7 @@ import (
 func (s *YAxCServer) Start() {
 	log.Info("Starting YAxC server on", s.BindAddress)
 
-	cfg := fiber.Config{}
+	cfg := fiber.Config{ReadTimeout: 10 * time.Second}
 
 	if s.ProxyHeader != "" {
 		if s.ProxyHeader == "$proxy" {

internal/server/type.go

@@ -40,9 +40,10 @@ type YAxCConfig struct {
 	MinTTL     time.Duration // == MaxTTL -> cannot specify TTL
 	MaxTTL     time.Duration // == MinTTL -> cannot specify TTL
 	// Other
-	MaxBodyLength    int
+	MaxBodyLength    int64
 	EnableEncryption bool
 	ProxyHeader      string
+	JWTSign          []byte
 }
 
 type YAxCServer struct {
@@ -91,6 +92,5 @@ func NewServer(cfg *YAxCConfig) (s *YAxCServer) {
 	}
 
 	log.Info("Encryption:", s.EnableEncryption)
-
 	return
 }

internal/whitelist/claim.go

@@ -0,0 +1,9 @@
+package whitelist
+
+import "github.com/dgrijalva/jwt-go"
+
+type Claim struct {
+	MaxBody  int64 `json:"max_body"`
+	RandomID int   `json:"random_id"`
+	jwt.StandardClaims
+}

internal/whitelist/func.go

@@ -0,0 +1,48 @@
+package whitelist
+
+import (
+	"errors"
+	"fmt"
+	"github.com/darmiel/yaxc/internal/common"
+	"github.com/dgrijalva/jwt-go"
+	"math/rand"
+	"time"
+)
+
+func init() {
+	rand.Seed(time.Now().UnixNano())
+}
+
+var ErrParse = errors.New("error parsing claims (cast)")
+
+func GenerateToken(secret, audience, issuer string, maxBody int64) (string, error) {
+	randomId := rand.Intn(9999999)
+	claims := &Claim{
+		// attributes
+		MaxBody: maxBody,
+		// generate random id
+		RandomID: randomId,
+		StandardClaims: jwt.StandardClaims{
+			Audience: audience,
+			IssuedAt: time.Now().Unix(),
+			Issuer:   issuer,
+		},
+	}
+	token := jwt.NewWithClaims(jwt.SigningMethodHS512, claims)
+	return token.SignedString([]byte(secret))
+}
+
+func ValidateToken(secret []byte, signed string) (claims *Claim, err error) {
+	fmt.Println(common.StyleDebug(), "checking token", signed, "with secret:", string(secret))
+	var token *jwt.Token
+	if token, err = jwt.ParseWithClaims(signed, &Claim{}, func(token *jwt.Token) (interface{}, error) {
+		return secret, nil
+	}); err != nil {
+		return
+	}
+	var ok bool
+	if claims, ok = token.Claims.(*Claim); !ok {
+		return nil, ErrParse
+	}
+	return
+}

swagger.yml

@@ -77,6 +77,8 @@ paths:
       responses:
         "200":
           description: "OK"
+        "401":
+          description: "Invalid Auth-Type in Authorization-Header. Only accepts JWT"
         "406":
           description: "Invalid Path"
         "413":
@@ -93,6 +95,10 @@ paths:
           description: "Error decoding Base64"
         "506":
           description: "Invalid Base64-Mode. Available: encode, decode"
+        "509":
+          description: "The whitelist is currently not enabled on the server"
+        "510":
+          description: "Error validating token"
 
     #
     # GET /{anywhere}