Support Preferred Example Header

- Added support for example=something in Prefer request header
- Will fallback to a random example if not found
- Correctly support the Prefer header syntax specification, status doesn't have to be the first item anymore

Author: Jamie Curnow

apisprout.go

@@ -12,6 +12,7 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strconv"
 	"strings"
 	"time"
@@ -137,21 +138,32 @@ func addParameter(flags *pflag.FlagSet, name, short string, def interface{}, des
 
 // getTypedExample will return an example from a given media type, if such an
 // example exists. If multiple examples are given, then one is selected at
-// random.
-func getTypedExample(mt *openapi3.MediaType) (interface{}, error) {
+// random unless an "example" item exists in the Prefer header
+func getTypedExample(mt *openapi3.MediaType, prefer map[string]string) (interface{}, error) {
 	if mt.Example != nil {
 		return mt.Example, nil
 	}
 
 	if len(mt.Examples) > 0 {
+		// If preferred example requested and it it exists, return it
+		preferredExample := ""
+		if mapContainsKey(prefer, "example") {
+			preferredExample = prefer["example"]
+			if _, ok := mt.Examples[preferredExample]; ok {
+				return mt.Examples[preferredExample].Value.Value, nil
+			}
+		}
+
 		// Choose a random example to return.
 		keys := make([]string, 0, len(mt.Examples))
 		for k := range mt.Examples {
 			keys = append(keys, k)
 		}
 
-		selected := keys[rand.Intn(len(keys))]
-		return mt.Examples[selected].Value.Value, nil
+		if len(keys) > 0 {
+			selected := keys[rand.Intn(len(keys))]
+			return mt.Examples[selected].Value.Value, nil
+		}
 	}
 
 	if mt.Schema != nil {
@@ -163,11 +175,12 @@ func getTypedExample(mt *openapi3.MediaType) (interface{}, error) {
 }
 
 // getExample tries to return an example for a given operation.
-func getExample(negotiator *ContentNegotiator, prefer string, op *openapi3.Operation) (int, string, map[string]*openapi3.HeaderRef, interface{}, error) {
+// Using the Prefer http header, the consumer can specify the type of response they want.
+func getExample(negotiator *ContentNegotiator, prefer map[string]string, op *openapi3.Operation) (int, string, map[string]*openapi3.HeaderRef, interface{}, error) {
 	var responses []string
 	var blankHeaders = make(map[string]*openapi3.HeaderRef)
 
-	if prefer == "" {
+	if !mapContainsKey(prefer, "status") {
 		// First, make a list of responses ordered by successful (200-299 status code)
 		// before other types.
 		success := make([]string, 0)
@@ -181,10 +194,10 @@ func getExample(negotiator *ContentNegotiator, prefer string, op *openapi3.Opera
 		}
 		responses = append(success, other...)
 	} else {
-		if op.Responses[prefer] == nil {
+		if op.Responses[prefer["status"]] == nil {
 			return 0, "", blankHeaders, nil, ErrNoExample
 		}
-		responses = []string{prefer}
+		responses = []string{prefer["status"]}
 	}
 
 	// Now try to find the first example we can and return it!
@@ -207,7 +220,7 @@ func getExample(negotiator *ContentNegotiator, prefer string, op *openapi3.Opera
 				continue
 			}
 
-			example, err := getTypedExample(content)
+			example, err := getTypedExample(content, prefer)
 			if err == nil {
 				return status, mt, response.Value.Headers, example, nil
 			}
@@ -311,6 +324,66 @@ func load(uri string, data []byte) (swagger *openapi3.Swagger, router *openapi3f
 	return
 }
 
+// parsePreferHeader takes the value of a prefer header and splits it out into key value pairs
+//
+// HTTP Prefer header specification examples:
+// - Prefer: status=200; example="something"
+// - Prefer: example=something;status=200;
+// - Prefer: example="somet,;hing";status=200;
+//
+// As part of the Prefer specification, it is completely valid to specify
+// multiple Prefer headers in a single request, however we won't be
+// supporting that for the moment and only the first Prefer header
+// will be used.
+func parsePreferHeader(value string) map[string]string {
+	prefer := map[string]string{}
+	if value != "" {
+		// In the event that something is quoted, we want to pull those items out of the string
+		// and save them for later, so they don't conflict with other splitting logic.
+
+		quotedRegex := regexp.MustCompile(`"[^"]*"`)
+		splitRegex := regexp.MustCompile(`(,|;| )`)
+		wilcardRegex := regexp.MustCompile(`%%([0-9]+)%%`)
+
+		quotedStrings := quotedRegex.FindAllString(value, -1)
+		if len(quotedStrings) > 0 {
+			// replace each quoted string with a placehoder
+			for idx, quotedString := range quotedStrings {
+				value = strings.Replace(value, quotedString, fmt.Sprintf("%%%%%v%%%%", idx), 1)
+			}
+		}
+
+		pairs := splitRegex.Split(value, -1)
+		for _, pair := range pairs {
+			pair = strings.TrimSpace(pair)
+			if pair != "" {
+				// Put any wildcards back
+				wildcardStrings := wilcardRegex.FindAllStringSubmatch(pair, -1)
+				for _, wildcard := range wildcardStrings {
+					quotedIdx, _ := strconv.Atoi(wildcard[1])
+					pair = strings.Replace(pair, wildcard[0], quotedStrings[quotedIdx], 1)
+				}
+
+				// Determine the key and valid for this argument
+				if strings.Contains(pair, "=") {
+					eqIdx := strings.Index(pair, "=")
+					prefer[pair[:eqIdx]] = strings.Trim(pair[eqIdx+1:], `"`)
+				} else {
+					prefer[pair] = ""
+				}
+			}
+		}
+	}
+	return prefer
+}
+
+func mapContainsKey(dict map[string]string, key string) bool {
+	if _, ok := dict[key]; ok {
+		return true
+	}
+	return false
+}
+
 // server loads an OpenAPI file and runs a mock server using the paths and
 // examples defined in the file.
 func server(cmd *cobra.Command, args []string) {
@@ -535,12 +608,7 @@ func server(cmd *cobra.Command, args []string) {
 			}
 		}
 
-		prefer := req.Header.Get("Prefer")
-		if strings.HasPrefix(prefer, "status=") {
-			prefer = prefer[7:10]
-		} else {
-			prefer = ""
-		}
+		prefer := parsePreferHeader(req.Header.Get("Prefer"))
 
 		status, mediatype, headers, example, err := getExample(negotiator, prefer, route.Operation)
 		if err != nil {

apisprout_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/getkin/kin-openapi/openapi3"
@@ -81,3 +82,81 @@ func TestAddLocalServers(t *testing.T) {
 		})
 	}
 }
+
+func TestParsePreferHeader(t *testing.T) {
+	tests := []struct {
+		name   string
+		header string
+		want   map[string]string
+	}{
+		{
+			name:   "Single",
+			header: "status=200",
+			want: map[string]string{
+				"status": "200",
+			},
+		},
+		{
+			name:   "Single Quotes",
+			header: "status=\"200\"",
+			want: map[string]string{
+				"status": "200",
+			},
+		},
+		{
+			name:   "Single Quotes Space",
+			header: "example=\"in progress\"",
+			want: map[string]string{
+				"example": "in progress",
+			},
+		},
+		{
+			name:   "Multiple Semicolon",
+			header: "status=200;example=complete",
+			want: map[string]string{
+				"status":  "200",
+				"example": "complete",
+			},
+		},
+		{
+			name:   "Multiple Semi Space",
+			header: "status=200; example=complete",
+			want: map[string]string{
+				"status":  "200",
+				"example": "complete",
+			},
+		},
+		{
+			name:   "Multiple Comma",
+			header: "status=200,example=complete",
+			want: map[string]string{
+				"status":  "200",
+				"example": "complete",
+			},
+		},
+		{
+			name:   "Multiple Comma Space",
+			header: "status=200, example=complete",
+			want: map[string]string{
+				"status":  "200",
+				"example": "complete",
+			},
+		},
+		{
+			name:   "Mixed Pairs",
+			header: "example=complete; foo, status=\"200\",",
+			want: map[string]string{
+				"example": "complete",
+				"foo":     "",
+				"status":  "200",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := parsePreferHeader(tt.header); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("parsePreferHeader() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}