Update registration logic to use connection-service latency record (#62)

* Update registration logic to use connection-service latency record

Query connection service url from bastion and then proxy the register
request through the connection service's latency record to support
geo-ip based registration.

Author: Sebby Lipman

VERSION

@@ -1 +1 @@
-3.0.732.28
+3.0.732.29

agent/bzero/registration/registration.go

@@ -3,7 +3,6 @@ package registration
 import (
 	"bytes"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -21,8 +20,9 @@ const (
 	BZeroConfigStorage    = "BZeroConfig"
 	BZeroRegErrorExitCode = 234
 
-	registrationEndpoint = "api/v2/targets/ssm/register"
-	prodServiceUrl       = "https://cloud.bastionzero.com/" // default
+	registrationEndpoint         = "targets/ssm/register"
+	prodServiceUrl               = "https://cloud.bastionzero.com/" // default
+	getConnectionServiceEndpoint = "api/v2/connection-service/url"
 )
 
 // This is the data sent to the Reg API
@@ -43,6 +43,10 @@ type BZeroRegResponse struct {
 	OrgProvider      string `json:"externalOrganizationProvider"`
 }
 
+type GetConnectionServiceResponse struct {
+	ConnectionServiceUrl string `json:"connectionServiceUrl"`
+}
+
 // Attempts to register as many times as is acceptable
 func Register(log logger.T, apiKey string, envName string, envId string, targetName string, serviceUrl string) (BZeroRegResponse, error) {
 	var response BZeroRegResponse
@@ -60,19 +64,8 @@ func Register(log logger.T, apiKey string, envName string, envId string, targetN
 		EnvName:    envName,
 	}
 
-	// Build Registration Endpoint
-
-	if serviceUrl == "" {
-		serviceUrl = prodServiceUrl
-	}
-	u, err := url.Parse(serviceUrl)
-	if err != nil {
-		return response, fmt.Errorf("could not parse service url: %s error: %s", serviceUrl, err)
-	}
-	u.Path = path.Join(u.Path, registrationEndpoint)
-
 	// Register with BastionZero
-	resp, err := post(log, regInfo, u.String())
+	resp, err := sendRegisterRequest(log, regInfo, serviceUrl)
 	if err != nil {
 		return response, err
 	}
@@ -93,7 +86,75 @@ func Register(log logger.T, apiKey string, envName string, envId string, targetN
 	return response, nil
 }
 
-func post(log logger.T, regInfo BZeroRegRequest, regUrl string) (*http.Response, error) {
+func sendRegisterRequest(log logger.T, regInfo BZeroRegRequest, serviceUrl string) (*http.Response, error) {
+	// Declare our variables
+	var response *http.Response
+
+	// Marshal the regInfo data so we don't do it every time
+	regInfoBytes, err := json.Marshal(regInfo)
+	if err != nil {
+		return response, fmt.Errorf("could not marshal registration request")
+	}
+
+	// Build Registration Endpoint
+	if serviceUrl == "" {
+		serviceUrl = prodServiceUrl
+	}
+
+	log.Infof("Using service url %s", serviceUrl)
+
+	// Get connection service url from bastion
+	connectionServiceUrl, connectionServiceUrlErr := getConnectionServiceUrlFromServiceUrl(log, serviceUrl)
+	if connectionServiceUrlErr != nil {
+		return &http.Response{}, connectionServiceUrlErr
+	}
+
+	u, err := url.Parse(connectionServiceUrl)
+	if err != nil {
+		return response, fmt.Errorf("could not parse connection service url: %s error: %s", connectionServiceUrl, err)
+	}
+	u.Path = path.Join(u.Path, registrationEndpoint)
+
+	log.Infof("Registration Request Body: %s", string(regInfoBytes))
+	req, err := http.NewRequest("POST", u.String(), bytes.NewBuffer(regInfoBytes))
+	if err != nil {
+		return response, err
+	}
+
+	resp, err := sendRequestWithRetry(log, req)
+	if err != nil {
+		return resp, err
+	}
+
+	return resp, nil
+}
+
+func missingResponseFields(resp BZeroRegResponse) ([]string, bool) {
+	// Print out a specific message if missing registration data
+	missing := []string{}
+	if resp.ActivationId == "" {
+		missing = append(missing, "Activation ID")
+	}
+	if resp.ActivationCode == "" {
+		missing = append(missing, "Activation Code")
+	}
+	if resp.ActivationRegion == "" {
+		missing = append(missing, "Activation Region")
+	}
+	if resp.SSMTargetId == "" {
+		missing = append(missing, "SSM Target ID")
+	}
+	if resp.OrgID == "" {
+		missing = append(missing, "Organization ID")
+	}
+	if resp.OrgProvider == "" {
+		missing = append(missing, "Organization Provider")
+	}
+
+	return missing, len(missing) == 0
+}
+
+func sendRequestWithRetry(log logger.T, req *http.Request) (*http.Response, error) {
 	// Default params
 	// Ref: https://github.com/cenkalti/backoff/blob/a78d3804c2c84f0a3178648138442c9b07665bda/exponential.go#L76
 	// DefaultInitialInterval     = 500 * time.Millisecond
@@ -110,33 +171,19 @@ func post(log logger.T, regInfo BZeroRegRequest, regUrl string) (*http.Response,
 	// Make our ticker
 	ticker := backoff.NewTicker(backoffParams)
 
-	// Declare our variables
-	var response *http.Response
-
-	// Marshal the regInfo data so we don't do it every time
-	regInfoBytes, err := json.Marshal(regInfo)
-	if err != nil {
-		return response, fmt.Errorf("could not marshal registration request")
-	}
-
-	// Keep looping through our ticker, waiting for it to tell us when to retry
 	for range ticker.C {
-		// Make our Client
+
+		// Make our http Client
 		var httpClient = &http.Client{
 			Timeout: time.Second * 10,
 		}
 
-		// Build request
-		req, err := http.NewRequest("POST", regUrl, bytes.NewBuffer(regInfoBytes))
-		if err != nil {
-			return response, fmt.Errorf("Error creating new http request: %v", err)
-		}
-
 		// Headers
 		req.Header.Add("Accept", "application/json")
 		req.Header.Add("Content-Type", "application/json")
 
-		response, err = httpClient.Do(req)
+		log.Infof("Sending request to: %s", req.URL)
+		response, err := httpClient.Do(req)
 
 		// If the status code is unauthorized, do not attempt to retry
 		if response.StatusCode == http.StatusInternalServerError ||
@@ -146,8 +193,6 @@ func post(log logger.T, regInfo BZeroRegRequest, regUrl string) (*http.Response,
 			response.StatusCode == http.StatusUnsupportedMediaType {
 
 			ticker.Stop()
-			log.Infof("Registration Endpoint: %s", regUrl)
-			log.Infof("Registration Request Body: %s", string(regInfoBytes))
 			return response, fmt.Errorf("received response code: %d, not retrying", response.StatusCode)
 		}
 
@@ -159,30 +204,37 @@ func post(log logger.T, regInfo BZeroRegRequest, regUrl string) (*http.Response,
 		return response, err
 	}
 
-	return nil, errors.New("unable to make post request")
+	return nil, fmt.Errorf("Failed to successfully make request to: %s", req.URL)
 }
 
-func missingResponseFields(resp BZeroRegResponse) ([]string, bool) {
-	// Print out a specific message if missing registration data
-	missing := []string{}
-	if resp.ActivationId == "" {
-		missing = append(missing, "Activation ID")
-	}
-	if resp.ActivationCode == "" {
-		missing = append(missing, "Activation Code")
+func getConnectionServiceUrlFromServiceUrl(log logger.T, serviceUrl string) (string, error) {
+	// Make request to bastion to get connection service url
+	u, err := url.Parse(serviceUrl)
+	if err != nil {
+		return "", fmt.Errorf("could not parse service url: %s error: %s", serviceUrl, err)
 	}
-	if resp.ActivationRegion == "" {
-		missing = append(missing, "Activation Region")
+	u.Path = path.Join(u.Path, getConnectionServiceEndpoint)
+
+	req, err := http.NewRequest("GET", u.String(), nil)
+	if err != nil {
+		return "", err
 	}
-	if resp.SSMTargetId == "" {
-		missing = append(missing, "SSM Target ID")
+
+	resp, err := sendRequestWithRetry(log, req)
+	if err != nil {
+		return "", err
 	}
-	if resp.OrgID == "" {
-		missing = append(missing, "Organization ID")
+
+	// Unmarshal the response
+	respBytes, readAllErr := ioutil.ReadAll(resp.Body)
+	if readAllErr != nil {
+		return "", fmt.Errorf("error reading body on get connection service url request: %v", readAllErr)
 	}
-	if resp.OrgProvider == "" {
-		missing = append(missing, "Organization Provider")
+
+	var getConnectionServiceResponse GetConnectionServiceResponse
+	if err := json.Unmarshal(respBytes, &getConnectionServiceResponse); err != nil {
+		return "", fmt.Errorf("malformed getConnectionService response: %s", err)
 	}
 
-	return missing, len(missing) == 0
+	return getConnectionServiceResponse.ConnectionServiceUrl, nil
 }

core/agent_parser.go

@@ -166,7 +166,7 @@ func bzeroInit(log logger.T) (exitCode int) {
 func bzeroRegistration(log logger.T) (exitCode int) {
 	// check to see if the agent is already registered
 	if !force && registration.InstanceID() != "" {
-		log.Info("Agent is already registered, it will not re-register unless provided the -f flag. Restarting...")
+		log.Info("Agent is already registered, it will not re-register unless provided the -y flag. Restarting...")
 		return restartAgent(log)
 	}