From a1591d21d739c0d920f09d7c257158646446a9a9 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Tue, 2 Sep 2014 14:23:43 -0700
Subject: [PATCH] agent: Support for checkpoint

---
 command/agent/command.go     | 41 ++++++++++++++++++++++++++++++++++++
 command/agent/config.go      |  8 +++++++
 command/agent/config_test.go | 22 ++++++++++++++++---
 3 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/command/agent/command.go b/command/agent/command.go
index 8d4c0b0ba..8058bfb0a 100644
--- a/command/agent/command.go
+++ b/command/agent/command.go
@@ -7,6 +7,7 @@ import (
 	"net"
 	"os"
 	"os/signal"
+	"path/filepath"
 	"regexp"
 	"runtime"
 	"strings"
@@ -15,6 +16,7 @@ import (
 
 	"github.com/armon/go-metrics"
 	"github.com/hashicorp/consul/watch"
+	"github.com/hashicorp/go-checkpoint"
 	"github.com/hashicorp/go-syslog"
 	"github.com/hashicorp/logutils"
 	"github.com/mitchellh/cli"
@@ -293,9 +295,48 @@ func (c *Command) setupAgent(config *Config, logOutput io.Writer, logWriter *log
 		c.dnsServer = server
 	}
 
+	// Setup update checking
+	if !config.DisableUpdateCheck {
+		updateParams := &checkpoint.CheckParams{
+			Product: "consul",
+			Version: fmt.Sprintf("%s%s", config.Version, config.VersionPrerelease),
+		}
+		if !config.DisableAnonymousSignature {
+			updateParams.SignatureFile = filepath.Join(config.DataDir, "checkpoint-signature")
+		}
+
+		// Schedule a periodic check with expected interval of 24 hours
+		checkpoint.CheckInterval(updateParams, 24*time.Hour, c.checkpointResults)
+
+		// Do an immediate check within the next 30 seconds
+		go func() {
+			time.Sleep(randomStagger(30 * time.Second))
+			c.checkpointResults(checkpoint.Check(updateParams))
+		}()
+	}
+
 	return nil
 }
 
+// checkpointResults is used to handler periodic results from our update checker
+func (c *Command) checkpointResults(results *checkpoint.CheckResponse, err error) {
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Failed to check for updates: %v", err))
+		return
+	}
+	if results.Outdated {
+		c.Ui.Error(fmt.Sprintf("Newer Consul version available: %s", results.CurrentVersion))
+	}
+	for _, alert := range results.Alerts {
+		switch alert.Level {
+		case "info":
+			c.Ui.Info(fmt.Sprintf("Bulletin [%s]: %s (%s)", alert.Level, alert.Message, alert.URL))
+		default:
+			c.Ui.Error(fmt.Sprintf("Bulletin [%s]: %s (%s)", alert.Level, alert.Message, alert.URL))
+		}
+	}
+}
+
 // startupJoin is invoked to handle any joins specified to take place at start time
 func (c *Command) startupJoin(config *Config) error {
 	if len(config.StartJoin) == 0 {
diff --git a/command/agent/config.go b/command/agent/config.go
index f2d55cf66..a85591e62 100644
--- a/command/agent/config.go
+++ b/command/agent/config.go
@@ -255,6 +255,14 @@ type Config struct {
 	// feature. This is for security to prevent unknown scripts from running.
 	DisableRemoteExec bool `mapstructure:"disable_remote_exec"`
 
+	// DisableUpdateCheck is used to turn off the automatic update and
+	// security bulletin checking.
+	DisableUpdateCheck bool `mapstructure:"disable_update_check"`
+
+	// DisableAnonymousSignature is used to turn off the anonymous signature
+	// send with the update check. This is used to deduplicate messages.
+	DisableAnonymousSignature bool `mapstructure:"disable_anonymous_signature"`
+
 	// AEInterval controls the anti-entropy interval. This is how often
 	// the agent attempts to reconcile it's local state with the server'
 	// representation of our state. Defaults to every 60s.
diff --git a/command/agent/config_test.go b/command/agent/config_test.go
index 4ef8f463f..e12178af2 100644
--- a/command/agent/config_test.go
+++ b/command/agent/config_test.go
@@ -447,6 +447,20 @@ func TestDecodeConfig(t *testing.T) {
 	if config.Addresses.RPC != "127.0.0.1" {
 		t.Fatalf("bad: %#v", config)
 	}
+
+	// Disable updates
+	input = `{"disable_update_check": true, "disable_anonymous_signature": true}`
+	config, err = DecodeConfig(bytes.NewReader([]byte(input)))
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	if !config.DisableUpdateCheck {
+		t.Fatalf("bad: %#v", config)
+	}
+	if !config.DisableAnonymousSignature {
+		t.Fatalf("bad: %#v", config)
+	}
 }
 
 func TestDecodeConfig_Service(t *testing.T) {
@@ -613,9 +627,11 @@ func TestMergeConfig(t *testing.T) {
 				"handler": "foobar",
 			},
 		},
-		DisableRemoteExec: true,
-		StatsiteAddr:      "127.0.0.1:7250",
-		StatsdAddr:        "127.0.0.1:7251",
+		DisableRemoteExec:         true,
+		StatsiteAddr:              "127.0.0.1:7250",
+		StatsdAddr:                "127.0.0.1:7251",
+		DisableUpdateCheck:        true,
+		DisableAnonymousSignature: true,
 	}
 
 	c := MergeConfig(a, b)