From 00bedd402086171563c04432a02a8f11012785a1 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Mon, 21 Apr 2014 15:02:36 -0700
Subject: [PATCH] agent: Fix decoding of checks. Fixes #60

---
 command/agent/agent_endpoint.go   | 20 ++++++++++++++++++--
 command/agent/catalog_endpoint.go |  4 ++--
 command/agent/http.go             | 16 ++++++++++++++--
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/command/agent/agent_endpoint.go b/command/agent/agent_endpoint.go
index 63528a4fac..6d31d786ed 100644
--- a/command/agent/agent_endpoint.go
+++ b/command/agent/agent_endpoint.go
@@ -55,7 +55,11 @@ func (s *HTTPServer) AgentForceLeave(resp http.ResponseWriter, req *http.Request
 
 func (s *HTTPServer) AgentRegisterCheck(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
 	var args CheckDefinition
-	if err := decodeBody(req, &args); err != nil {
+	// Fixup the type decode of TTL or Interval
+	decodeCB := func(raw interface{}) error {
+		return FixupCheckType(raw)
+	}
+	if err := decodeBody(req, &args, decodeCB); err != nil {
 		resp.WriteHeader(400)
 		resp.Write([]byte(fmt.Sprintf("Request decode failed: %v", err)))
 		return nil, nil
@@ -108,7 +112,19 @@ func (s *HTTPServer) AgentCheckFail(resp http.ResponseWriter, req *http.Request)
 
 func (s *HTTPServer) AgentRegisterService(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
 	var args ServiceDefinition
-	if err := decodeBody(req, &args); err != nil {
+	// Fixup the type decode of TTL or Interval if a check if provided
+	decodeCB := func(raw interface{}) error {
+		rawMap, ok := raw.(map[string]interface{})
+		if !ok {
+			return nil
+		}
+		check, ok := rawMap["check"]
+		if !ok {
+			return nil
+		}
+		return FixupCheckType(check)
+	}
+	if err := decodeBody(req, &args, decodeCB); err != nil {
 		resp.WriteHeader(400)
 		resp.Write([]byte(fmt.Sprintf("Request decode failed: %v", err)))
 		return nil, nil
diff --git a/command/agent/catalog_endpoint.go b/command/agent/catalog_endpoint.go
index ec5fc7821e..bb7d96d0bd 100644
--- a/command/agent/catalog_endpoint.go
+++ b/command/agent/catalog_endpoint.go
@@ -9,7 +9,7 @@ import (
 
 func (s *HTTPServer) CatalogRegister(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
 	var args structs.RegisterRequest
-	if err := decodeBody(req, &args); err != nil {
+	if err := decodeBody(req, &args, nil); err != nil {
 		resp.WriteHeader(400)
 		resp.Write([]byte(fmt.Sprintf("Request decode failed: %v", err)))
 		return nil, nil
@@ -30,7 +30,7 @@ func (s *HTTPServer) CatalogRegister(resp http.ResponseWriter, req *http.Request
 
 func (s *HTTPServer) CatalogDeregister(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
 	var args structs.DeregisterRequest
-	if err := decodeBody(req, &args); err != nil {
+	if err := decodeBody(req, &args, nil); err != nil {
 		resp.WriteHeader(400)
 		resp.Write([]byte(fmt.Sprintf("Request decode failed: %v", err)))
 		return nil, nil
diff --git a/command/agent/http.go b/command/agent/http.go
index 8e187f94fd..67afbe3b9f 100644
--- a/command/agent/http.go
+++ b/command/agent/http.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"github.com/hashicorp/consul/consul/structs"
+	"github.com/mitchellh/mapstructure"
 	"io"
 	"log"
 	"net"
@@ -151,9 +152,20 @@ func (s *HTTPServer) Index(resp http.ResponseWriter, req *http.Request) {
 }
 
 // decodeBody is used to decode a JSON request body
-func decodeBody(req *http.Request, out interface{}) error {
+func decodeBody(req *http.Request, out interface{}, cb func(interface{}) error) error {
+	var raw interface{}
 	dec := json.NewDecoder(req.Body)
-	return dec.Decode(out)
+	if err := dec.Decode(&raw); err != nil {
+		return err
+	}
+
+	// Invoke the callback prior to decode
+	if cb != nil {
+		if err := cb(raw); err != nil {
+			return err
+		}
+	}
+	return mapstructure.Decode(raw, out)
 }
 
 // setIndex is used to set the index response header