Add a prefix for RTAMO in the storage key (#109)

This is needed to fix an issue that currently blocks Extended RTAMO, see: #108

By having this prefix, it should be possible to enforce a referer header check on the CDN itself.

attributioncode/validator.go

@@ -93,6 +93,12 @@ func (c *Code) URLEncode() string {
 	return url.QueryEscape(c.rawURLVals.Encode())
 }
 
+// FromRTAMO returns true when the content parameter contains a prefix for
+// RTAMO, and false otherwise.
+func (c *Code) FromRTAMO() bool {
+	return rtamo.MatchString(c.Content)
+}
+
 // Validator validates and returns santized attribution codes
 type Validator struct {
 	HMACKey string
@@ -187,7 +193,7 @@ func (v *Validator) Validate(code, sig, refererHeader string) (*Code, error) {
 		rawURLVals: vals,
 	}
 
-	if fromRTAMO(attributionCode.Content) {
+	if attributionCode.FromRTAMO() {
 		refererMatch := mozillaOrg.MatchString(refererHeader)
 
 		if !refererMatch {
@@ -217,7 +223,3 @@ func checkMAC(key, msg, msgMAC []byte) error {
 	}
 	return nil
 }
-
-func fromRTAMO(content string) bool {
-	return rtamo.MatchString(content)
-}

attributioncode/validator_test.go

@@ -152,13 +152,17 @@ func TestFromRTAMO(t *testing.T) {
 	validCodes := []string{"rta:123", "rta:abc"}
 
 	for _, v := range invalidCodes {
-		if fromRTAMO(v) {
+		c := Code{Content: v}
+
+		if c.FromRTAMO() {
 			t.Errorf("Invalid code matched regex: %s", v)
 		}
 	}
 
 	for _, v := range validCodes {
-		if !fromRTAMO(v) {
+		c := Code{Content: v}
+
+		if !c.FromRTAMO() {
 			t.Errorf("Valid code did not match regex: %s", v)
 		}
 	}

stubservice/stubhandlers/redirect.go

@@ -17,6 +17,10 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+// This prefix is prepended to the product value when we construct a storage key
+// and the attribution code contains data for RTAMO.
+const rtamoProductPrefix = "rtamo-"
+
 // redirectHandler serves redirects to modified stub binaries
 type redirectHandler struct {
 	CDNPrefix string
@@ -60,6 +64,14 @@ func (s *redirectHandler) ServeStub(w http.ResponseWriter, req *http.Request, co
 		return errors.Wrap(err, "QueryUnescape")
 	}
 
+	if code.FromRTAMO() {
+		product = rtamoProductPrefix + product
+		logrus.WithFields(logrus.Fields{
+			"prefix":  rtamoProductPrefix,
+			"product": product,
+		}).Info("Updated product value in storage key for RTAMO")
+	}
+
 	key := (s.KeyPrefix + "builds/" +
 		s3PathEscape(product) + "/" +
 		s3PathEscape(lang) + "/" +

stubservice/stubhandlers/stubhandler_test.go

@@ -139,18 +139,24 @@ func TestRedirectFull(t *testing.T) {
 		NewRedirectHandler(storage, server.URL+"/cdn/", ""),
 		&attributioncode.Validator{})
 
-	runTest := func(attributionCode, expectedCode string) {
+	runTest := func(attributionCode, referer string, expectedLocation string, expectedCode string) {
 		expectedCodeRegexp := regexp.MustCompile(expectedCode)
+		expectedLocationRegexp := regexp.MustCompile(expectedLocation)
 		recorder := httptest.NewRecorder()
 		base64Code := base64.URLEncoding.WithPadding('.').EncodeToString([]byte(attributionCode))
 		req := httptest.NewRequest("GET", `http://test/?product=firefox-stub&os=win&lang=en-US&attribution_code=`+url.QueryEscape(base64Code), nil)
+		req.Header.Set("Referer", referer)
 		svc.ServeHTTP(recorder, req)
 
-		if recorder.HeaderMap.Get("Location") == "" {
+		location := recorder.HeaderMap.Get("Location")
+		if location == "" {
 			t.Fatal("Location is not set")
 		}
+		if !expectedLocationRegexp.MatchString(location) {
+			t.Fatalf("Unexpected location, got: %s", location)
+		}
 
-		resp, err := http.Get(recorder.HeaderMap.Get("Location"))
+		resp, err := http.Get(location)
 		if err != nil {
 			t.Fatal("request failed", err)
 		}
@@ -173,14 +179,43 @@ func TestRedirectFull(t *testing.T) {
 		}
 	}
 
+	emptyReferer := ""
 	runTest(
 		`campaign=%28not+set%29&content=%28not+set%29&medium=organic&source=www.google.com`,
+		emptyReferer,
+		`/cdn/builds/firefox-stub/en-US/win/`,
 		`campaign%3D%2528not%2Bset%2529%26content%3D%2528not%2Bset%2529%26dltoken%3D[\w\d-]+%26medium%3Dorganic%26source%3Dwww.google.com`,
 	)
 	runTest(
 		`campaign=%28not+set%29&content=%28not+set%29&medium=organic&source=www.notinwhitelist.com`,
+		emptyReferer,
+		`/cdn/builds/firefox-stub/en-US/win/`,
 		`campaign%3D%2528not%2Bset%2529%26content%3D%2528not%2Bset%2529%26dltoken%3D[\w\d-]+%26medium%3Dorganic%26source%3D%2528other%2529`,
 	)
+	// We expect the product to be prefixed in the location URL below because the
+	// attribution code contains data for RTAMO and the referer header contains
+	// the right value.
+	runTest(
+		`campaign=fxa-cta-123&content=rta:value&medium=referral&source=addons.mozilla.org`,
+		`https://www.mozilla.org/`,
+		`/cdn/builds/rtamo-firefox-stub/en-US/win/`,
+		`campaign%3Dfxa-cta-123%26content%3Drta%253Avalue%26dltoken%3D[\w\d-]+%26medium%3Dreferral%26source%3Daddons.mozilla.org`,
+	)
+	// We expect no prefix because the attribution data is not related to RTAMO.
+	runTest(
+		`campaign=some-campaign&content=not-for-rtamo&medium=referral&source=addons.mozilla.org`,
+		`https://www.mozilla.org/`,
+		`/cdn/builds/firefox-stub/en-US/win/`,
+		`campaign%3Dsome-campaign%26content%3Dnot-for-rtamo%26dltoken%3D[\w\d-]+%26medium%3Dreferral%26source%3Daddons.mozilla.org`,
+	)
+	// This should not return a modified installer because the referer is not the
+	// expected one.
+	runTest(
+		`campaign=fxa-cta-123&content=rta:value&medium=referral&source=addons.mozilla.org`,
+		`https://example.org/`,
+		`\?lang=en-US&os=win&product=firefox-stub`,
+		``,
+	)
 }
 
 func TestDirectFull(t *testing.T) {