Skip to content

Commit

Permalink
feat: custom json and base64 encoders for Token and Parser
Browse files Browse the repository at this point in the history
  • Loading branch information
zhouyiheng.go committed Apr 2, 2023
1 parent b88a60f commit cd88b49
Show file tree
Hide file tree
Showing 9 changed files with 230 additions and 21 deletions.
13 changes: 13 additions & 0 deletions encoder.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
package jwt

// Base64Encoder is an interface that allows to implement custom Base64 encoding/decoding algorithms.
type Base64Encoder interface {
EncodeToString(src []byte) string
DecodeString(s string) ([]byte, error)
}

// JSONEncoder is an interface that allows to implement custom JSON encoding/decoding algorithms.
type JSONEncoder interface {
Marshal(v any) ([]byte, error)
Unmarshal(data []byte, v any) error
}
26 changes: 26 additions & 0 deletions encoder_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
package jwt_test

import (
"encoding/base64"
"encoding/json"
)

type customJSONEncoder struct{}

func (s *customJSONEncoder) Marshal(v any) ([]byte, error) {
return json.Marshal(v)
}

func (s *customJSONEncoder) Unmarshal(data []byte, v any) error {
return json.Unmarshal(data, v)
}

type customBase64Encoder struct{}

func (s *customBase64Encoder) EncodeToString(data []byte) string {
return base64.StdEncoding.EncodeToString(data)
}

func (s *customBase64Encoder) DecodeString(data string) ([]byte, error) {
return base64.RawURLEncoding.DecodeString(data)
}
34 changes: 34 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM=
github.com/bytedance/sonic v1.8.6 h1:aUgO9S8gvdN6SyW2EhIpAw5E4ChworywIEndZCkCVXk=
github.com/bytedance/sonic v1.8.6/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZXU064P/U=
github.com/chenzhuoyu/base64x v0.0.0-20211019084208-fb5309c8db06/go.mod h1:DH46F32mSOjUmXrMHnKwZdA8wcEefY7UVqBKYGjpdQY=
github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 h1:qSGYFH7+jGhDF8vLC+iwCD4WpbV1EBDSzWkJODFLams=
github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311/go.mod h1:b583jCggY9gE99b6G5LEC39OIiVsWj+R97kbl5odCEk=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4=
github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/segmentio/asm v1.2.0 h1:9BQrFxC+YOHJlTlHGkTrFWf59nbL3XnCoFLTwDCI7ys=
github.com/segmentio/asm v1.2.0/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=
github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
golang.org/x/arch v0.0.0-20210923205945-b76863e36670 h1:18EFjUmQOcUvxNYSkA6jO9VAiXCnxFY6NyDX0bHDmkU=
golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
golang.org/x/sys v0.0.0-20220412211240-33da011f77ad h1:ntjMns5wyP/fN65tdBD4g8J5w8n015+iIIs9rtjXkY0=
golang.org/x/sys v0.0.0-20220412211240-33da011f77ad/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4=
46 changes: 36 additions & 10 deletions parser.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,17 +12,25 @@ type Parser struct {
// If populated, only these methods will be considered valid.
validMethods []string

// Use JSON Number format in JSON decoder.
// Use JSON Number format in JSON decoder. This field is disabled when using a custom json encoder.
useJSONNumber bool

// Skip claims validation during token parsing.
skipClaimsValidation bool

validator *validator

// This field is disabled when using a custom base64 encoder.
decodeStrict bool

// This field is disabled when using a custom base64 encoder.
decodePaddingAllowed bool

// Custom base64 encoder.
base64Encoder Base64Encoder

// Custom json encoder.
jsonEncoder JSONEncoder
}

// NewParser creates a new Parser with the specified options
Expand Down Expand Up @@ -135,7 +143,12 @@ func (p *Parser) ParseUnverified(tokenString string, claims Claims) (token *Toke
}
return token, parts, newError("could not base64 decode header", ErrTokenMalformed, err)
}
if err = json.Unmarshal(headerBytes, &token.Header); err != nil {
if p.jsonEncoder != nil {
err = p.jsonEncoder.Unmarshal(headerBytes, &token.Header)
} else {
err = json.Unmarshal(headerBytes, &token.Header)
}
if err != nil {
return token, parts, newError("could not JSON decode header", ErrTokenMalformed, err)
}

Expand All @@ -146,21 +159,30 @@ func (p *Parser) ParseUnverified(tokenString string, claims Claims) (token *Toke
if claimBytes, err = p.DecodeSegment(parts[1]); err != nil {
return token, parts, newError("could not base64 decode claim", ErrTokenMalformed, err)
}
dec := json.NewDecoder(bytes.NewBuffer(claimBytes))
if p.useJSONNumber {
dec.UseNumber()
}

// JSON Decode. Special case for map type to avoid weird pointer behavior
if c, ok := token.Claims.(MapClaims); ok {
err = dec.Decode(&c)
mapClaims, isMapClaims := token.Claims.(MapClaims)
if p.jsonEncoder != nil {
if isMapClaims {
err = p.jsonEncoder.Unmarshal(claimBytes, &mapClaims)
} else {
err = p.jsonEncoder.Unmarshal(claimBytes, &claims)
}
} else {
err = dec.Decode(&claims)
decoder := json.NewDecoder(bytes.NewBuffer(claimBytes))
if p.useJSONNumber {
decoder.UseNumber()
}
if isMapClaims {
err = decoder.Decode(&mapClaims)
} else {
err = decoder.Decode(&claims)
}
}
// Handle decode error
if err != nil {
return token, parts, newError("could not JSON decode claim", ErrTokenMalformed, err)
}

// Lookup signature method
if method, ok := token.Header["alg"].(string); ok {
if token.Method = GetSigningMethod(method); token.Method == nil {
Expand All @@ -177,6 +199,10 @@ func (p *Parser) ParseUnverified(tokenString string, claims Claims) (token *Toke
// take into account whether the [Parser] is configured with additional options,
// such as [WithStrictDecoding] or [WithPaddingAllowed].
func (p *Parser) DecodeSegment(seg string) ([]byte, error) {
if p.base64Encoder != nil {
return p.base64Encoder.DecodeString(seg)
}

encoding := base64.RawURLEncoding

if p.decodePaddingAllowed {
Expand Down
13 changes: 13 additions & 0 deletions parser_option.go
Original file line number Diff line number Diff line change
Expand Up @@ -125,3 +125,16 @@ func WithStrictDecoding() ParserOption {
p.decodeStrict = true
}
}

// WithJSONEncoder supports
func WithJSONEncoder(enc JSONEncoder) ParserOption {
return func(p *Parser) {
p.jsonEncoder = enc
}
}

func WithBase64Encoder(enc Base64Encoder) ParserOption {
return func(p *Parser) {
p.base64Encoder = enc
}
}
Loading

0 comments on commit cd88b49

Please sign in to comment.