pir.go

package pir

import (
	"fmt"
	"os"
	"runtime"
	"runtime/debug"
	"runtime/pprof"
	"time"
	//	"math"
)

// Defines the interface for PIR with preprocessing schemes
type PIR interface {
	Name() string

	PickParams(N, d, n, logq uint64) Params
	PickParamsGivenDimensions(l, m, n, logq uint64) Params

	GetBW(info DBinfo, p Params)

	Init(info DBinfo, p Params) State
	MyInit(info DBinfo, p Params) State
	InitCompressed(info DBinfo, p Params) (State, CompressedState)
	DecompressState(info DBinfo, p Params, comp CompressedState) State

	Setup(DB *Database, shared State, p Params) (State, Msg)
	MySetup(DB *Database, shared State, p Params) (State, Msg)
	FakeSetup(DB *Database, p Params) (State, float64) // used for benchmarking online phase

	Query(i uint64, shared State, p Params, info DBinfo) (State, Msg)
	MyQuery(i []uint64, shared State, p Params, info DBinfo) (State, Msg)

	Answer(DB *Database, query MsgSlice, server State, shared State, p Params) Msg
	MyAnswer(DB *Database, query MsgSlice, server State, shared State, p Params) Msg

	Recover(i uint64, batch_index uint64, offline Msg, query Msg, answer Msg, shared State, client State,
		p Params, info DBinfo) uint64
	MyRecover(i []uint64, batch_index uint64, offline Msg, query Msg, answer Msg, shared State, client State,
		p Params, info DBinfo) []uint64

	Reset(DB *Database, p Params) // reset DB to its correct state, if modified during execution
}

// Run PIR's online phase, with a random preprocessing (to skip the offline phase).
// Gives accurate bandwidth and online time measurements.
func RunFakePIR(pi PIR, DB *Database, p Params, i []uint64,
	f *os.File, profile bool) (float64, float64, float64, float64) {
	fmt.Printf("Executing %s\n", pi.Name())
	//fmt.Printf("Memory limit: %d\n", debug.SetMemoryLimit(math.MaxInt64))
	debug.SetGCPercent(-1)

	num_queries := uint64(len(i))
	if DB.Data.Rows/num_queries < DB.Info.Ne {
		panic("Too many queries to handle!")
	}
	shared_state := pi.Init(DB.Info, p)

	fmt.Println("Setup...")
	server_state, bw := pi.FakeSetup(DB, p)
	offline_comm := bw
	runtime.GC()

	fmt.Println("Building query...")
	start := time.Now()
	var query MsgSlice
	for index, _ := range i {
		_, q := pi.Query(i[index], shared_state, p, DB.Info)
		query.Data = append(query.Data, q)
	}
	printTime(start)
	online_comm := float64(query.Size() * uint64(p.Logq) / (8.0 * 1024.0))
	fmt.Printf("\t\tOnline upload: %f KB\n", online_comm)
	bw += online_comm
	runtime.GC()

	fmt.Println("Answering query...")
	if profile {
		pprof.StartCPUProfile(f)
	}
	start = time.Now()
	answer := pi.Answer(DB, query, server_state, shared_state, p)
	elapsed := printTime(start)
	if profile {
		pprof.StopCPUProfile()
	}
	rate := printRate(p, elapsed, len(i))
	online_down := float64(answer.Size() * uint64(p.Logq) / (8.0 * 1024.0))
	fmt.Printf("\t\tOnline download: %f KB\n", online_down)
	bw += online_down
	online_comm += online_down

	runtime.GC()
	debug.SetGCPercent(100)
	pi.Reset(DB, p)

	if offline_comm+online_comm != bw {
		panic("Should not happen!")
	}

	return rate, bw, offline_comm, online_comm
}

// Run full PIR scheme (offline + online phases).
func RunPIR(pi PIR, DB *Database, p Params, i []uint64) (float64, float64) {
	fmt.Printf("Executing %s\n", pi.Name())
	hstart := time.Now()
	//fmt.Printf("Memory limit: %d\n", debug.SetMemoryLimit(math.MaxInt64))
	debug.SetGCPercent(-1)

	num_queries := uint64(len(i))
	if DB.Data.Rows/num_queries < DB.Info.Ne {
		panic("Too many queries to handle!")
	}
	// batch_sz := DB.Data.Rows / (DB.Info.Ne * num_queries) * DB.Data.Cols
	bw := float64(0)

	shared_state := pi.Init(DB.Info, p) // 根据数据库DB和LWE相关参数，创造A随机矩阵
	fmt.Printf("shared_state.Data: %v\n", *shared_state.Data[0])

	fmt.Println("Setup...")
	start := time.Now()
	server_state, offline_download := pi.Setup(DB, shared_state, p) // 计算H矩阵，并将DB元素映射到[0，p]
	printTime(start)
	comm := float64(offline_download.Size() * uint64(p.Logq) / (8.0 * 1024.0))
	fmt.Printf("\t\tOffline download: %f KB\n", comm)
	bw += comm
	runtime.GC()

	fmt.Println("Building query...")
	start = time.Now()
	var client_state []State // holding secrets
	var query MsgSlice       // holding queries
	for index, _ := range i {
		// index_to_query := i[index] + uint64(index)*batch_sz
		index_to_query := i[index]
		cs, q := pi.Query(index_to_query, shared_state, p, DB.Info) // 依次制作query语句，qu = As + e + ΔUi
		client_state = append(client_state, cs)
		query.Data = append(query.Data, q)
	}
	runtime.GC()
	printTime(start)
	comm = float64(query.Size() * uint64(p.Logq) / (8.0 * 1024.0))
	fmt.Printf("\t\tOnline upload: %f KB\n", comm)
	bw += comm
	runtime.GC()

	fmt.Println("Answering query...")
	start = time.Now()
	answer := pi.Answer(DB, query, server_state, shared_state, p) // ans = DB * qu
	elapsed := printTime(start)
	rate := printRate(p, elapsed, len(i))
	comm = float64(answer.Size() * uint64(p.Logq) / (8.0 * 1024.0))
	fmt.Printf("\t\tOnline download: %f KB\n", comm)
	bw += comm
	runtime.GC()

	pi.Reset(DB, p)
	fmt.Println("Reconstructing...")
	start = time.Now()
	for index, _ := range i {
		// index_to_query := i[index] + uint64(index)*batch_sz
		index_to_query := i[index]
		val := pi.Recover(index_to_query, uint64(index), offline_download,
			query.Data[index], answer, shared_state,
			client_state[index], p, DB.Info) // 返回指定下标的元素
		if DB.GetElem(index_to_query) != val {
			fmt.Printf("Batch %d (querying index %d -- row should be >= %d): Got %d instead of %d\n",
				index, index_to_query, DB.Data.Rows/4, val, DB.GetElem(index_to_query))
			panic("Reconstruct failed!")
		}
		fmt.Println("Simple PIR ---- Query Element Index:", index_to_query, "\tElement in Database:", DB.GetElem(index_to_query), "\tGet Element:", val)
	}
	fmt.Println("Success!")
	fmt.Println(pi.Name(), " Process Duration:")

	printTime(hstart)
	runtime.GC()
	debug.SetGCPercent(100)
	return rate, bw
}

// Run full PIR scheme (offline + online phases), where the transmission of the A matrix is compressed.
func RunPIRCompressed(pi PIR, DB *Database, p Params, i []uint64) (float64, float64) {
	fmt.Printf("Executing %s\n", pi.Name())
	//fmt.Printf("Memory limit: %d\n", debug.SetMemoryLimit(math.MaxInt64))
	debug.SetGCPercent(-1)

	num_queries := uint64(len(i))
	if DB.Data.Rows/num_queries < DB.Info.Ne {
		panic("Too many queries to handle!")
	}
	batch_sz := DB.Data.Rows / (DB.Info.Ne * num_queries) * DB.Data.Cols
	bw := float64(0)

	server_shared_state, comp_state := pi.InitCompressed(DB.Info, p)
	client_shared_state := pi.DecompressState(DB.Info, p, comp_state)

	fmt.Println("Setup...")
	start := time.Now()
	server_state, offline_download := pi.Setup(DB, server_shared_state, p)
	printTime(start)
	comm := float64(offline_download.Size() * uint64(p.Logq) / (8.0 * 1024.0))
	fmt.Printf("\t\tOffline download: %f KB\n", comm)
	bw += comm
	runtime.GC()

	fmt.Println("Building query...")
	start = time.Now()
	var client_state []State
	var query MsgSlice
	for index, _ := range i {
		index_to_query := i[index] + uint64(index)*batch_sz
		cs, q := pi.Query(index_to_query, client_shared_state, p, DB.Info)
		client_state = append(client_state, cs)
		query.Data = append(query.Data, q)
	}
	runtime.GC()
	printTime(start)
	comm = float64(query.Size() * uint64(p.Logq) / (8.0 * 1024.0))
	fmt.Printf("\t\tOnline upload: %f KB\n", comm)
	bw += comm
	runtime.GC()

	fmt.Println("Answering query...")
	start = time.Now()
	answer := pi.Answer(DB, query, server_state, server_shared_state, p)
	elapsed := printTime(start)
	rate := printRate(p, elapsed, len(i))
	comm = float64(answer.Size() * uint64(p.Logq) / (8.0 * 1024.0))
	fmt.Printf("\t\tOnline download: %f KB\n", comm)
	bw += comm
	runtime.GC()

	pi.Reset(DB, p)
	fmt.Println("Reconstructing...")
	start = time.Now()

	for index, _ := range i {
		index_to_query := i[index] + uint64(index)*batch_sz
		val := pi.Recover(index_to_query, uint64(index), offline_download,
			query.Data[index], answer, client_shared_state,
			client_state[index], p, DB.Info)

		if DB.GetElem(index_to_query) != val {
			fmt.Printf("Batch %d (querying index %d -- row should be >= %d): Got %d instead of %d\n",
				index, index_to_query, DB.Data.Rows/4, val, DB.GetElem(index_to_query))
			panic("Reconstruct failed!")
		}
	}
	fmt.Println("Success!")
	printTime(start)

	runtime.GC()
	debug.SetGCPercent(100)
	return rate, bw
}

// Run My PIR scheme (offline + online phases).
func RunMyPIR(pi PIR, DB *Database, p Params, i []uint64) (float64, float64) {
	fmt.Printf("Executing %s\n", pi.Name())
	hstart := time.Now()
	//fmt.Printf("Memory limit: %d\n", debug.SetMemoryLimit(math.MaxInt64))
	debug.SetGCPercent(-1)

	num_hash := uint64(len(i))
	if num_hash > DB.Data.Cols*DB.Info.Packing {
		panic("Too many hashes!")
	}
	// fmt.Println("pir.go 1")
	bw := float64(0)
	start := time.Now()
	shared_state_A := pi.MyInit(DB.Info, p) // 根据数据库DB和LWE相关参数，创造A随机矩阵
	printTime(start)

	fmt.Println("Setup...")
	start = time.Now()
	server_state, offline_download := pi.MySetup(DB, shared_state_A, p) // 计算H矩阵，并将DB元素映射到[0，p]
	printTime(start)
	comm := float64(offline_download.Size() * uint64(p.Logq) / (8.0 * 1024.0))
	fmt.Printf("\t\tOffline download: %f KB\n", comm)
	bw += comm
	runtime.GC()

	fmt.Println("Building query...")
	start = time.Now()
	var client_state []State // holding secrets
	var query MsgSlice       // holding queries

	cs, q := pi.MyQuery(i, shared_state_A, p, DB.Info) // 依次制作query语句，qu = As + e + ΔUi

	client_state = append(client_state, cs)
	// client_state = append(client_state, MakeState(cs.Data[0]))
	// err := cs.Data[1]
	query.Data = append(query.Data, q)

	runtime.GC()
	printTime(start)
	comm = float64(query.Size() * uint64(p.Logq) / (8.0 * 1024.0))
	fmt.Printf("\t\tOnline upload: %f KB\n", comm)
	bw += comm
	runtime.GC()

	fmt.Println("Answering query...")
	start = time.Now()
	answer := pi.MyAnswer(DB, query, server_state, shared_state_A, p) // ans = DB * qu
	elapsed := printTime(start)
	rate := printRate(p, elapsed, len(i))
	comm = float64(answer.Size() * uint64(p.Logq) / (8.0 * 1024.0))
	fmt.Printf("\t\tOnline download: %f KB\n", comm)
	bw += comm
	runtime.GC()

	fmt.Println("Reconstructing...")
	start = time.Now()
	vals := pi.MyRecover(i, uint64(0), offline_download,
		query.Data[0], answer, shared_state_A,
		client_state[0], p, DB.Info) // 返回指定下标的元素
	// if DB.GetElem(index_to_query) != val {
	// 	fmt.Printf("Batch %d (querying index %d -- row should be >= %d): Got %d instead of %d\n",
	// 		index, index_to_query, DB.Data.Rows/4, val, DB.GetElem(index_to_query))
	// 	panic("Reconstruct failed!")
	// }
	// fmt.Println("Simple PIR ---- Query Element Index:", index_to_query, "\tElement in Database:", DB.GetElem(index_to_query), "\tGet Element:", val)

	printTime(start)
	fmt.Println("ans vector:", vals)
	fmt.Println("Success!")
	fmt.Println(pi.Name(), " Process Duration:")

	printTime(hstart)
	runtime.GC()
	debug.SetGCPercent(100)
	return rate, bw
}