Kela/server/dht.go

package main

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

// 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)
	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 value for a key from the DHT
func dhtGet(key string) ([]byte, error) {
	keyPos := keyPos(key)
	var wg sync.WaitGroup
	var ret []byte
	bestTime := 0
	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 {
				// TODO: Remove this server from DHT?
				// For sanity reasons this might be a bad idea
				return
			}
			b, err := io.ReadAll(resp.Body)
			if err != nil {
				return
			}
			err = verify(key, b)
			if err != nil {
				return
			}
			valTime, err := strconv.Atoi(strings.Split(string(b), "\n")[0])
			if err != nil {
				return
			}
			if ret == nil || valTime > bestTime {
				ret = b
				bestTime = valTime
			}
		}()
	}
	wg.Wait()
	if ret == nil {
		return nil, errors.New("id not in kvstore")
	}
	return ret, nil
}

// Put a key-value pair into the DHT
func dhtPut(key string, val []byte) error {
	err := verify(key, val)
	if err != nil {
		return err
	}
	keyPos := keyPos(key)
	for i := 0; i < 5 && i < len(peerHashes); i++ {
		j := hashToDomain[peerHashes[(keyPos+i)%len(peerHashes)]]
		go func() {
			http.Post(j + "/dht/" + key + "?direct", "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.Form.Get("direct") != "" {
		// Directly modify kvstore
		if keyPos(key) - myPos >= 5 {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		if r.Method == "GET" {
			mu.Lock()
			val, ok := kvstore[key + fmt.Sprint(time.Now().Unix() / 600)]
			mu.Unlock()
			if !ok || verify(key, val) != nil {
				w.WriteHeader(http.StatusNotFound)
				return
			}
			w.Write(val)
		} else if r.Method == "POST" {
			val, err := io.ReadAll(r.Body)
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				return
			}
			mu.Lock()
			kvstore[key + fmt.Sprint(time.Now().Unix() / 600)] = val
			mu.Unlock()
		}
		return
	}
	if r.Method == "GET" {
		val, err := dhtGet(key)
		if err != nil {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		w.Write([]byte(val))
	} else if r.Method == "POST" {
		val, err := io.ReadAll(r.Body)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		err = dhtPut(key, val)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusOK)
	}
}