permlib/fieldperm.go

package permlib

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
)

type filterResponseWriter struct {
	http.ResponseWriter
	buf    bytes.Buffer
	req    *http.Request
	perms  []string
	engine *Engine
	code   int
}

func newFilterResponseWriter(w http.ResponseWriter, req *http.Request, perms []string, e *Engine) *filterResponseWriter {
	return &filterResponseWriter{
		ResponseWriter: w,
		req:            req,
		perms:          perms,
		engine:         e,
		code:           200,
	}
}

func (fw *filterResponseWriter) WriteHeader(code int) {
	fw.code = code
}

func (fw *filterResponseWriter) Write(b []byte) (int, error) {
	return fw.buf.Write(b)
}

func (fw *filterResponseWriter) flush() {
	body := fw.buf.Bytes()

	if len(body) > 0 {
		body = fw.engine.filterResponseBody(body, fw.req, fw.perms)
	}

	fw.ResponseWriter.WriteHeader(fw.code)
	fw.ResponseWriter.Write(body)
}

func (e *Engine) filterResponseBody(body []byte, req *http.Request, perms []string) []byte {
	key := req.Method + " " + req.URL.Path
	fm, ok := e.fieldPerms[key]
	if !ok || fm.Response == nil {
		return body
	}
	return filterResponseByMap(body, fm.Response, perms)
}

func filterResponseByMap(body []byte, node *FieldNode, perms []string) []byte {
	if node == nil {
		return body
	}
	permSet := toPermSet(perms)
	body = bytes.TrimSpace(body)
	if len(body) == 0 {
		return body
	}

	if body[0] == '[' {
		return filterArray(body, node, permSet)
	}
	return filterObject(body, node, permSet)
}

func filterObject(raw json.RawMessage, node *FieldNode, permSet map[string]bool) json.RawMessage {
	if node == nil {
		return raw
	}
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(raw, &obj); err != nil {
		return raw
	}
	for jsonName, permCode := range node.Fields {
		if !permSet[permCode] {
			delete(obj, jsonName)
		}
	}
	for jsonName, child := range node.Nested {
		v, ok := obj[jsonName]
		if !ok {
			continue
		}
		v = bytes.TrimSpace(v)
		if len(v) == 0 {
			continue
		}
		if v[0] == '[' {
			obj[jsonName] = filterArray(v, child, permSet)
		} else if v[0] == '{' {
			obj[jsonName] = filterObject(v, child, permSet)
		}
	}
	result, _ := json.Marshal(obj)
	return result
}

func filterArray(raw json.RawMessage, node *FieldNode, permSet map[string]bool) json.RawMessage {
	var arr []json.RawMessage
	if err := json.Unmarshal(raw, &arr); err != nil {
		return raw
	}
	for i, item := range arr {
		arr[i] = filterObject(item, node, permSet)
	}
	result, _ := json.Marshal(arr)
	return result
}

func toPermSet(perms []string) map[string]bool {
	m := make(map[string]bool, len(perms))
	for _, p := range perms {
		m[p] = true
	}
	return m
}

func readBody(req *http.Request) ([]byte, error) {
	if req.Body == nil {
		return nil, nil
	}
	body, err := io.ReadAll(req.Body)
	req.Body.Close()
	if err != nil {
		return nil, err
	}
	req.Body = io.NopCloser(bytes.NewReader(body))
	return body, nil
}

func restoreBody(req *http.Request, body []byte) {
	req.Body = io.NopCloser(bytes.NewReader(body))
}