gen.go

package openapi

import (
	"cmp"
	"fmt"
	"net/http"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"unicode"

	kin "github.com/getkin/kin-openapi/openapi3"
	"github.com/getkin/kin-openapi/openapi3conv"
	kingen "github.com/getkin/kin-openapi/openapi3gen"
	"github.com/go-chi/chi/v5"
)

const ver = "3.1.0"

// SpecConfig configures how the spec is built.
type SpecConfig struct {
	// StripPrefixes strips the given prefixes from the routes.
	StripPrefixes []string

	// ObjPkgSegments determines the maximum number of
	// package segments to use to identify an object.
	ObjPkgSegments int

	// OpenAPIVersion sets the OpenAPI version string in the output document,
	// e.g. "3.0.0" or "3.1.0" (default). When a 3.1.x version is requested
	// the document is upgraded via openapi3conv.Upgrade, which rewrites
	// 3.0-specific schema constructs (nullable, boolean exclusive bounds,
	// singular example) into their 3.1 equivalents.
	OpenAPIVersion string
}

// BuildSpec builds openapi v3 spec from the given chi router.
func BuildSpec(r chi.Routes, cfg SpecConfig) (kin.T, error) {
	cfg.OpenAPIVersion = cmp.Or(cfg.OpenAPIVersion, ver)

	gen := newGenerator(cfg.OpenAPIVersion)
	gen.objPkgSegments = cfg.ObjPkgSegments

	err := chi.Walk(r, func(method, route string, handler http.Handler, middlewares ...func(http.Handler) http.Handler) error {
		for _, prefix := range cfg.StripPrefixes {
			if !strings.HasPrefix(route, prefix) {
				continue
			}
			route = strings.TrimPrefix(route, prefix)
		}

		var (
			op    Operation
			found bool
		)
		for _, m := range middlewares {
			h := m(opHandler{})
			if oph, ok := h.(opHandler); ok {
				op = op.Merge(oph.Op)
				found = true
			}
		}

		if o, ok := opReg.Op(handler); ok {
			op = op.Merge(o)
			found = true
		}

		if !found || op.id == "" {
			return nil
		}

		return gen.AddOperation(method, route, op)
	})
	if err != nil {
		return kin.T{}, err
	}

	doc := gen.doc
	if strings.HasPrefix(cfg.OpenAPIVersion, "3.1") {
		openapi3conv.Upgrade(&doc)
		doc.OpenAPI = cfg.OpenAPIVersion
	}

	return doc, nil
}

type generator struct {
	doc kin.T
	gen *kingen.Generator

	objPkgSegments int
}

func newGenerator(ver string) *generator {
	comp := kin.NewComponents()
	comp.Schemas = kin.Schemas{}
	comp.SecuritySchemes = kin.SecuritySchemes{}

	return &generator{
		doc: kin.T{
			OpenAPI:    ver,
			Components: &comp,
		},
		gen: kingen.NewGenerator(kingen.SchemaCustomizer(customizer(ver))),
	}
}

func (g *generator) AddOperation(method, path string, op Operation) error {
	params, err := g.toParams(op.params)
	if err != nil {
		return fmt.Errorf("generating parameters for %s %q: %w", method, path, err)
	}

	reqBody, err := g.toRequestBody(op.reads, op.consumes)
	if err != nil {
		return fmt.Errorf("generating request body for %s %q: %w", method, path, err)
	}

	responses, err := g.toResponses(op.returns, op.produces)
	if err != nil {
		return fmt.Errorf("generating responses for %s %q: %w", method, path, err)
	}

	secReqs, err := g.addSecuritySchemes(op.security)
	if err != nil {
		return fmt.Errorf("generating security requirement for %s %q: %w", method, path, err)
	}

	g.doc.AddOperation(path, method, &kin.Operation{
		Summary:     op.doc,
		Description: op.description,
		OperationID: op.id,
		Tags:        op.tags,
		Deprecated:  op.deprecated,
		Parameters:  params,
		RequestBody: reqBody,
		Responses:   responses,
		Security:    secReqs,
	})
	return nil
}

func (g *generator) schema(obj any) (*kin.SchemaRef, error) {
	t := reflect.TypeOf(obj)
	if t.Kind() == reflect.Pointer {
		t = t.Elem()
	}
	if t.Kind() != reflect.Struct {
		return g.gen.NewSchemaRefForValue(obj, g.doc.Components.Schemas)
	}

	name := t.Name()
	if from, to := strings.Index(name, "["), strings.LastIndex(name, "]"); from != -1 && from < to {
		// name: "Object[github.com/org/repo/pkg.struct]".
		parts := strings.Split(name[from+1:to], "/")
		parts = strings.Split(strings.Join(parts, "."), ".")
		if l := len(parts); l > g.objPkgSegments {
			parts = parts[l-g.objPkgSegments-1:]
		}

		switch {
		case len(parts) == 0:
			// "Object".
			name = name[:from]
		default:
			// "Object[struct.pkg.repo.org.github.com]".
			name = name[:from] + "[" + strings.Join(parts, ".") + "]"
		}
	}
	if path := t.PkgPath(); path != "" && g.objPkgSegments > 0 {
		parts := strings.Split(path, "/")
		if l := len(parts); l > g.objPkgSegments {
			parts = parts[l-g.objPkgSegments:]
		}
		name = strings.Join(parts, ".") + "." + name
	}

	if _, ok := g.doc.Components.Schemas[name]; ok {
		return &kin.SchemaRef{Ref: "#/components/schemas/" + name}, nil
	}

	schema, err := g.gen.NewSchemaRefForValue(obj, g.doc.Components.Schemas)
	if err != nil {
		return nil, err
	}
	if schema.Value == nil || !isExported(t.Name()) {
		return schema, nil
	}

	g.doc.Components.Schemas[name] = schema

	return &kin.SchemaRef{Ref: "#/components/schemas/" + name}, nil
}

func isExported(name string) bool {
	runes := []rune(name)
	if len(runes) == 0 {
		return false
	}
	return !unicode.IsLower(runes[0])
}

func (g *generator) toParams(params []Parameter) (kin.Parameters, error) {
	if len(params) == 0 {
		return nil, nil
	}

	ret := make([]*kin.ParameterRef, len(params))
	for i, param := range params {
		var (
			schema *kin.SchemaRef
			err    error
		)
		switch {
		case param.typ != "":
			schema = &kin.SchemaRef{
				Value: &kin.Schema{
					Type: &kin.Types{param.typ},
				},
			}
		case param.dataType != nil:
			schema, err = g.schema(param.dataType)
			if err != nil {
				return nil, err
			}
		}

		ret[i] = &kin.ParameterRef{Value: &kin.Parameter{
			Name:        param.name,
			In:          param.in,
			Description: param.description,
			Required:    param.required,
			Schema:      schema,
			Deprecated:  param.deprecated,
		}}
	}
	return ret, nil
}

func (g *generator) toRequestBody(obj any, mediaTypes []string) (*kin.RequestBodyRef, error) {
	if obj == nil || len(mediaTypes) == 0 {
		//nolint:nilnil
		return nil, nil
	}

	schema, err := g.schema(obj)
	if err != nil {
		return nil, err
	}

	content := kin.Content{}
	for _, mime := range mediaTypes {
		content[mime] = &kin.MediaType{Schema: schema}
	}

	return &kin.RequestBodyRef{Value: &kin.RequestBody{
		Required: true,
		Content:  content,
	}}, nil
}

func (g *generator) toResponses(res []Response, mediaTypes []string) (*kin.Responses, error) {
	if len(res) == 0 {
		return nil, nil
	}

	responses := &kin.Responses{}
	for _, r := range res {
		if r.writes == nil {
			responses.Set(strconv.Itoa(r.code), &kin.ResponseRef{Value: &kin.Response{
				Description: &r.description,
			}})
			continue
		}

		schema, err := g.schema(r.writes)
		if err != nil {
			return nil, err
		}

		useMediaTypes := mediaTypes
		if len(r.mediaTypes) > 0 {
			useMediaTypes = r.mediaTypes
		}

		content := kin.Content{}
		for _, mime := range useMediaTypes {
			content[mime] = &kin.MediaType{Schema: schema}
		}

		headers := make(kin.Headers, len(r.headers))
		for _, name := range r.headers {
			headers[name] = &kin.HeaderRef{
				Value: &kin.Header{
					Parameter: kin.Parameter{
						Name: name,
						In:   kin.ParameterInHeader,
					},
				},
			}
		}

		responses.Set(strconv.Itoa(r.code), &kin.ResponseRef{Value: &kin.Response{
			Description: &r.description,
			Content:     content,
			Headers:     headers,
		}})
	}
	return responses, nil
}

// addSecuritySchemes derives a security scheme from the given Security struct and returns the security requirements,
// which act as a reference from an endpoint to its security scheme.
func (g *generator) addSecuritySchemes(secs map[string]Security) (*kin.SecurityRequirements, error) {
	if len(secs) == 0 {
		return nil, nil //nolint:nilnil
	}

	var reqs kin.SecurityRequirements
	for name, sec := range secs {
		switch sec.Type {
		case secTypeBasic:
			g.doc.Components.SecuritySchemes[name] = &kin.SecuritySchemeRef{
				Value: &kin.SecurityScheme{
					Type:   "http",
					Scheme: "basic",
				},
			}
			reqs = append(reqs, kin.SecurityRequirement{
				name: []string{},
			})
		case secTypeBearer:
			g.doc.Components.SecuritySchemes[name] = &kin.SecuritySchemeRef{
				Value: &kin.SecurityScheme{
					BearerFormat: sec.BearerFormat,
					Type:         "http",
					Scheme:       "bearer",
				},
			}
			reqs = append(reqs, kin.SecurityRequirement{
				name: []string{},
			})
		case secTypeAPIKey:
			g.doc.Components.SecuritySchemes[name] = &kin.SecuritySchemeRef{
				Value: &kin.SecurityScheme{
					Type: "apiKey",
					Name: sec.APIKeyName,
					In:   sec.APIKeyIn,
				},
			}
			reqs = append(reqs, kin.SecurityRequirement{
				name: []string{},
			})
		default:
			return nil, fmt.Errorf("unsupported security type %q", sec.Type)
		}
	}

	return &reqs, nil
}

type openAPIType interface {
	OpenAPISchemaType() []string
	OpenAPISchemaFormat() string
}

type oneOfTypes interface {
	OpenAPIV3OneOfTypes() []string
}

type docable interface {
	Docs() map[string]string
}

type exemplar interface {
	Example() any
}

type attrable interface {
	Attributes() map[string]string
}

type formatable interface {
	Formats() map[string]string
}

type enumerable interface {
	Enums() map[string][]string
}

// keyable is implemented by map types to declare allowed keys for the generated
// schema. The customizer applies the returned keys as a propertyNames constraint
// (OpenAPI 3.1), which restricts the keys of map/object schemas.
//
// Only applied for OpenAPI 3.1+ builds; silently omitted when building 3.0.x documents.
type keyable interface {
	AllowedKeys() []string
}

// customizer is a schema customizer factory that returns a SchemaCustomizerFn
// configured for the given OpenAPI version. Returned by BuildSpec and exposed
// for tests so the output schema for a keyable map can be inspected without
// going through the full chi walk.
func customizer(ver string) kingen.SchemaCustomizerFn {
	return func(name string, t reflect.Type, _ reflect.StructTag, schema *kin.Schema) error {
		v := reflect.New(t).Elem().Interface()

		if obj, ok := v.(openAPIType); ok {
			if err := applyType(name, schema, obj); err != nil {
				return err
			}
		}

		if obj, ok := v.(oneOfTypes); ok {
			applyOneOfTypes(schema, obj)
		}

		if obj, ok := v.(docable); ok {
			applyDocs(schema, obj)
		}

		if obj, ok := v.(exemplar); ok {
			applyExemplar(schema, obj)
		}

		if obj, ok := v.(attrable); ok {
			applyAttrs(schema, obj)
		}

		if obj, ok := v.(formatable); ok {
			applyFormats(schema, obj)
		}

		if obj, ok := v.(enumerable); ok {
			applyEnums(schema, obj)
		}

		// propertyNames is an OpenAPI 3.1 feature; only apply it when
		// the caller is building a 3.1.x document.
		if strings.HasPrefix(ver, "3.1") && t.Kind() == reflect.Map {
			if obj, ok := v.(keyable); ok {
				keys := obj.AllowedKeys()
				if len(keys) > 0 {
					schema.PropertyNames = &kin.SchemaRef{
						Value: &kin.Schema{
							Type: &kin.Types{"string"},
							Enum: stringsToAny(keys),
						},
					}
				}
			}
		}

		return nil
	}
}

func applyType(name string, schema *kin.Schema, obj openAPIType) error {
	typs := obj.OpenAPISchemaType()
	if len(typs) == 0 {
		return fmt.Errorf("type %q defines open api types by returns none", name)
	}
	schema.Type = &kin.Types{typs[0]}
	schema.Format = obj.OpenAPISchemaFormat()
	return nil
}

func applyOneOfTypes(schema *kin.Schema, obj oneOfTypes) {
	typs := obj.OpenAPIV3OneOfTypes()
	refs := make(kin.SchemaRefs, 0, len(typs))
	for _, typ := range typs {
		refs = append(refs, &kin.SchemaRef{Value: &kin.Schema{Type: &kin.Types{typ}}})
	}
	schema.OneOf = refs
}

func applyDocs(schema *kin.Schema, obj docable) {
	docs := obj.Docs()
	for k, prop := range schema.Properties {
		doc, ok := docs[k]
		if !ok {
			continue
		}
		if prop.Value == nil {
			continue
		}

		prop.Value.Description = doc
	}
}

func applyExemplar(schema *kin.Schema, obj exemplar) {
	example := obj.Example()
	if example == nil {
		return
	}
	schema.Example = example
}

func applyAttrs(schema *kin.Schema, obj attrable) {
	attrs := obj.Attributes()
	var required []string
	for k, prop := range schema.Properties {
		attr := attrs[k]
		if prop.Value == nil {
			continue
		}

		switch attr {
		case "readonly":
			prop.Value.ReadOnly = true
		case "required":
			required = append(required, k)
		}
	}
	if len(required) > 0 {
		sort.Strings(required)
		schema.Required = required
	}
}

func applyFormats(schema *kin.Schema, obj formatable) {
	fmts := obj.Formats()
	for k, prop := range schema.Properties {
		fmt := fmts[k]
		if prop.Value == nil {
			continue
		}

		prop.Value.WithFormat(fmt)
	}
}

func stringsToAny(s []string) []any {
	if len(s) == 0 {
		return nil
	}
	out := make([]any, len(s))
	for i, v := range s {
		out[i] = v
	}
	return out
}

func applyEnums(schema *kin.Schema, obj enumerable) {
	enums := obj.Enums()
	for k, prop := range schema.Properties {
		enum := enums[k]
		if len(enum) == 0 || prop.Value == nil {
			continue
		}

		enumValues := stringsToAny(enum)

		// For arrays, enum constraints belong to the item schema
		if prop.Value.Type.Is(kin.TypeArray) {
			if prop.Value.Items == nil {
				continue
			}
			// Items.Value is always non-nil here: kin-openapi only sets Value=nil for
			// $ref-promoted structs, which are never array element types for enum fields
			prop.Value.Items.Value.Enum = enumValues
			continue
		}

		// For scalar types, apply enum directly
		prop.Value.Enum = enumValues
	}
}