def test_client_agg(n_clients):
seed = prio.PRGSeed()
skA, pkA = prio.create_keypair()
skB, pkB = prio.create_keypair()
# the config is shared across all actors
config = prio.Config(133, pkA, pkB, b"test_batch")
sA = prio.Server(config, prio.PRIO_SERVER_A, skA, seed)
sB = prio.Server(config, prio.PRIO_SERVER_B, skB, seed)
client = prio.Client(config)
n_data = config.num_data_fields()
data_items = bytes([(i % 3 == 1) or (i % 5 == 1) for i in range(n_data)])
for i in range(n_clients):
for_server_a, for_server_b = client.encode(data_items)
# Setup verification
vA = sA.create_verifier(for_server_a)
vB = sB.create_verifier(for_server_b)
# Produce a packet1 and send to the other party
p1A = vA.create_verify1()
p1B = vB.create_verify1()
# Produce packet2 and send to the other party
p2A = vA.create_verify2(p1A, p1B)
p2B = vB.create_verify2(p1A, p1B)
assert vA.is_valid(p2A, p2B)
assert vB.is_valid(p2A, p2B)
sA.aggregate(vA)
sB.aggregate(vB)
t_a = sA.total_shares()
t_b = sB.total_shares()
output = prio.total_share_final(config, t_a, t_b)
expected = [item * n_clients for item in list(data_items)]
assert list(output) == expected
async def main():
Prio_init()
n_clients = 4
n_data = 133
server_secret = prio.PRGSeed()
skA, pkA = prio.create_keypair()
skB, pkB = prio.create_keypair()
cfg = prio.Config(n_data, pkA, pkB, b"test_batch")
sA = prio.Server(cfg, prio.PRIO_SERVER_A, skA, server_secret)
sB = prio.Server(cfg, prio.PRIO_SERVER_B, skB, server_secret)
data_items = bytes([(i % 3 == 1) or (i % 5 == 1) for i in range(n_data)])
logger.info("Starting asyncio prio pipeline.")
client = prio.Client(cfg)
queue_a = asyncio.Queue()
queue_b = asyncio.Queue()
await client_produce(client, data_items, queue_a, queue_b, n_clients)
consumers = asyncio.ensure_future(
asyncio.gather(server_consume(sA, queue_a, queue_b),
server_consume(sB, queue_b, queue_a)))
await asyncio.gather(queue_a.join(), queue_b.join())
t_a = sA.total_shares()
t_b = sB.total_shares()
output = prio.total_share_final(cfg, t_a, t_b)
expected = [item * n_clients for item in list(data_items)]
assert list(output) == expected
consumers.cancel()
Prio_clear()
logger.info("Done!")
async def run_server(pubkey, pvtkey, pubkey_other, server_id, n_fields,
batch_id, shared_seed):
connection = await aio_pika.connect_robust(
"amqp://*****:*****@rabbitmq:5672/")
channel = await connection.channel()
queue = await channel.declare_queue(f"prio.{server_id}")
pk = prio.PublicKey().import_hex(pubkey)
sk = prio.PrivateKey().import_hex(pvtkey, pubkey)
pk_other = prio.PublicKey().import_hex(pubkey_other)
seed = prio.PRGSeed()
seed.instance = shared_seed
config = prio.Config(n_fields, pk, pk_other, batch_id)
server = prio.Server(config, server_id, sk, seed)
cache = {}
async for message in queue:
with message.process():
pid = message.message_id
v, p1, p2 = cache.get(pid, (None, None, None))
def log(line):
logger.info("Message {}: {}".format(pid, line))
ptype = message.type
routing_key = "prio.{}".format(get_other_server(server_id))
if (ptype == "verify1" and not p1) or (ptype == "verify2"
and not p2):
log("Re-queuing message!")
await channel.default_exchange.publish(
aio_pika.Message(
body=message.body,
message_id=message.message_id,
type=message.type,
),
routing_key="prio.{}".format(server_id),
)
elif ptype == "data":
log("Generating verify packet 1")
v = server.create_verifier(message.body)
p1 = v.create_verify1()
await channel.default_exchange.publish(
aio_pika.Message(
body=pickle.dumps(p1),
message_id=message.message_id,
type="verify1",
),
routing_key=routing_key,
)
elif ptype == "verify1":
log("Generating verify packet 2")
p2 = v.create_verify2(p1, pickle.loads(message.body))
await channel.default_exchange.publish(
aio_pika.Message(
body=pickle.dumps(p2),
message_id=message.message_id,
type="verify2",
),
routing_key=routing_key,
)
elif ptype == "verify2":
if v.is_valid(p2, pickle.loads(message.body)):
log("Aggregate data")
server.aggregate(v)
else:
log("Invalid data")
del cache[pid]
else:
log("Bad message type {}".format(ptype))
cache[pid] = (v, p1, p2)
from prio_processor.prio import wrapper as prio
from prio import PrioContext
import sys
with PrioContext():
skA, pkA = prio.create_keypair()
skB, pkB = prio.create_keypair()
n_data = 133
batch_id = b"test_batch"
cfg = prio.Config(n_data, pkA, pkB, batch_id)
server_secret = prio.PRGSeed()
sA = prio.Server(cfg, prio.PRIO_SERVER_A, skA, server_secret)
sB = prio.Server(cfg, prio.PRIO_SERVER_B, skB, server_secret)
client = prio.Client(cfg)
data_items = bytes([(i % 3 == 1) or (i % 5 == 1) for i in range(n_data)])
for_server_a, for_server_b = client.encode(data_items)
# Setup verification
vA = sA.create_verifier(for_server_a)
vB = sB.create_verifier(for_server_b)
# Produce a packet1 and send to the other party
p1A = vA.create_verify1()
p1B = vB.create_verify1()
def serverB(seed, config, serverB_keypair):
sk, _ = serverB_keypair
return prio.Server(config, prio.PRIO_SERVER_B, sk, seed)