1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111

use super::prove::ProveFuture;
use super::verify::VerifyFuture;
use crate::Error;

use std::convert::TryInto;
use std::future::Future;
use std::io::{self, Write};
use std::os::unix::net::UnixStream;
use std::pin::Pin;
use std::task::{Context, Poll};

use dusk_tlv::{TlvReader, TlvWriter};
use dusk_uds::{Message, TaskProvider};

macro_rules! try_result_future {
    ($e:expr) => {
        match $e {
            Ok(a) => a,
            Err(e) => {
                error!("Error resolving the request: {}", e);
                return Poll::Ready(Message::Error);
            }
        }
    };
}

macro_rules! try_poll {
    ($e:expr, $c:expr) => {
        match Pin::new(&mut $e).poll($c) {
            Poll::Ready(a) => {
                trace!("Request resolved");
                a
            }
            Poll::Pending => {
                trace!("Pending request discarded");
                return Poll::Pending;
            }
        }
    };
}

pub struct MainFuture {
    socket: Option<UnixStream>,
}

impl Default for MainFuture {
    fn default() -> Self {
        MainFuture { socket: None }
    }
}

impl Clone for MainFuture {
    fn clone(&self) -> Self {
        MainFuture::default()
    }
}

impl TaskProvider for MainFuture {
    fn set_socket(&mut self, socket: UnixStream) {
        self.socket.replace(socket);
    }
}

impl Future for MainFuture {
    type Output = Message;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match &mut self.socket {
            Some(s) => {
                let mut reader = TlvReader::new(s);
                // Fetch the full request
                let request = reader.next().transpose();
                let request = try_result_future!(request);
                let request = request.ok_or(Error::Io(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    "The request was not provided",
                )));
                let request = try_result_future!(request);
                let s = reader.into_inner();

                let opcode = request[0];

                // Proof
                if opcode == 1 {
                    let proof = try_poll!(ProveFuture::new(&request[1..]), cx);
                    let proof = try_result_future!(proof);
                    let proof: Vec<u8> = try_result_future!(proof.try_into());

                    let mut writer = TlvWriter::new(s);
                    try_result_future!(writer.write(proof.as_slice()));

                    Poll::Ready(Message::Success)
                // Verify
                } else if opcode == 2 {
                    let verify = try_poll!(VerifyFuture::new(&request[1..]), cx).is_ok();
                    let verify = if verify { 0x01u8 } else { 0x00u8 };

                    let mut writer = TlvWriter::new(s);
                    try_result_future!(writer.write(&[verify]));

                    Poll::Ready(Message::Success)
                // Undefined operation
                } else {
                    try_result_future!(Err(Error::Other("Undefined operation code".to_owned())))
                }
            }

            None => try_result_future!(Err(Error::Other("No socket provided".to_owned()))),
        }
    }
}