permlib/fieldperm.go

package permlib

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

type fieldPermEntry struct {
	jsonName string
	permCode string
}

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
	if fm, ok := e.fieldPerms[key]; ok && len(fm.Response) > 0 {
		return filterResponseByMap(body, fm.Response, perms)
	}
	return body
}

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

	entries := make([]fieldPermEntry, 0, len(fieldMap))
	for jsonField, permCode := range fieldMap {
		entries = append(entries, fieldPermEntry{jsonName: jsonField, permCode: permCode})
	}

	if body[0] == '[' {
		var arr []json.RawMessage
		if err := json.Unmarshal(body, &arr); err != nil {
			return body
		}
		for i, item := range arr {
			arr[i] = filterObject(item, entries, permSet)
		}
		result, _ := json.Marshal(arr)
		return result
	}
	return filterObject(body, entries, permSet)
}

func filterObject(raw json.RawMessage, entries []fieldPermEntry, permSet map[string]bool) json.RawMessage {
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(raw, &obj); err != nil {
		return raw
	}
	for _, entry := range entries {
		if _, has := obj[entry.jsonName]; has && !permSet[entry.permCode] {
			delete(obj, entry.jsonName)
		}
	}
	result, _ := json.Marshal(obj)
	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))
}