Initial `dataonly` work.

http/logical.go

@@ -147,7 +147,9 @@ func handleLogical(core *vault.Core, dataOnly bool, prepareRequestCallback Prepa
 }
 
 func respondLogical(w http.ResponseWriter, r *http.Request, req *logical.Request, dataOnly bool, resp *logical.Response) {
-	var httpResp interface{}
+	var httpResp *logical.HTTPResponse
+	var ret interface{}
+
 	if resp != nil {
 		if resp.Redirect != "" {
 			// If we have a redirect, redirect! We use a 307 code
@@ -156,11 +158,6 @@ func respondLogical(w http.ResponseWriter, r *http.Request, req *logical.Request
 			return
 		}
 
-		if dataOnly {
-			respondOk(w, resp.Data)
-			return
-		}
-
 		// Check if this is a raw response
 		if _, ok := resp.Data[logical.HTTPContentType]; ok {
 			respondRaw(w, r, req.Path, resp)
@@ -168,7 +165,7 @@ func respondLogical(w http.ResponseWriter, r *http.Request, req *logical.Request
 		}
 
 		if resp.WrapInfo != nil && resp.WrapInfo.Token != "" {
-			httpResp = logical.HTTPResponse{
+			httpResp = &logical.HTTPResponse{
 				WrapInfo: &logical.HTTPWrapInfo{
 					Token:           resp.WrapInfo.Token,
 					TTL:             int(resp.WrapInfo.TTL.Seconds()),
@@ -177,14 +174,22 @@ func respondLogical(w http.ResponseWriter, r *http.Request, req *logical.Request
 				},
 			}
 		} else {
-			sanitizedHttp := logical.SanitizeResponse(resp)
-			sanitizedHttp.RequestID = req.ID
-			httpResp = sanitizedHttp
+			httpResp = logical.SanitizeResponse(resp)
+			httpResp.RequestID = req.ID
+		}
+
+		ret = httpResp
+
+		if dataOnly {
+			injector := logical.HTTPSysInjector{
+				Response: httpResp,
+			}
+			ret = injector
 		}
 	}
 
 	// Respond
-	respondOk(w, httpResp)
+	respondOk(w, ret)
 	return
 }
 

logical/sanitize.go

@@ -1,6 +1,11 @@
 package logical
 
-import "time"
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"time"
+)
 
 // This logic was pulled from the http package so that it can be used for
 // encoding wrapped responses as well. It simply translates the logical request
@@ -60,3 +65,46 @@ type HTTPWrapInfo struct {
 	CreationTime    time.Time `json:"creation_time"`
 	WrappedAccessor string    `json:"wrapped_accessor,omitempty"`
 }
+
+type HTTPSysInjector struct {
+	Response *HTTPResponse
+}
+
+func (h HTTPSysInjector) MarshalJSON() ([]byte, error) {
+	j, err := json.Marshal(h.Response)
+	if err != nil {
+		return nil, err
+	}
+
+	// Fast path no data or empty data
+	if h.Response.Data == nil || len(h.Response.Data) == 0 {
+		return j, nil
+	}
+
+	// Marshaling a response will always be a JSON object, meaning it will
+	// always start with '{', so we hijack this to prepend necessary values
+
+	// Make a guess at the capacity, and write the object opener
+	buf := bytes.NewBuffer(make([]byte, 0, len(j)+len(h.Response.Data)*100))
+	buf.WriteRune('{')
+
+	for k, v := range h.Response.Data {
+		// Marshal each key/value individually
+		mk, err := json.Marshal(k)
+		if err != nil {
+			return nil, err
+		}
+		mv, err := json.Marshal(v)
+		if err != nil {
+			return nil, err
+		}
+		// Write into the final buffer. We'll never have a valid response
+		// without any fields so we can unconditionally add a comma after each.
+		buf.WriteString(fmt.Sprintf("%s: %s, ", mk, mv))
+	}
+
+	// Add the rest, without the first '{'
+	buf.Write(j[1:])
+
+	return buf.Bytes(), nil
+}