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
Co-Authored-By: Christian Banse <oxisto@aybaze.com>
  • Loading branch information
zhouyiheng.go and oxisto committed Sep 13, 2023
1 parent 1e76606 commit 6207357
Show file tree
Hide file tree
Showing 7 changed files with 180 additions and 17 deletions.
17 changes: 17 additions & 0 deletions encoder.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
package jwt

// Base64Encoder is an interface that allows to implement custom Base64 encoding
// algorithms.
type Base64EncodeFunc func(src []byte) string

// Base64Decoder is an interface that allows to implement custom Base64 decoding
// algorithms.
type Base64DecodeFunc func(s string) ([]byte, error)

// JSONEncoder is an interface that allows to implement custom JSON encoding
// algorithms.
type JSONMarshalFunc func(v any) ([]byte, error)

// JSONUnmarshal is an interface that allows to implement custom JSON unmarshal
// algorithms.
type JSONUnmarshalFunc func(data []byte, v any) error
29 changes: 24 additions & 5 deletions parser.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,17 +12,22 @@ 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

unmarshalFunc JSONUnmarshalFunc
base64DecodeFunc Base64DecodeFunc
}

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

// Choose our JSON decoder. If no custom function is supplied, we use the standard library.
var unmarshal JSONUnmarshalFunc
if p.unmarshalFunc != nil {
unmarshal = p.unmarshalFunc
} else {
unmarshal = json.Unmarshal
}

err = unmarshal(headerBytes, &token.Header)
if err != nil {
return token, parts, newError("could not JSON decode header", ErrTokenMalformed, err)
}

Expand All @@ -162,13 +177,13 @@ func (p *Parser) ParseUnverified(tokenString string, claims Claims) (token *Toke

// If `useJSONNumber` is enabled then we must use *json.Decoder to decode
// the claims. However, this comes with a performance penalty so only use
// it if we must and, otherwise, simple use json.Unmarshal.
// it if we must and, otherwise, simple use our decode function.
if !p.useJSONNumber {
// JSON Unmarshal. Special case for map type to avoid weird pointer behavior.
if c, ok := token.Claims.(MapClaims); ok {
err = json.Unmarshal(claimBytes, &c)
err = unmarshal(claimBytes, &c)
} else {
err = json.Unmarshal(claimBytes, &claims)
err = unmarshal(claimBytes, &claims)
}
} else {
dec := json.NewDecoder(bytes.NewBuffer(claimBytes))
Expand Down Expand Up @@ -200,6 +215,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.base64DecodeFunc != nil {
return p.base64DecodeFunc(seg)
}

encoding := base64.RawURLEncoding

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

// WithJSONUnmarshal supports a custom [JSONUnmarshal] to use in parsing the JWT.
func WithJSONUnmarshal(f JSONUnmarshalFunc) ParserOption {
return func(p *Parser) {
p.unmarshalFunc = f
}
}

// WithBase64Decoder supports a custom [Base64Decoder] to use in parsing the JWT.
func WithBase64Decoder(f Base64DecodeFunc) ParserOption {
return func(p *Parser) {
p.base64DecodeFunc = f
}
}
Loading

0 comments on commit 6207357

Please sign in to comment.