package helper import ( "fmt" "io" "mime/multipart" "strconv" "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" ) type derivedPricingFunc func(req, resp map[string]any) (any, error) var derivedPricingFuncs = map[string]derivedPricingFunc{} type ResolvedDimensionValues struct { Values map[string]any Missing map[string]bool RawInput map[string]any } type DimensionResolver struct { c *gin.Context request map[string]any response map[string]any } func NewDimensionResolver(c *gin.Context) (*DimensionResolver, error) { contentType := c.Request.Header.Get("Content-Type") req := map[string]any{} switch { case strings.HasPrefix(contentType, "application/json"): storage, err := common.GetBodyStorage(c) if err != nil { return nil, err } body, err := storage.Bytes() if err != nil { return nil, err } if len(strings.TrimSpace(string(body))) > 0 { if err := common.Unmarshal(body, &req); err != nil { return nil, err } } if err := resetRequestBody(c, storage); err != nil { return nil, err } case strings.Contains(contentType, gin.MIMEMultipartPOSTForm): form, err := common.ParseMultipartFormReusable(c) if err != nil { return nil, err } req = requestMapFromMultipart(form) default: return nil, fmt.Errorf("unsupported content-type for pricing dimensions: %s", contentType) } return &DimensionResolver{ c: c, request: req, response: map[string]any{}, }, nil } func (r *DimensionResolver) ResolveDimensions(dims []types.PricingDimension) (*ResolvedDimensionValues, error) { resolved := &ResolvedDimensionValues{ Values: make(map[string]any, len(dims)), Missing: make(map[string]bool), RawInput: types.CloneMapAny(r.request), } for _, dim := range dims { value, exists, err := r.resolveSource(dim.Source) if err != nil { return nil, err } if !exists || !hasNonEmptyValue(value) { if dim.Optional { resolved.Values[dim.Key] = dim.Default continue } resolved.Values[dim.Key] = nil resolved.Missing[dim.Key] = true continue } converted, err := convertPricingDimensionValue(value, dim.Type) if err != nil { return nil, fmt.Errorf("dimension %q: %w", dim.Key, err) } resolved.Values[dim.Key] = converted } return resolved, nil } func IsAllowedPricingSource(source string) bool { if strings.HasPrefix(source, "response.") { return false } if strings.HasPrefix(source, "request.") { return len(strings.TrimPrefix(source, "request.")) > 0 } if _, ok := derivedPricingFuncs[strings.TrimPrefix(source, "derived.")]; ok && strings.HasPrefix(source, "derived.") { return true } return false } func registerDerivedPricingFunc(name string, fn derivedPricingFunc) { derivedPricingFuncs[name] = fn } func (r *DimensionResolver) resolveSource(source string) (any, bool, error) { if strings.HasPrefix(source, "response.") { return nil, false, fmt.Errorf("pricing source %q is unsupported in schema v1", source) } if strings.HasPrefix(source, "request.") { return getNestedValue(r.request, strings.TrimPrefix(source, "request.")) } if strings.HasPrefix(source, "derived.") { name := strings.TrimPrefix(source, "derived.") fn, ok := derivedPricingFuncs[name] if !ok { return nil, false, fmt.Errorf("derived pricing source %q is not registered", source) } value, err := fn(r.request, r.response) if err != nil { return nil, false, err } return value, hasNonEmptyValue(value), nil } return nil, false, fmt.Errorf("pricing source %q is not allowed", source) } func getNestedValue(src map[string]any, path string) (any, bool, error) { if path == "" { return nil, false, nil } var cur any = src for _, part := range strings.Split(path, ".") { m, ok := cur.(map[string]any) if !ok { return nil, false, nil } cur, ok = m[part] if !ok { return nil, false, nil } } return cur, true, nil } func convertPricingDimensionValue(value any, typ string) (any, error) { switch typ { case "string": if s, ok := value.(string); ok { return s, nil } return fmt.Sprintf("%v", value), nil case "number": switch v := value.(type) { case float64: return v, nil case float32: return float64(v), nil case int: return float64(v), nil case int64: return float64(v), nil case int32: return float64(v), nil case string: f, err := strconv.ParseFloat(strings.TrimSpace(v), 64) if err != nil { return nil, fmt.Errorf("value %q is not a number", v) } return f, nil default: return nil, fmt.Errorf("value is not a number") } case "boolean": switch v := value.(type) { case bool: return v, nil case string: b, err := strconv.ParseBool(strings.TrimSpace(v)) if err != nil { return nil, fmt.Errorf("value %q is not a boolean", v) } return b, nil default: return nil, fmt.Errorf("value is not a boolean") } default: return nil, fmt.Errorf("unknown dimension type %q", typ) } } func requestMapFromMultipart(form *multipart.Form) map[string]any { req := map[string]any{} for key, vals := range form.Value { if len(vals) == 1 { assignNestedValue(req, key, vals[0]) continue } copied := append([]string(nil), vals...) assignNestedValue(req, key, copied) } for key, files := range form.File { if len(files) == 1 { assignNestedValue(req, key, files[0].Filename) continue } names := make([]string, 0, len(files)) for _, file := range files { names = append(names, file.Filename) } assignNestedValue(req, key, names) } return req } func assignNestedValue(dst map[string]any, key string, value any) { parts := strings.Split(key, ".") cur := dst for i, part := range parts { if i == len(parts)-1 { cur[part] = value return } next, ok := cur[part].(map[string]any) if !ok { next = map[string]any{} cur[part] = next } cur = next } } func resetRequestBody(c *gin.Context, storage common.BodyStorage) error { if _, err := storage.Seek(0, io.SeekStart); err != nil { return err } c.Request.Body = io.NopCloser(storage) return nil }