permlib/middleware.go

package permlib

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
)

func (e *Engine) authMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, req *http.Request) {
		authHeader := req.Header.Get("Authorization")
		if authHeader == "" {
			e.cfg.Callbacks.OnError(w, req, 401, 401, "未登录")
			return
		}

		tokenStr := strings.TrimPrefix(authHeader, "Bearer ")
		if tokenStr == authHeader {
			e.cfg.Callbacks.OnError(w, req, 401, 401, "token格式错误")
			return
		}

		info, err := e.verifyAndGetUser(req.Context(), tokenStr)
		if err != nil {
			e.cfg.Callbacks.OnError(w, req, 401, 401, "token无效或已过期")
			return
		}

		apiCode := e.resolvePermCode(req)
		if apiCode != "" && !containsPerm(info.Perms, apiCode) {
			e.cfg.Callbacks.OnError(w, req, 403, 403, fmt.Sprintf("无权访问: %s", apiCode))
			return
		}

		user := &UserInfo{
			UserId:      info.UserId,
			Username:    info.Username,
			ProductCode: info.ProductCode,
			MemberType:  info.MemberType,
			Perms:       info.Perms,
		}
		ctx := withUser(req.Context(), user)

		req, rejected := e.filterRequest(req, info.Perms)
		if rejected {
			e.cfg.Callbacks.OnError(w, req, 403, 403, "包含无权写入的字段")
			return
		}

		dataCode := e.resolveDataCode(req, apiCode)
		hasDataPerm := dataCode == "" || containsPerm(info.Perms, dataCode)

		var wrappedNext http.HandlerFunc
		if e.hasRespFilter(req) {
			wrappedNext = func(w http.ResponseWriter, req *http.Request) {
				rw := newFilterResponseWriter(w, req, info.Perms, e)
				next(rw, req)
				rw.flush()
			}
		} else {
			wrappedNext = next
		}

		e.cfg.Callbacks.OnSuccess(w, req.WithContext(ctx), wrappedNext, user, hasDataPerm)
	}
}

func (e *Engine) resolvePermCode(req *http.Request) string {
	key := req.Method + " " + req.URL.Path
	if decl, ok := e.routePerms[key]; ok {
		return decl.PermCode
	}
	return ""
}

func (e *Engine) resolveDataCode(req *http.Request, apiCode string) string {
	key := req.Method + " " + req.URL.Path
	if decl, ok := e.routePerms[key]; ok {
		return decl.DataCode
	}
	if apiCode != "" {
		return apiToDataCode(apiCode)
	}
	return ""
}

func (e *Engine) hasRespFilter(req *http.Request) bool {
	key := req.Method + " " + req.URL.Path
	if fm, ok := e.fieldPerms[key]; ok {
		return len(fm.Response) > 0
	}
	return false
}

func (e *Engine) filterRequest(req *http.Request, perms []string) (*http.Request, bool) {
	key := req.Method + " " + req.URL.Path
	if fm, ok := e.fieldPerms[key]; ok && len(fm.Request) > 0 {
		return filterRequestByMap(req, fm.Request, perms, e.cfg.FieldWriteMode)
	}
	return req, false
}

func (e *Engine) verifyAndGetUser(ctx context.Context, token string) (*cachedUser, error) {
	if u, ok := e.cache.get(token); ok {
		return u, nil
	}

	v, err, _ := e.cache.sf.Do(token, func() (interface{}, error) {
		resp, err := e.client.verifyToken(ctx, token)
		if err != nil {
			return nil, err
		}
		if !resp.Valid {
			return nil, fmt.Errorf("token验证失败")
		}
		u := &cachedUser{
			UserId:      resp.UserId,
			Username:    resp.Username,
			ProductCode: resp.ProductCode,
			MemberType:  resp.MemberType,
			Perms:       resp.Perms,
		}
		e.cache.set(token, u)
		return u, nil
	})
	if err != nil {
		return nil, err
	}
	return v.(*cachedUser), nil
}

func containsPerm(perms []string, code string) bool {
	for _, p := range perms {
		if p == code {
			return true
		}
	}
	return false
}

func filterRequestByMap(req *http.Request, fieldMap map[string]string, perms []string, mode FieldWriteMode) (*http.Request, bool) {
	if req.Body == nil {
		return req, false
	}

	body, err := readBody(req)
	if err != nil || len(body) == 0 {
		return req, false
	}

	var obj map[string]json.RawMessage
	if err := json.Unmarshal(body, &obj); err != nil {
		restoreBody(req, body)
		return req, false
	}

	permSet := toPermSet(perms)
	for jsonField, permCode := range fieldMap {
		if _, has := obj[jsonField]; has && !permSet[permCode] {
			if mode == FieldWriteReject {
				restoreBody(req, body)
				return req, true
			}
			delete(obj, jsonField)
		}
	}

	filtered, _ := json.Marshal(obj)
	restoreBody(req, filtered)
	req.ContentLength = int64(len(filtered))
	return req, false
}