Resolve RFC 6901 JSON Pointers Safely

package main

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
)

func escapeToken(token string) string {
	token = strings.ReplaceAll(token, "~", "~0")
	return strings.ReplaceAll(token, "/", "~1")
}

func decodeToken(token string) (string, error) {
	var builder strings.Builder
	for index := 0; index < len(token); index++ {
		if token[index] != '~' {
			builder.WriteByte(token[index])
			continue
		}
		if index+1 >= len(token) {
			return "", fmt.Errorf("dangling ~ in token %q", token)
		}
		switch token[index+1] {
		case '0':
			builder.WriteByte('~')
		case '1':
			builder.WriteByte('/')
		default:
			return "", fmt.Errorf("invalid escape ~%c in token %q", token[index+1], token)
		}
		index++
	}
	return builder.String(), nil
}

func parsePointer(pointer string) ([]string, error) {
	if pointer == "" {
		return nil, nil
	}
	if !strings.HasPrefix(pointer, "/") {
		return nil, fmt.Errorf("pointer must start with /: %q", pointer)
	}
	rawParts := strings.Split(pointer[1:], "/")
	tokens := make([]string, 0, len(rawParts))
	for _, part := range rawParts {
		token, err := decodeToken(part)
		if err != nil {
			return nil, err
		}
		tokens = append(tokens, token)
	}
	return tokens, nil
}

func Resolve(document any, pointer string) (any, error) {
	tokens, err := parsePointer(pointer)
	if err != nil {
		return nil, err
	}
	current := document
	for _, token := range tokens {
		switch node := current.(type) {
		case map[string]any:
			value, ok := node[token]
			if !ok {
				return nil, fmt.Errorf("missing key %q", token)
			}
			current = value
		case []any:
			if token == "-" {
				return nil, fmt.Errorf("'-' is only valid when appending")
			}
			index, err := strconv.Atoi(token)
			if err != nil || index < 0 || index >= len(node) {
				return nil, fmt.Errorf("invalid array index %q", token)
			}
			current = node[index]
		default:
			return nil, fmt.Errorf("cannot descend into %T at %q", current, token)
		}
	}
	return current, nil
}

func main() {
	payload := `{"items":[{"name":"first"},{"name":"second"}],"a/b":{"m~n":42}}`
	document := map[string]any{}
	if err := json.Unmarshal([]byte(payload), &document); err != nil {
		panic(err)
	}
	name, err := Resolve(document, "/items/1/name")
	if err != nil || name != "second" {
		panic("failed to resolve array element")
	}
	value, err := Resolve(document, "/a~1b/m~0n")
	if err != nil || value != float64(42) {
		panic("failed to resolve escaped keys")
	}
	fmt.Println(name, value, escapeToken("a/b~c"))
	if _, err := Resolve(document, "/items/9"); err == nil {
		panic("out-of-range index should fail")
	}
	if _, err := Resolve(document, "/bad~2escape"); err == nil {
		panic("invalid escape should fail")
	}
}