Kela/server/dht.go

package main

import (
	"bytes"
	"crypto/ed25519"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

var myHash string
var myPos int
var hashToDomain map[string]string
var peerHashes []string
var kvstore map[string][]byte

// Get the sha256sum of string as a URL-safe unpadded base64 string
func sha256sum(s string) string {
	b := sha256.Sum256([]byte(s))
	return base64.RawURLEncoding.EncodeToString(b[:])
}

// Try to peer with another server
func addPeer(peer string) error {
	peerHash := sha256sum(peer)
	// Check if already peered
	mu.Lock()
	_, ok := hashToDomain[peerHash]
	mu.Unlock()
	if ok {
		return nil
	}
	mu.Lock()
	hashToDomain[peerHash] = peer
	mu.Unlock()

	// Try request to peer
	log.Printf("%s trying to peer with %s", me, peer)
	resp, err := http.Get(peer + "/peer?peer=" + me)
	if err != nil {
		// Request failed, delete peer
		mu.Lock()
		delete(hashToDomain, peerHash)
		mu.Unlock()
		return err
	}

	log.Printf("%s successfully peered with %s", me, peer)
	mu.Lock()
	i := sort.SearchStrings(peerHashes, peerHash)
	peerHashes = append(peerHashes, "")
	copy(peerHashes[i+1:], peerHashes[i:])
	peerHashes[i] = peerHash
	myPos = sort.SearchStrings(peerHashes, me)

	// Distribute keys to new server
	for id, user := range users {
		phase := time.Now().Unix() / 600
		if keyPos(id + "\n" + fmt.Sprint(phase))-myPos < 5 {
			go http.Post(peer+"/dht/"+id+"?phase="+fmt.Sprint(phase)+"&direct=true", "application/octet-stream", bytes.NewBuffer(user.dhtVal))
		}
		if keyPos(id + "\n" + fmt.Sprint(phase+1))-myPos < 5 {
			go http.Post(peer+"/dht/"+id+"?phase="+fmt.Sprint(phase+1)+"&direct=true", "application/octet-stream", bytes.NewBuffer(user.dhtVal))
		}
	}
	mu.Unlock()

	// Read response body
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	// Try adding all peers of this peer
	newPeers := strings.Split(string(body), "\n")
	for _, newPeer := range newPeers[:len(newPeers)-1] {
		go addPeer(newPeer)
	}
	return nil
}

// Handle incoming peer requests
func peerHandler(w http.ResponseWriter, r *http.Request) {
	r.ParseForm()
	peer := r.Form.Get("peer")
	if peer == "" {
		w.WriteHeader(http.StatusBadRequest)
		return
	}
	for _, p := range hashToDomain {
		fmt.Fprintf(w, "%s\n", p)
	}
	go addPeer(peer)
}

// Find the position of a key in the DHT
func keyPos(key string) int {
	keyPos := sort.SearchStrings(peerHashes, sha256sum(key))
	if keyPos < myPos {
		keyPos += len(peerHashes)
	}
	return keyPos
}

// Get the timestamp of this val
func timestamp(val []byte) int {
	if len(val) < ed25519.SignatureSize {
		return 0
	}
	message := string(val[:len(val)-ed25519.SignatureSize])
	timestamp, err := strconv.Atoi(strings.Split(message, "\n")[0])
	if err != nil {
		return 0
	}
	return timestamp
}

// Get the value for a key from the DHT
func dhtGet(key, direct string) []byte {
	phase := fmt.Sprint(time.Now().Unix() / 600)
	keyPos := keyPos(key + "\n" + phase)
	if direct != "" && keyPos-myPos < 5 {
		// Directly read from kvstore
		mu.Lock()
		val, ok := kvstore[key+"\n"+phase]
		mu.Unlock()
		if !ok || verify(key, val) != nil {
			return nil
		}
		return val
	}

	// Contact 5 servers that store this key-value pair
	var mu sync.Mutex
	var wg sync.WaitGroup
	var latest []byte
	for i := 0; i < 5 && i < len(peerHashes); i++ {
		wg.Add(1)
		j := hashToDomain[peerHashes[(keyPos+i)%len(peerHashes)]]
		go func() {
			defer wg.Done()
			resp, err := http.Get(j + "/dht/" + key + "?direct")
			if err != nil {
				return
			}
			val, err := io.ReadAll(resp.Body)
			if err != nil {
				return
			}
			err = verify(key, val)
			if err != nil {
				return
			}
			mu.Lock()
			if latest == nil || timestamp(val) > timestamp(latest) {
				latest = val
			}
			mu.Unlock()
		}()
	}
	// Wait for all to finish or time out
	wg.Wait()
	return latest
}

// Post a key-value pair into the DHT
func dhtPost(key, phase, direct string, val []byte) error {
	err := verify(key, val)
	if err != nil {
		return err
	}
	if phase == "" {
		phase = fmt.Sprint(time.Now().Unix() / 600)
	}
	user, ok := users[key]
	if ok {
		curPhase, err := strconv.Atoi(phase)
		if err != nil {
			return err
		}
		nextPhase := time.Now().Unix()/600 + 1
		if int64(curPhase) < nextPhase && user.phase < nextPhase {
			user.phase = nextPhase
			go dhtPost(key, fmt.Sprint(nextPhase), "", val)
		}
	}

	keyPos := keyPos(key + "\n" + phase)
	if direct != "" && keyPos-myPos < 5 {
		// Directly write to kvstore
		mu.Lock()
		curVal, ok := kvstore[key+"\n"+phase]
		if !ok || timestamp(val) > timestamp(curVal) {
			kvstore[key+"\n"+phase] = val
		}
		mu.Unlock()
		return nil
	}

	// Contact 5 servers that store this key-value pair
	for i := 0; i < 5 && i < len(peerHashes); i++ {
		j := hashToDomain[peerHashes[(keyPos+i)%len(peerHashes)]]
		go http.Post(j+"/dht/"+key+"?phase="+phase+"&direct=true", "application/octet-stream", bytes.NewBuffer(val))
	}
	return nil
}

// Handle DHT requests
func dhtHandler(w http.ResponseWriter, r *http.Request) {
	key := r.URL.Path[5:]
	r.ParseForm()
	if r.Method == "GET" {
		val := dhtGet(key, r.Form.Get("direct"))
		if val == nil {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		w.Write(val)
	} else if r.Method == "POST" {
		val, err := io.ReadAll(r.Body)
		if err != nil || dhtPost(key, r.Form.Get("phase"), r.Form.Get("direct"), val) != nil {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusOK)
	}
}

// Clean out offline peers
func cleanPeers() {
	for true {
		mu.Lock()
		peer := hashToDomain[peerHashes[rand.Intn(len(peerHashes))]]
		mu.Unlock()
		_, err := http.Get(peer)
		if err != nil {
			// Bad response, so remove peer
			mu.Lock()
			i := sort.SearchStrings(peerHashes, sha256sum(peer))
			peerHashes = append(peerHashes[:i], peerHashes[i+1:]...)
			myPos = sort.SearchStrings(peerHashes, me)

			// Distribute keys on this server to other servers
			if len(peerHashes) >= 5 {
				for id, user := range users {
					phase := time.Now().Unix() / 600
					kpos := keyPos(id + "\n" + fmt.Sprint(phase))
					if kpos-i < 5 {
						server := hashToDomain[peerHashes[(kpos+4)%len(peerHashes)]]
						go http.Post(server+"/dht/"+id+"?phase="+fmt.Sprint(phase)+"&direct=true", "application/octet-stream", bytes.NewBuffer(user.dhtVal))
					}
					kpos = keyPos(id + "\n" + fmt.Sprint(phase+1))
					if kpos-i < 5 {
						server := hashToDomain[peerHashes[(kpos+4)%len(peerHashes)]]
						go http.Post(server+"/dht/"+id+"?phase="+fmt.Sprint(phase+1)+"&direct=true", "application/octet-stream", bytes.NewBuffer(user.dhtVal))
					}
				}
			}
			mu.Unlock()
		}
		time.Sleep(5 * time.Second)
	}
}

// Clean out old keys from the KVStore every minute
func cleanKVStore() {
	for true {
		// This locks the mutex for a while which sucks
		mu.Lock()
		for key := range kvstore {
			timestamp, err := strconv.Atoi(strings.Split(key, "\n")[1])
			if err != nil || int64(timestamp)+1 < time.Now().Unix()/600 {
				delete(kvstore, key)
			}
		}
		mu.Unlock()
		time.Sleep(time.Minute)
	}
}

// Redistribute key-value pairs periodically
func redistributeKeys() {
	for true {
		mu.Lock()
		for id, user := range users {
			nextPhase := time.Now().Unix()/600 + 1
			if user.phase < nextPhase {
				go dhtPost(id, fmt.Sprint(nextPhase), "", user.dhtVal)
			}
		}
		mu.Unlock()
		time.Sleep(time.Duration(rand.Intn(300)) * time.Second)
	}
}