Add a mechanism to control the sticky cookie value

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/mailgun/multibuf v0.0.0-20150714184110-565402cd71fb
 	github.com/mailgun/timetools v0.0.0-20141028012446-7e6055773c51
 	github.com/mailgun/ttlmap v0.0.0-20170619185759-c1c17f74874f
+	github.com/segmentio/fasthash v1.0.3
 	github.com/sirupsen/logrus v1.4.2
 	github.com/stretchr/testify v1.5.1
 	github.com/vulcand/predicate v1.1.0

go.sum

@@ -26,6 +26,8 @@ github.com/mailgun/ttlmap v0.0.0-20170619185759-c1c17f74874f h1:ZZYhg16XocqSKPGN
 github.com/mailgun/ttlmap v0.0.0-20170619185759-c1c17f74874f/go.mod h1:8heskWJ5c0v5J9WH89ADhyal1DOZcayll8fSbhB+/9A=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/segmentio/fasthash v1.0.3 h1:EI9+KE1EwvMLBWwjpRDc+fEM+prwxDYbslddQGtrmhM=
+github.com/segmentio/fasthash v1.0.3/go.mod h1:waKX8l2N8yckOgmSsXJi7x1ZfdKZ4x7KRMzBtS3oedY=
 github.com/sirupsen/logrus v1.4.2 h1:SPIRibHv4MatM3XXNO2BJeFLZwZ2LvZgfQ5+UNI2im4=
 github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=

roundrobin/stickycookie/aes_value.go

@@ -0,0 +1,138 @@
+package stickycookie
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"encoding/base64"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+	"net/url"
+	"strconv"
+	"strings"
+	"time"
+)
+
+// AESValue manages hashed sticky value.
+type AESValue struct {
+	block cipher.AEAD
+	ttl   time.Duration
+}
+
+// NewAESValue takes a fixed-size key and returns an CookieValue or an error.
+// Key size must be exactly one of 16, 24, or 32 bytes to select AES-128, AES-192, or AES-256.
+func NewAESValue(key []byte, ttl time.Duration) (*AESValue, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	gcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, err
+	}
+
+	return &AESValue{block: gcm, ttl: ttl}, nil
+}
+
+// Get hashes the sticky value.
+func (v *AESValue) Get(raw *url.URL) string {
+	base := raw.String()
+	if v.ttl > 0 {
+		base = fmt.Sprintf("%s|%d", base, time.Now().UTC().Add(v.ttl).Unix())
+	}
+
+	// Nonce is the 64bit nanosecond-resolution time, plus 32bits of crypto/rand, for 96bits (12Bytes).
+	// Theoretically, if 2^32 calls were made in 1 nanoseconds, there might be a repeat.
+	// Adds ~765ns, and 4B heap in 1 alloc
+	nonce := make([]byte, 12)
+	binary.PutVarint(nonce, time.Now().UnixNano())
+
+	rpend := make([]byte, 4)
+	if _, err := io.ReadFull(rand.Reader, rpend); err != nil {
+		// This is a near-impossible error condition on Linux systems.
+		// An error here means rand.Reader (and thus getrandom(2), and thus /dev/urandom) returned
+		// less than 4 bytes of data. /dev/urandom is guaranteed to always return the number of
+		// bytes requested up to 512 bytes on modern kernels. Behaviour on non-Linux systems
+		// varies, of course.
+		panic(err)
+	}
+
+	for i := 0; i < 4; i++ {
+		nonce[i+8] = rpend[i]
+	}
+
+	obfuscated := v.block.Seal(nil, nonce, []byte(base), nil)
+	// We append the 12byte nonce onto the end of the message
+	obfuscated = append(obfuscated, nonce...)
+	obfuscatedStr := base64.RawURLEncoding.EncodeToString(obfuscated)
+
+	return obfuscatedStr
+}
+
+// FindURL gets url from array that match the value.
+func (v *AESValue) FindURL(raw string, urls []*url.URL) (*url.URL, error) {
+	rawURL, err := v.fromValue(raw)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, u := range urls {
+		ok, err := areURLEqual(rawURL, u)
+		if err != nil {
+			return nil, err
+		}
+
+		if ok {
+			return u, nil
+		}
+	}
+
+	return nil, nil
+}
+
+func (v *AESValue) fromValue(obfuscatedStr string) (string, error) {
+	obfuscated, err := base64.RawURLEncoding.DecodeString(obfuscatedStr)
+	if err != nil {
+		return "", err
+	}
+
+	// The first len-12 bytes is the ciphertext, the last 12 bytes is the nonce
+	n := len(obfuscated) - 12
+	if n <= 0 {
+		// Protect against range errors causing panics
+		return "", errors.New("post-base64-decoded string is too short")
+	}
+
+	nonce := obfuscated[n:]
+	obfuscated = obfuscated[:n]
+
+	raw, err := v.block.Open(nil, nonce, []byte(obfuscated), nil)
+	if err != nil {
+		return "", err
+	}
+
+	if v.ttl > 0 {
+		rawParts := strings.Split(string(raw), "|")
+		if len(rawParts) < 2 {
+			return "", fmt.Errorf("TTL set but cookie doesn't contain an expiration: '%s'", raw)
+		}
+
+		// validate the ttl
+		i, err := strconv.ParseInt(rawParts[1], 10, 64)
+		if err != nil {
+			return "", err
+		}
+
+		if time.Now().UTC().After(time.Unix(i, 0).UTC()) {
+			strTime := time.Unix(i, 0).UTC().String()
+			return "", fmt.Errorf("TTL expired: '%s' (%s)\n", raw, strTime)
+		}
+
+		raw = []byte(rawParts[0])
+	}
+
+	return string(raw), nil
+}

roundrobin/stickycookie/cookie_value.go

@@ -0,0 +1,23 @@
+package stickycookie
+
+import "net/url"
+
+// CookieValue interface to manage the sticky cookie value format.
+// It will be used by the load balancer to generate the sticky cookie value and to retrieve the matching url.
+type CookieValue interface {
+	// Get converts raw value to an expected sticky format.
+	Get(*url.URL) string
+
+	// FindURL gets url from array that match the value.
+	FindURL(string, []*url.URL) (*url.URL, error)
+}
+
+// areURLEqual compare a string to a url and check if the string is the same as the url value.
+func areURLEqual(normalized string, u *url.URL) (bool, error) {
+	u1, err := url.Parse(normalized)
+	if err != nil {
+		return false, err
+	}
+
+	return u1.Scheme == u.Scheme && u1.Host == u.Host && u1.Path == u.Path, nil
+}

roundrobin/stickycookie/fallback_value.go

@@ -0,0 +1,37 @@
+package stickycookie
+
+import (
+	"errors"
+	"net/url"
+)
+
+// FallbackValue manages hashed sticky value.
+type FallbackValue struct {
+	from CookieValue
+	to   CookieValue
+}
+
+// NewFallbackValue creates a new FallbackValue
+func NewFallbackValue(from CookieValue, to CookieValue) (*FallbackValue, error) {
+	if from == nil || to == nil {
+		return nil, errors.New("from and to are mandatory")
+	}
+
+	return &FallbackValue{from: from, to: to}, nil
+}
+
+// Get hashes the sticky value.
+func (v *FallbackValue) Get(raw *url.URL) string {
+	return v.to.Get(raw)
+}
+
+// FindURL gets url from array that match the value.
+// If it is a symmetric algorithm, it decodes the URL, otherwise it compares the ciphered values.
+func (v *FallbackValue) FindURL(raw string, urls []*url.URL) (*url.URL, error) {
+	findURL, err := v.from.FindURL(raw, urls)
+	if findURL != nil {
+		return findURL, err
+	}
+
+	return v.to.FindURL(raw, urls)
+}