Merge pull request #470 from kkaneda/kkaneda/permission_empty_key

Change MarshalResponse and UnmarshalRequest to gracefully handle a proto type check error

server/permission_test.go

@@ -24,6 +24,7 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"testing"
 
 	"github.com/cockroachdb/cockroach/proto"
 	"github.com/cockroachdb/cockroach/storage/engine"
@@ -268,3 +269,60 @@ func ExamplePermContentTypes() {
 	// - readwrite
 	// - writeonly
 }
+
+// TestPermEmptyKey verifies that the Accept header can be used
+// to control the format of the response when a key is empty.
+func TestPermEmptyKey(t *testing.T) {
+	httpServer := startAdminServer()
+	defer httpServer.Close()
+
+	config := &proto.PermConfig{}
+	err := yaml.Unmarshal([]byte(testPermConfig), config)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	body, err := json.MarshalIndent(config, "", "  ")
+	if err != nil {
+		t.Fatal(err)
+	}
+	keys := []string{"key0", "key1"}
+	for _, key := range keys {
+		req, err := http.NewRequest("POST", fmt.Sprintf("%s://%s%s/%s", adminScheme, testContext.Addr, permPathPrefix, key), bytes.NewReader(body))
+		req.Header.Set("Content-Type", "application/json")
+		if _, err = sendAdminRequest(req); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	testCases := []struct {
+		accept, expBody string
+	}{
+		{"application/json", `[
+  "",
+  "key0",
+  "key1"
+]`},
+		{"text/yaml", `- ""
+- key0
+- key1
+`},
+		{"application/x-protobuf", `[
+  "",
+  "key0",
+  "key1"
+]`},
+	}
+
+	for i, test := range testCases {
+		req, err := http.NewRequest("GET", fmt.Sprintf("%s://%s%s", adminScheme, testContext.Addr, permPathPrefix), nil)
+		req.Header.Set("Accept", test.accept)
+		body, err = sendAdminRequest(req)
+		if err != nil {
+			t.Fatalf("%d: %s", i, err)
+		}
+		if string(body) != test.expBody {
+			t.Errorf("%d: expected %q; got %q", i, test.expBody, body)
+		}
+	}
+}

util/http.go

@@ -121,7 +121,11 @@ func UnmarshalRequest(r *http.Request, body []byte, value interface{}, allowed [
 		}
 	case ProtoContentType, AltProtoContentType:
 		if isAllowed(ProtoEncoding, allowed) {
-			return gogoproto.Unmarshal(body, value.(gogoproto.Message))
+			msg, ok := value.(gogoproto.Message)
+			if !ok {
+				return Errorf("unable to convert %+v to protobuf", value)
+			}
+			return gogoproto.Unmarshal(body, msg)
 		}
 	case YAMLContentType, AltYAMLContentType:
 		if isAllowed(YAMLEncoding, allowed) {
@@ -138,8 +142,8 @@ func UnmarshalRequest(r *http.Request, body []byte, value interface{}, allowed [
 // is used. The value parameter is marshalled using the response
 // encoding and the resulting body and content type are returned. If
 // the encoding could not be determined by either header, the response
-// is marshalled using JSON. An error is returned on marshalling
-// failure.
+// is marshalled using JSON. Falls back to JSON when the protobuf format
+// cannot be used for the given value.
 func MarshalResponse(r *http.Request, value interface{}, allowed []EncodingType) (
 	body []byte, contentType string, err error) {
 	// TODO(spencer): until there's a nice (free) way to parse the
@@ -179,6 +183,13 @@ func MarshalResponse(r *http.Request, value interface{}, allowed []EncodingType)
 		}
 	}
 
+	// Reset protoIdx if value cannot be converted to a protocol message
+	if protoIdx < math.MaxInt32 {
+		if _, ok := value.(gogoproto.Message); !ok {
+			protoIdx = int32(math.MaxInt32)
+		}
+	}
+
 	if protoIdx < jsonIdx && protoIdx < yamlIdx {
 		// Protobuf-encode the config.
 		contentType = ProtoContentType

util/http_test.go

@@ -193,3 +193,31 @@ func TestMarshalResponse(t *testing.T) {
 		}
 	}
 }
+
+// TestProtoEncodingError verifies that MarshalResponse and
+// UnmarshalRequest gracefully handles a protocol message type error.
+func TestProtoEncodingError(t *testing.T) {
+	req, err := http.NewRequest("GET", "http://foo.com", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	req.Header.Add(util.ContentTypeHeader, "application/x-protobuf")
+	reqBody := []byte("foo")
+	var value string
+	err = util.UnmarshalRequest(req, reqBody, value, []util.EncodingType{util.ProtoEncoding})
+	if err == nil {
+		t.Errorf("%d: unexpected success")
+	}
+
+	req.Header.Add(util.AcceptHeader, "application/x-protobuf")
+	body, cType, err := util.MarshalResponse(req, value, []util.EncodingType{util.ProtoEncoding})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if cType != "application/json" {
+		t.Errorf("unexpected content type; got %s", cType)
+	}
+	if !bytes.Equal(body, body) {
+		t.Errorf("unexpected boy; got %s", body)
+	}
+}