AnonymousTopicBasedPubSubCommunicationforMicroblogging/follower/follower.go

package main

//#cgo CFLAGS: -fopenmp -O2
//#cgo LDFLAGS: -lcrypto -lm -fopenmp
//#include "../c/dpf.h"
//#include "../c/okv.h"
//#include "../c/dpf.c"
//#include "../c/okv.c"
import "C"

//ssssssssssssss

import (
	"2PPS/lib"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"strconv"
	"sync"
	"time"
	"unsafe"

	"golang.org/x/crypto/nacl/box"
)

const leader string = "127.0.0.1:4442"

//needs to be changed at leader/follower/client at the same time
const neededSubscriptions = 1

//this stores all neccessary information for each client
type clientKeys struct {
	SharedSecret [32]byte
	PirQuery     [][]byte
}

//uses clients publicKey as key
var clientData = make(map[[32]byte]clientKeys)
var topicList []byte
var topicAmount int

var followerPrivateKey *[32]byte
var followerPublicKey *[32]byte
var leaderPublicKey *[32]byte

var numThreads int = 12

//has to be dividable by 32
var dataLength int = 64
var dbWriteSize int = 4

var round int = 1
var startTime time.Time

var maxTimePerRound time.Duration = 5 * time.Second

func main() {

	generatedPublicKey, generatedPrivateKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		panic(err)
	}

	followerPrivateKey = generatedPrivateKey
	followerPublicKey = generatedPublicKey

	/*
		if len(os.Args) != 4 {
			fmt.Println("try again: numThreads, dataLength, numRows")
			return
		}
		numThreads, _ = strconv.Atoi(os.Args[2])
		dataLength, _ = strconv.Atoi(os.Args[3])
		numRows, _ = strconv.Atoi(os.Args[4])
	*/

	C.initializeServer(C.int(numThreads))

	followerConnectionPrivateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		panic(err)
	}

	// Generate a pem block with the private key
	keyPem := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(followerConnectionPrivateKey),
	})

	tml := x509.Certificate{
		// you can add any attr that you need
		NotBefore: time.Now(),
		NotAfter:  time.Now().AddDate(5, 0, 0),
		// you have to generate a different serial number each execution
		SerialNumber: big.NewInt(123123),
		Subject: pkix.Name{
			CommonName:   "New Name",
			Organization: []string{"New Org."},
		},
		BasicConstraintsValid: true,
	}
	cert, err := x509.CreateCertificate(rand.Reader, &tml, &tml, &followerConnectionPrivateKey.PublicKey, followerConnectionPrivateKey)
	if err != nil {
		panic(err)
	}

	// Generate a pem block with the certificate
	certPem := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: cert,
	})

	tlsCert, err := tls.X509KeyPair(certPem, keyPem)
	if err != nil {
		panic(err)
	}

	config := &tls.Config{Certificates: []tls.Certificate{tlsCert}}

	fmt.Println("start leader")

	//listens for leader
	lnLeader, err := tls.Listen("tcp", ":4443", config)
	if err != nil {
		panic(err)
	}
	defer lnLeader.Close()

	leaderConnection, err := lnLeader.Accept()
	if err != nil {
		panic(err)
	}

	//send publicKey to leader
	writeToConn(leaderConnection, followerPublicKey[:])

	//receives leader PublicKey
	var tmpLeaderPubKey [32]byte
	_, err = leaderConnection.Read(tmpLeaderPubKey[:])
	if err != nil {
		panic(err)
	}

	leaderPublicKey = &tmpLeaderPubKey

	//setup ends here

	//locks access to DB
	var m sync.Mutex
	wg := &sync.WaitGroup{}

	for {

		//phase1
		fmt.Println("phase1")

		//create write db for this round
		for i := 0; i < dbWriteSize; i++ {
			C.createDb(C.int(0), C.int(dataLength))
		}

		//receives the virtualAddresses
		virtualAddresses := make([]int, dbWriteSize+1)
		for i := 0; i <= dbWriteSize; i++ {
			virtualAddress := readFromConn(leaderConnection, 4)
			virtualAddresses[i] = byteToInt(virtualAddress)
		}

		for i := 0; i < numThreads; i++ {
			wg.Add(1)
			leaderConnection, err := lnLeader.Accept()
			if err != nil {
				panic(err)
			}
			leaderConnection.SetDeadline(time.Time{})
			startTime = time.Now()
			go phase1(i, leaderConnection, m, wg, virtualAddresses)
		}
		wg.Wait()

		//phase2
		fmt.Println("phase2")
		leaderConnection, err := lnLeader.Accept()
		if err != nil {
			panic(err)
		}
		leaderConnection.SetDeadline(time.Time{})

		phase2(leaderConnection)

		//phase3

		fmt.Println("phase3")

		if round == 1 {
			//addTestTweets()
		}

		//no tweets -> continue to phase 1 and mb get tweets
		topicList, topicAmount = lib.GetTopicList(0)
		if len(topicList) == 0 {
			continue
		}

		for i := 0; i < numThreads; i++ {
			wg.Add(1)
			leaderConnection, err := lnLeader.Accept()
			if err != nil {
				panic(err)
			}
			leaderConnection.SetDeadline(time.Time{})
			startTime = time.Now()
			go phase3(leaderConnection, wg)
		}

		wg.Wait()

		lib.CleanUpdbR(round)

		round++

	}
}

func phase1(id int, leaderWorkerConnection net.Conn, m sync.Mutex, wg *sync.WaitGroup, virtualAddresses []int) {

	gotClient := make([]byte, 1)

	for {
		gotClient = readFromConn(leaderWorkerConnection, 1)

		//this worker is done
		if gotClient[0] == 0 {
			wg.Done()
			return
		}

		//setup the worker-specific db
		dbSize := int(C.dbSize)
		db := make([][]byte, dbSize)
		for i := 0; i < dbSize; i++ {
			db[i] = make([]byte, int(C.db[i].dataSize))
		}

		//gets clients publicKey
		var clientPublicKey *[32]byte
		var tmpClientPublicKey [32]byte
		_, err := leaderWorkerConnection.Read(tmpClientPublicKey[:])
		if err != nil {
			panic(err)
		}
		clientPublicKey = &tmpClientPublicKey

		clientKeys := clientData[tmpClientPublicKey]
		clientKeys, pirQuery := handlePirQuery(clientKeys, leaderWorkerConnection, 0, tmpClientPublicKey, true)
		getSendVirtualAddress(pirQuery[0], virtualAddresses, clientKeys.SharedSecret, leaderWorkerConnection)

		clientData[*clientPublicKey] = clientKeys

		//gets dpfQuery from leader
		dpfLengthBytes := readFromConn(leaderWorkerConnection, 4)
		dpfLength := byteToInt(dpfLengthBytes)

		dpfQueryBEncrypted := make([]byte, dpfLength)

		dpfQueryBEncrypted = readFromConn(leaderWorkerConnection, dpfLength)

		//decrypt dpfQueryB for sorting into db
		var decryptNonce [24]byte
		copy(decryptNonce[:], dpfQueryBEncrypted[:24])
		dpfQueryB, ok := box.Open(nil, dpfQueryBEncrypted[24:], &decryptNonce, clientPublicKey, followerPrivateKey)
		if !ok {
			panic("dpfQueryB decryption not ok")
		}

		ds := int(C.db[0].dataSize)
		dataShareFollower := make([]byte, ds)
		pos := C.getUint128_t(C.int(virtualAddresses[dbWriteSize]))
		C.evalDPF(C.ctx[id], (*C.uchar)(&dpfQueryB[0]), pos, C.int(ds), (*C.uchar)(&dataShareFollower[0]))

		dataShareLeader := make([]byte, ds)

		writeToConn(leaderWorkerConnection, dataShareFollower)

		dataShareLeader = readFromConn(leaderWorkerConnection, ds)

		auditXOR := make([]byte, ds)
		passedAudit := true
		for i := 0; i < ds; i++ {
			auditXOR[i] = dataShareLeader[i] ^ dataShareFollower[i]

			//client tried to write to a position that is not a virtuallAddress
			if auditXOR[i] != 0 {
				passedAudit = false
			}
		}

		if passedAudit {
			//run dpf, xor into local db
			for i := 0; i < dbSize; i++ {
				ds := int(C.db[i].dataSize)
				dataShare := make([]byte, ds)
				pos := C.getUint128_t(C.int(virtualAddresses[i]))
				C.evalDPF(C.ctx[id], (*C.uchar)(&dpfQueryB[0]), pos, C.int(ds), (*C.uchar)(&dataShare[0]))
				for j := 0; j < ds; j++ {
					db[i][j] = db[i][j] ^ dataShare[j]
				}
			}

			//xor the worker's DB into the main DB
			for i := 0; i < dbSize; i++ {
				m.Lock()
				C.xorIn(C.int(i), (*C.uchar)(&db[i][0]))
				m.Unlock()
			}
		}
	}
}

func phase2(leaderWorkerConnection net.Conn) {

	//gets current seed
	seedFollower := make([]byte, 16)
	C.readSeed((*C.uchar)(&seedFollower[0]))

	//get data
	dbSize := int(C.dbSize)
	tmpdbFollower := make([][]byte, dbSize)
	for i := range tmpdbFollower {
		tmpdbFollower[i] = make([]byte, dataLength)
	}
	for i := 0; i < dbSize; i++ {
		C.readData(C.int(i), (*C.uchar)(&tmpdbFollower[i][0]))
	}

	//receive seed from leader
	seedLeader := readFromConn(leaderWorkerConnection, 16)

	//receive data from leader
	tmpdbLeader := make([][]byte, dbSize)
	for i := range tmpdbLeader {
		tmpdbLeader[i] = make([]byte, dataLength)
	}
	for i := 0; i < dbSize; i++ {
		tmpdbLeader[i] = readFromConn(leaderWorkerConnection, dataLength)
	}

	//writes seed to leader
	writeToConn(leaderWorkerConnection, seedFollower)

	//write data to leader
	for i := 0; i < dbSize; i++ {
		writeToConn(leaderWorkerConnection, tmpdbFollower[i])
	}
	//put together the db
	tmpdb := make([][]byte, dbSize)
	for i := range tmpdb {
		tmpdb[i] = make([]byte, dataLength)
	}

	//get own Ciphers
	ciphersFollowers := make([]*C.uchar, dbSize)
	for i := 0; i < dbSize; i++ {
		ciphersFollowers[i] = (*C.uchar)(C.malloc(16))
	}
	for i := 0; i < dbSize; i++ {
		C.getCipher(0, C.int(i), ciphersFollowers[i])
	}

	//receive ciphers from leader
	ciphersLeader := make([]byte, dbSize*16)
	for i := 0; i < dbSize; i++ {
		_, err := leaderWorkerConnection.Read(ciphersLeader[i*16:])
		if err != nil {
			panic(err)
		}
	}

	//send own Ciphers to leader
	for i := 0; i < dbSize; i++ {
		writeToConn(leaderWorkerConnection, C.GoBytes(unsafe.Pointer(ciphersFollowers[i]), 16))
	}

	//put in ciphers from leader
	for i := 0; i < dbSize; i++ {
		C.putCipher(0, C.int(i), (*C.uchar)(&ciphersLeader[i*16]))
	}

	for i := 0; i < dbSize; i++ {
		C.decryptRow(C.int(i), (*C.uchar)(&tmpdb[i][0]), (*C.uchar)(&tmpdbLeader[i][0]), (*C.uchar)(&tmpdbFollower[i][0]), (*C.uchar)(&seedLeader[0]), (*C.uchar)(&seedFollower[0]))
	}

	var tweets []lib.Tweet
	for i := 0; i < dbSize; i++ {
		//discard cover message
		if tmpdb[i][0] == 0 {
			continue
		} else {
			//reconstruct tweet
			var position int = 0
			var topics []string
			var topic string
			var text string
			for _, letter := range tmpdb[i] {
				if string(letter) == ";" {
					if topic != "" {
						topics = append(topics, topic)
						topic = ""
					}
					position++
				} else {
					if position == 0 {
						if string(letter) == "," {
							topics = append(topics, topic)
							topic = ""
						} else {
							topic = topic + string(letter)
						}
					} else if position == 1 {
						text = text + string(letter)
					}
				}
			}

			tweet := lib.Tweet{"", -1, topics, text, round}

			tweets = append(tweets, tweet)
		}
	}

	//fmt.Println("tweets recovered: ", tweets)

	//sort into read db
	lib.NewEntries(tweets, 0)

	//reset write db after the tweets were moved to read db
	C.resetDb()

	//gets current dbWriteSize from leader
	dbWriteSizeBytes := readFromConn(leaderWorkerConnection, 4)
	dbWriteSize = byteToInt(dbWriteSizeBytes)
}

func addTestTweets() {
	//creates test tweets
	tweets := make([]lib.Tweet, 5)

	for i := range tweets {
		j := i
		if i == 1 {
			j = 0
		}
		text := "Text " + strconv.Itoa(i)
		var topics []string
		topics = append(topics, "Topic "+strconv.Itoa(j))
		tweets[i] = lib.Tweet{"", -1, topics, text, i}
	}

	lib.NewEntries(tweets, 0)
}

func phase3(leaderWorkerConnection net.Conn, wg *sync.WaitGroup) {

	gotClient := make([]byte, 1)

	for {
		gotClient = readFromConn(leaderWorkerConnection, 1)

		//this worker is done
		if gotClient[0] == 0 {
			wg.Done()
			return
		}

		subPhase := readFromConn(leaderWorkerConnection, 1)

		var clientPublicKey [32]byte
		_, err := leaderWorkerConnection.Read(clientPublicKey[:])
		if err != nil {
			panic(err)
		}

		//gets the client data
		clientKeys := clientData[clientPublicKey]

		if subPhase[0] == 0 || subPhase[0] == 1 {
			clientKeys, _ = handlePirQuery(clientKeys, leaderWorkerConnection, int(subPhase[0]), clientPublicKey, false)
		}

		getSendTweets(clientKeys, nil, leaderWorkerConnection)

		wantsArchive := readFromConn(leaderWorkerConnection, 1)

		if wantsArchive[0] == 1 {
			_, archiveQuerys := handlePirQuery(clientKeys, leaderWorkerConnection, -1, clientPublicKey, false)
			getSendTweets(clientKeys, archiveQuerys, leaderWorkerConnection)

		}

		//saves clientKeys
		clientData[clientPublicKey] = clientKeys
	}
}

func getSendTweets(clientKeys clientKeys, archiveQuerys [][]byte, leaderWorkerConnection net.Conn) {
	tmpNeededSubscriptions := neededSubscriptions
	if archiveQuerys != nil {
		tmpNeededSubscriptions = len(archiveQuerys)
	}
	for i := 0; i < tmpNeededSubscriptions; i++ {
		//gets all requested tweets
		var tweets []byte
		if archiveQuerys == nil {
			tweets = lib.GetTweets(clientKeys.PirQuery[i], dataLength, 0)
		} else {
			tweets = lib.GetTweets(archiveQuerys[i], dataLength, 1)
		}

		//expand sharedSecret so it is of right length
		expandBy := len(tweets) / 32
		var expandedSharedSecret []byte
		for i := 0; i < expandBy; i++ {
			expandedSharedSecret = append(expandedSharedSecret, clientKeys.SharedSecret[:]...)
		}

		//Xor's sharedSecret with all tweets
		lib.Xor(expandedSharedSecret[:], tweets)

		lib.Xor(tweets, expandedSharedSecret[:])

		//sends tweets to leader
		tweetsLengthBytes := intToByte(len(tweets))

		writeToConn(leaderWorkerConnection, tweetsLengthBytes)

		writeToConn(leaderWorkerConnection, tweets)
	}
}

func handlePirQuery(clientKeys clientKeys, leaderWorkerConnection net.Conn, subPhase int, clientPublicKey [32]byte, doAuditing bool) (clientKeys, [][]byte) {

	archiveNeededSubscriptions := make([]byte, 4)
	if subPhase == -1 {
		archiveNeededSubscriptions = readFromConn(leaderWorkerConnection, 4)
	}

	//gets the msg length
	msgLengthBytes := readFromConn(leaderWorkerConnection, 4)
	msgLength := byteToInt(msgLengthBytes)

	//gets the message
	message := readFromConn(leaderWorkerConnection, msgLength)

	var decryptNonce [24]byte
	copy(decryptNonce[:], message[:24])
	decrypted, ok := box.Open(nil, message[24:], &decryptNonce, &clientPublicKey, followerPrivateKey)
	if !ok {
		panic("pirQuery decryption not ok")
	}

	//gets sharedSecret
	if subPhase == 0 {
		//bs!
		var newSharedSecret [32]byte
		for index := 0; index < 32; index++ {
			newSharedSecret[index] = decrypted[index]
		}
		clientKeys.SharedSecret = newSharedSecret
		decrypted = decrypted[32:]

		if doAuditing {
			result := make([][]byte, 1)
			result[0] = decrypted
			return clientKeys, result
		}

		//follower updates sharedSecret
	} else if subPhase == 1 {
		sharedSecret := clientKeys.SharedSecret
		sharedSecret = sha256.Sum256(sharedSecret[:])
		clientKeys.SharedSecret = sharedSecret
	}

	//follower expects pirQuery

	//transforms byteArray to ints of wanted topics
	tmpNeededSubscriptions := neededSubscriptions
	tmpTopicAmount := topicAmount
	if subPhase == -1 {
		tmpNeededSubscriptions = byteToInt(archiveNeededSubscriptions)
		_, tmpTopicAmount = lib.GetTopicList(1)
	}

	pirQueryFlattened := decrypted
	pirQuerys := make([][]byte, tmpNeededSubscriptions)
	for i := range pirQuerys {
		pirQuerys[i] = make([]byte, tmpTopicAmount)
	}
	for i := 0; i < tmpNeededSubscriptions; i++ {
		pirQuerys[i] = pirQueryFlattened[i*tmpTopicAmount : (i+1)*tmpTopicAmount]
	}

	//sets the pirQuery for the client in case whe are not archiving
	if subPhase != -1 {
		clientKeys.PirQuery = pirQuerys
	}

	return clientKeys, pirQuerys
}

func getSendVirtualAddress(pirQuery []byte, virtualAddresses []int, sharedSecret [32]byte, leaderWorkerConnection net.Conn) {
	//xores all requested addresses into virtuallAddress
	virtualAddress := make([]byte, 4)

	for index, num := range pirQuery {
		if num == 1 {
			currentAddress := intToByte(virtualAddresses[index])
			for i := 0; i < 4; i++ {
				virtualAddress[i] = virtualAddress[i] ^ currentAddress[i]
			}
		}
	}

	//xores the sharedSecret
	for i := 0; i < 4; i++ {
		virtualAddress[i] = virtualAddress[i] ^ sharedSecret[i]
	}

	writeToConn(leaderWorkerConnection, virtualAddress)
}

//sends the array to the connection
func writeToConn(connection net.Conn, array []byte) {
	_, err := connection.Write(array)
	if err != nil {
		panic(err)
	}
}

//reads an array which is returned and of size "size" from the connection
func readFromConn(connection net.Conn, size int) []byte {
	array := make([]byte, size)
	_, err := connection.Read(array)
	if err != nil {
		panic(err)
	}
	return array
}

func transformBytesToStringArray(topicsAsBytes []byte) []string {
	var topics []string
	var topic string
	var position int = 0
	for _, letter := range topicsAsBytes {
		if string(letter) == "," {
			topics[position] = topic
			topic = ""
			position++
		} else {
			topic = topic + string(letter)
		}
	}
	return topics
}

func byteToInt(myBytes []byte) (x int) {
	x = int(myBytes[3])<<24 + int(myBytes[2])<<16 + int(myBytes[1])<<8 + int(myBytes[0])
	return
}

func intToByte(myInt int) (retBytes []byte) {
	retBytes = make([]byte, 4)
	retBytes[3] = byte((myInt >> 24) & 0xff)
	retBytes[2] = byte((myInt >> 16) & 0xff)
	retBytes[1] = byte((myInt >> 8) & 0xff)
	retBytes[0] = byte(myInt & 0xff)
	return
}