hasher.go

package main

import (
	"crypto/md5"
	"crypto/rand"
	"fmt"
	"math/big"
	"net/http"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/gorilla/mux"
	"github.com/guregu/dynamo"
)

type keyring struct {
	keys   map[string]string
	keyIds []string
}

var mykeys *keyring

// We have a simple DynamoDB table of {KeyName, KeyVal}
type blob struct {
	KeyName string `dynamo:"KeyName"`
	KeyVal  string `dynamo:"KeyVal"`
	TTL     int64  `dynamo:"TTL",omitempty`
}

func newKeyring() *keyring {
	q := keyring{}
	q.keys = make(map[string]string)
	q.keyIds = make([]string, 0)
	fmt.Printf("%#v\n", q)
	return &q
}

func (k *keyring) AddKey(keyId string, keyVal string) {
	k.keys[keyId] = keyVal
	k.keyIds = append(k.keyIds, keyId)
}

func (k *keyring) GenerateKey() string {
	keyId, keyVal := generateKey()
	k.AddKey(keyId, keyVal)
	// If we generate a key, we need to put it into DynamoDB
	dydbCom <- blob{keyId, keyVal, 0}
	return keyId
}

func (k *keyring) CurrentKey() (string, string) {
	keyId := k.keyIds[len(k.keyIds)-1]
	key := k.keys[keyId]

	return keyId, key
}

func (k *keyring) GetKey(kn string) string {
	return k.keys[kn]
}

// Various channels for inter-goroutine communicationals.
var dydbCom, dydbRes chan blob
var ready chan bool

// We have a need for random strings of various lengths.
// Occasionally with separators between chunks of digits.
func randomString(l int, sep bool) string {
	letters := "abcedfghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"

	out := ""
	for x := 1; x < l+1; x++ {
		i, err := rand.Int(rand.Reader, big.NewInt(62))
		if err != nil {
			panic(err)
		}
		j := i.Uint64()
		out = out + letters[j:j+1]
		if sep && x%5 == 0 && x < 21 {
			out = out + "!"
		}
	}
	return out
}

// Our keys match /.....-.....-.....-.....-...../ with an
// associated Id matching /...../
func generateKey() (string, string) {
	id := randomString(5, false)
	val := randomString(25, true)

	return id, val
}

// Helper to add a new key because we need to do this on startup as a demo
func addNewKey() string {
	keyId, keyVal := generateKey()
	mykeys.AddKey(keyId, keyVal)
	dydbCom <- blob{keyId, keyVal, 0}

	return keyId
}

func newKeyHandler(w http.ResponseWriter, r *http.Request) {
	k := addNewKey()
	w.Write([]byte(fmt.Sprintf("OK %s", k)))
}

func hashHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]

	keyId, key := mykeys.CurrentKey()

	// This should be refactored...
	hashVal := fmt.Sprintf("%s:%s:%s", key, keyId, id)
	fmt.Printf("v=%s\n", hashVal)
	hashBytes := md5.Sum([]byte(hashVal))
	hashStr := fmt.Sprintf("%x", hashBytes)

	// Probably should use a real JSON marshaller here :)
	j := fmt.Sprintf("{\"hash\": \"%s\", \"key\":\"%s\",\"id\":%s}", hashStr, keyId, id)
	w.Write([]byte(j))
}

func makeHash(id string, keyId string) string {
	dydbCom <- blob{keyId, "", 0}
	kb := <-dydbRes

	// If we don't get a key value back, we've not found it
	// locally or in DynamoDB which means the check can never
	// succeed, in which case an empty string suffices.
	if kb.KeyVal == "" {
		return ""
	}

	key := kb.KeyVal

	// ...because this is the exact same code!
	hashVal := fmt.Sprintf("%s:%s:%s", key, keyId, id)
	hashBytes := md5.Sum([]byte(hashVal))
	hashStr := fmt.Sprintf("%x", hashBytes)

	return hashStr
}

// Check a hash given an Id and a KeyId.
// .../check/{id}/{key}/{hash}
func checkHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]
	keyId := vars["key"]
	hash := vars["hash"]

	retval := "NOT OK"

	// We can only check the hash if they request a key that exists.
	ourHash := makeHash(id, keyId)
	if ourHash == hash {
		retval = "OK"
	}
	w.Write([]byte(retval))
}

// Ultra-trivial HTTP endpoint for the healthcheck
func healthCheckHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("OK"))
}

// Ultra-trivial HTTP server
func startHTTP() {
	r := mux.NewRouter()
	r.HandleFunc("/healthcheck", healthCheckHandler)
	r.HandleFunc("/hash/{id}", hashHandler)
	r.HandleFunc("/check/{id}/{key}/{hash}", checkHandler)
	r.HandleFunc("/newkey", newKeyHandler)
	r.HandleFunc("/refresh", refreshHandler)
	http.Handle("/", r)
	http.ListenAndServe("localhost:9888", nil)
}

func refreshHandler(w http.ResponseWriter, r *http.Request) {
	dydbCom <- blob{"", "", -1}
}

func refreshKeys(t dynamo.Table) {
	// 'expired' records can persist for 48 hours on DynamoDB. Best to filter
	// out anything which might have expired but not been cleaned.
	now := time.Now().Unix()

	// Read the entire table into `keys`
	var tscan []blob
	err := t.Scan().Filter("'TTL' > ?", now).All(&tscan)
	if err != nil {
		panic(err)
	}

	q := newKeyring()

	for _, v := range tscan {
		fmt.Printf("D s=%s v=%s\n", v.KeyName, v.KeyVal)
		q.AddKey(v.KeyName, v.KeyVal)
	}

	mykeys = q
}

func initDynamoDB(command chan blob, results chan blob) {
	fmt.Println("Connecting to DynamoDB")
	db := dynamo.New(session.New(), &aws.Config{Region: aws.String("eu-west-1")})
	fmt.Println("Connected!")
	table := db.Table("HashKeys")

	refreshKeys(table)

	ready <- true

	// Loop FOREVER waiting for DynamoDB queries or updates.
	fmt.Printf("Waiting for commands\n")
	for {
		task := <-command
		// Currently this will never happen because we scan the entire
		// table when we startup and (in theory) we're the only person
		// who adds stuff to the table which means we're always in sync.
		// -but- if we ever run multiple copies of this, another one
		// might have added a key we're supposed to check which means
		// querying that from the database. But this needs more thinky.
		if task.TTL == -1 {
			refreshKeys(table)
		} else if task.KeyVal == "" {
			fmt.Printf("query for key=%s\n", task.KeyName)
			// Do we have this key cached locally?
			kv := mykeys.GetKey(task.KeyName)
			if kv != "" {
				task.KeyVal = kv
				fmt.Printf("Cached k=%s v=%s\n", task.KeyName, kv)
			} else {
				// We didn't find it locally, check DynamoDB
				err := table.Get("KeyName", task.KeyName).One(&task)
				// If we get an error or the item is missing, punt it.
				if err != nil {
					task.KeyVal = ""
				} else {
					// We got a real key, store it locally
					mykeys.AddKey(task.KeyName, task.KeyVal)
				}
				fmt.Printf("Queried DynamoDB for k=%s v=%s\n", task.KeyName, task.KeyVal)
			}
			fmt.Printf("%#v\n", task)
			results <- task
		} else {
			fmt.Printf("update for key=%s val=%s\n", task.KeyName, task.KeyVal)
			ttl := time.Now().Add(time.Hour)
			err := table.Put(blob{task.KeyName, task.KeyVal, ttl.Unix()}).Run()
			if err != nil {
				panic(err)
			}
		}
	}
}

func main() {
	mykeys = newKeyring()

	// I'm not convinced this is the best way to do it but we want
	// to localise all the Dynamo handling to a single point and
	// that requires a goroutine which we can't wait for easily.
	ready = make(chan bool)

	// Channels for DynamoDB handling.
	dydbCom = make(chan blob)
	dydbRes = make(chan blob)
	go initDynamoDB(dydbCom, dydbRes)

	// We need to wait for the DynamoDB connection because otherwise
	// keys we want to use might not present when a request arrives.
	<-ready

	// We always create a new key when we start up because we want
	// a new key to be the last one in the `keyIds` array and it's
	// less faff to force it ourself rather than sort Dynamo rows.
	_ = addNewKey()

	go startHTTP()

	// We do nothing but handle HTTP requests. IDLE LOOP.
	for {
		time.Sleep(15)
	}
}