Merge lp:~pedronis/ubuntu-push/acceptcli-enhancements into lp:ubuntu-push/automatic

[Samuele Pedroni]

* teach acceptance/kit.CliLoop to ask for hosts based on the deviceid hash to respect conistent hashing, explicitly specifying -addr still skips this,

* added some knobs to the acceptance.ClietnSession config

* teach acceptance/kit.CliLoop to ask for hosts based on the deviceid hash to respect conistent hashing, explicitly specifying -addr still skips this,

* added some knobs to the acceptance.ClietnSession config

[Guillermo Gonzalez (verterok)]

looks good.

review: Approve

=== modified file 'server/acceptance/acceptanceclient.go'
--- server/acceptance/acceptanceclient.go	2015-04-16 12:16:51 +0000
+++ server/acceptance/acceptanceclient.go	2015-06-04 16:25:32 +0000
@@ -49,6 +49,8 @@
 	cookie          string
 	cookieLock      sync.RWMutex
 	ReportSetParams bool
+	DontClose       bool
+	SlowStart       time.Duration
 	// connection
 	Connection net.Conn
 }
@@ -101,7 +103,10 @@
 // Run the session with the server, emits a stream of events.
 func (sess *ClientSession) Run(events chan<- string) error {
 	conn := sess.Connection
-	defer conn.Close()
+	if !sess.DontClose {
+		defer conn.Close()
+	}
+	time.Sleep(sess.SlowStart)
 	conn.SetDeadline(time.Now().Add(sess.ExchangeTimeout))
 	_, err := conn.Write(wireVersionBytes)
 	if err != nil {

=== modified file 'server/acceptance/cmd/acceptanceclient.go'
--- server/acceptance/cmd/acceptanceclient.go	2014-08-20 19:48:59 +0000
+++ server/acceptance/cmd/acceptanceclient.go	2015-06-04 16:25:32 +0000
@@ -36,7 +36,7 @@
 	kit.Defaults["auth_helper"] = ""
 	kit.Defaults["wait_for"] = ""
 	cfg := &configuration{}
-	kit.CliLoop(cfg, &cfg.Configuration, func(session *acceptance.ClientSession, cfgDir string) {
+	kit.CliLoop(cfg, &cfg.Configuration, func(session *acceptance.ClientSession, apiCli *kit.APIClient, cfgDir string) {
 		log.Printf("with: %#v", session)
 	}, func(url string) string {
 		if cfg.AuthHelper == "" {

=== modified file 'server/acceptance/kit/api.go'
--- server/acceptance/kit/api.go	2015-04-16 12:16:51 +0000
+++ server/acceptance/kit/api.go	2015-06-04 16:25:32 +0000
@@ -23,8 +23,10 @@
 	"crypto/tls"
 	"encoding/json"
 	"errors"
+	"io"
 	"io/ioutil"
 	"net/http"
+	"net/url"
 )
 
 // APIClient helps making api requests.
@@ -56,6 +58,20 @@
 
 var ErrNOk = errors.New("not ok")
 
+func readBody(respBody io.ReadCloser) (map[string]interface{}, error) {
+	defer respBody.Close()
+	body, err := ioutil.ReadAll(respBody)
+	if err != nil {
+		return nil, err
+	}
+	var res map[string]interface{}
+	err = json.Unmarshal(body, &res)
+	if err != nil {
+		return nil, &APIError{err.Error(), body}
+	}
+	return res, nil
+}
+
 // Post a API request.
 func (api *APIClient) PostRequest(path string, message interface{}) (map[string]interface{}, error) {
 	packedMessage, err := json.Marshal(message)
@@ -77,18 +93,36 @@
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
-	body, err := ioutil.ReadAll(resp.Body)
+	res, err := readBody(resp.Body)
 	if err != nil {
 		return nil, err
 	}
-	var res map[string]interface{}
-	err = json.Unmarshal(body, &res)
-	if err != nil {
-		return nil, &APIError{err.Error(), body}
-	}
+
 	if ok, _ := res["ok"].(bool); !ok {
 		return res, ErrNOk
 	}
 	return res, nil
 }
+
+// Get resource from API endpoint.
+func (api *APIClient) GetRequest(path string, params map[string]string) (map[string]interface{}, error) {
+	apiURL := api.ServerAPIURL + path
+	if len(params) != 0 {
+		vals := url.Values{}
+		for k, v := range params {
+			vals.Set(k, v)
+		}
+		apiURL += "?" + vals.Encode()
+	}
+	request, _ := http.NewRequest("GET", apiURL, nil)
+
+	resp, err := api.httpClient.Do(request)
+	if err != nil {
+		return nil, err
+	}
+	res, err := readBody(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+	return res, nil
+}

=== modified file 'server/acceptance/kit/cliloop.go'
--- server/acceptance/kit/cliloop.go	2015-04-15 15:35:12 +0000
+++ server/acceptance/kit/cliloop.go	2015-06-04 16:25:32 +0000
@@ -17,6 +17,7 @@
 package kit
 
 import (
+	"crypto/tls"
 	"flag"
 	"fmt"
 	"log"
@@ -26,6 +27,8 @@
 	"strings"
 	"time"
 
+	"launchpad.net/ubuntu-push/external/murmur3"
+
 	"launchpad.net/ubuntu-push/config"
 	"launchpad.net/ubuntu-push/server/acceptance"
 )
@@ -39,6 +42,9 @@
 	CertPEMFile string                `json:"cert_pem_file"`
 	Insecure    bool                  `json:"insecure" help:"disable checking of server certificate and hostname"`
 	Domain      string                `json:"domain" help:"domain for tls connect"`
+	// api config
+	APIURL         string `json:"api"`
+	APICertPEMFile string `json:"api_cert_pem_file"`
 	// run timeout
 	RunTimeout config.ConfigTimeDuration `json:"run_timeout"`
 	// flags
@@ -66,22 +72,24 @@
 var (
 	Name     = "acceptanceclient"
 	Defaults = map[string]interface{}{
-		"target":           "",
-		"addr":             ":0",
-		"exchange_timeout": "5s",
-		"cert_pem_file":    "",
-		"insecure":         false,
-		"domain":           "",
-		"run_timeout":      "0s",
-		"reportPings":      true,
-		"model":            "?",
-		"imageChannel":     "?",
-		"buildNumber":      -1,
+		"target":            "",
+		"addr":              ":0",
+		"exchange_timeout":  "5s",
+		"cert_pem_file":     "",
+		"insecure":          false,
+		"domain":            "",
+		"run_timeout":       "0s",
+		"reportPings":       true,
+		"model":             "?",
+		"imageChannel":      "?",
+		"buildNumber":       -1,
+		"api":               "",
+		"api_cert_pem_file": "",
 	}
 )
 
 // CliLoop parses command line arguments and runs a client loop.
-func CliLoop(totalCfg interface{}, cfg *Configuration, onSetup func(sess *acceptance.ClientSession, cfgDir string), auth func(string) string, waitFor func() string, onConnect func()) {
+func CliLoop(totalCfg interface{}, cfg *Configuration, onSetup func(sess *acceptance.ClientSession, apiCli *APIClient, cfgDir string), auth func(string) string, waitFor func() string, onConnect func()) {
 	flag.Usage = func() {
 		fmt.Fprintf(os.Stderr, "Usage: %s [options] <device id>\n", Name)
 		flag.PrintDefaults()
@@ -100,26 +108,57 @@
 	case narg < 1:
 		missingArg("device-id")
 	}
+	cfgDir := filepath.Dir(flag.Lookup("cfg@").Value.String())
+	// setup api
+	apiCli := &APIClient{}
+	var apiTLSConfig *tls.Config
+	if cfg.APICertPEMFile != "" || cfg.Insecure {
+		var err error
+		apiTLSConfig, err = MakeTLSConfig("", cfg.Insecure,
+			cfg.APICertPEMFile, cfgDir)
+		if err != nil {
+			log.Fatalf("api tls config: %v", err)
+		}
+	}
+	apiCli.SetupClient(apiTLSConfig, true, 1)
+	if cfg.APIURL == "" {
+		apiCli.ServerAPIURL = cfg.PickByTarget("api",
+			"https://push.ubuntu.com",
+			"https://push.staging.ubuntu.com")
+	} else {
+		apiCli.ServerAPIURL = cfg.APIURL
+	}
+	deviceId := flag.Arg(0)
 	addr := ""
+	domain := ""
 	if cfg.Addr == ":0" {
-		addr = cfg.PickByTarget("addr", "push-delivery.ubuntu.com:443",
-			"push-delivery.staging.ubuntu.com:443")
+		hash := murmur3.Sum64([]byte(deviceId))
+		hosts, err := apiCli.GetRequest("/delivery-hosts",
+			map[string]string{
+				"h": fmt.Sprintf("%x", hash),
+			})
+		if err != nil {
+			log.Fatalf("querying hosts: %v", err)
+		}
+		addr = hosts["hosts"].([]interface{})[0].(string)
+		domain = hosts["domain"].(string)
+		log.Printf("using: %s %s", addr, domain)
 	} else {
 		addr = cfg.Addr.HostPort()
+		domain = cfg.Domain
 	}
 	session := &acceptance.ClientSession{
 		ExchangeTimeout: cfg.ExchangeTimeout.TimeDuration(),
 		ServerAddr:      addr,
-		DeviceId:        flag.Arg(0),
+		DeviceId:        deviceId,
 		// flags
 		Model:        cfg.DeviceModel,
 		ImageChannel: cfg.ImageChannel,
 		BuildNumber:  cfg.BuildNumber,
 		ReportPings:  cfg.ReportPings,
 	}
-	cfgDir := filepath.Dir(flag.Lookup("cfg@").Value.String())
-	onSetup(session, cfgDir)
-	session.TLSConfig, err = MakeTLSConfig(cfg.Domain, cfg.Insecure, cfg.CertPEMFile, cfgDir)
+	onSetup(session, apiCli, cfgDir)
+	session.TLSConfig, err = MakeTLSConfig(domain, cfg.Insecure, cfg.CertPEMFile, cfgDir)
 	if err != nil {
 		log.Fatalf("tls config: %v", err)
 	}