diff --git a/generator/crypt.go b/generator/crypt.go index 4976f34..b099752 100644 --- a/generator/crypt.go +++ b/generator/crypt.go @@ -6,6 +6,7 @@ import ( "crypto/cipher" "crypto/rand" "encoding/base64" + "errors" "io" ) @@ -15,6 +16,23 @@ func pkcs7pad(data []byte, blockSize int) []byte { return append(data, padding...) } +func pkcs7unpad(data []byte) ([]byte, error) { + length := len(data) + if length == 0 { + return nil, errors.New("invalid padding size") + } + padLen := int(data[length-1]) + if padLen > length || padLen == 0 { + return nil, errors.New("invalid padding") + } + for _, v := range data[length-padLen:] { + if int(v) != padLen { + return nil, errors.New("invalid padding") + } + } + return data[:length-padLen], nil +} + func (g *Generator) Decrypt(s string) (string, error) { c, err := aes.NewCipher(g.encryptionKey) if err != nil { @@ -32,7 +50,12 @@ func (g *Generator) Decrypt(s string) (string, error) { cbc.CryptBlocks(cryptText, cryptText) - return string(cryptText), err + decrypted, err := pkcs7unpad(cryptText) + if err != nil { + return "", err + } + + return string(decrypted), err } func (g *Generator) generateBaseAesEncUrl(file []byte) (string, error) { diff --git a/generator/url.go b/generator/url.go index cfa6e33..e17b28e 100644 --- a/generator/url.go +++ b/generator/url.go @@ -41,9 +41,8 @@ func (g *Generator) GenerateUrl(file string, params []string, format Format) (st if params == nil || len(params) == 0 || params[0] == "" { params = []string{"raw:1"} - } else { - params = append(params, "sm:1") } + params = append(params, "sm:1") if PathPrefix != "" { file = PathPrefix + file