<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2015 dinowernli
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

use std::f64;

use bitstring::Bit;
use bitstring::Bitstring;

// Open questions:
// - How to declare some fields final
// - How to declare some fields mutable (can be changed for non-mut self)
// - Have the tree implement iter and simplify some methods with "fold".

/// An object capable of predicting observations and rewards based on
/// experience. Predictors have an abstract notion of history which
/// grows over time and represents the experience.
pub trait Predictor {
  /// Returns the size of the currently tracked history.
  fn history_size(&self) -> usize;

  /// Appends the provided bit string to the tracked history.
  fn update(&mut self, bitstring: &Bitstring);

  /// Reverts the context tree to a previous state by undoing update
  /// operations. The specified size must be at most the current size.
  fn revert_to_history_size(&mut self, target_size: usize);

  /// Returns the probability, given the current history, that "bits" are the
  /// next observed symbols.
  fn predict(&mut self, bits: &Bitstring) -> f64;
}

/// Context tree which computes a probability estimate for a sequence of
/// bits based on a mixture model of Markov chains. Lower-degree Markov
/// chains have higher priors in the mixture model, thus formalizing the
/// notion of Occam's Razor.
///
/// In order to get a probability for a sequence, all bits in the sequence
/// must be processed in-order.
pub struct ContextTree {
  root: Node,
  history: Bitstring,
  depth: usize,
}

impl Predictor for ContextTree {
  fn history_size(&self) -> usize {
    self.history.len()
  }

  fn revert_to_history_size(&mut self, target_size: usize) {
    assert!(self.history_size() >= target_size);
    while self.history_size() > target_size {
      self.revert_last();
    }
  }

  fn update(&mut self, bitstring: &Bitstring) {
    for bit in bitstring.bits() {
      self.update_bit(*bit);
    }
  }

  fn predict(&mut self, bits: &Bitstring) -> f64 {
    // If we don't have enough history, assume a uniform distribution.
    if !self.has_context() {
      let neg_len = -(bits.len() as i64);
      return (neg_len as f64).exp2();
    }

    // Add all the symbols and store the block probability before and after.
    let initial_size = self.history_size();
    let initial_log_prob = self.log_block_prob();

    self.update(bits);
    let final_log_prob = self.log_block_prob();
    self.revert_to_history_size(initial_size);

    // Let "b" be bits and "h" be our history. We use the equation:
    //      P[hb] = P[h] * P[b | h]
    // <==> log(P[hb]) = log(P[h]) + log(P[b | h])
    // <==> log(P[b | h]) = log(P[hb]) - log(P[h])
    let log_conditional_prob = final_log_prob - initial_log_prob;
    return log_conditional_prob.exp2();
  }
}

impl ContextTree {
  /// The depth of the tree is the distance between leaves and the root.
  /// For a given depth d, the tree will compute a mixture model of all
  /// Markov chains of degree at most d.
  pub fn create(depth: usize) -> Self {
    ContextTree {
      root: Node::create_root(depth, 0),
      history: Bitstring::new(),
      depth: depth,
    }
  }

  /// Returns the total number of nodes in the tree.
  pub fn size(&self) -> usize {
    self.root.size()
  }

  /// Adds a single bit to the history seen by this tree, updating all the
  /// involved nodes if necessary.
  fn update_bit(&mut self, bit: Bit) {
    self.for_node_in_context(&|node: &mut Node| {
      node.update(bit);
    });
    self.history.push(bit);
  }

  /// Removes a single bit from the history seen by this tree, updating all the
  /// involved nodes if necessary.
  fn revert_last(&mut self) {
    let bit = self.history.pop();
    self.for_node_in_context(&|node: &mut Node| {
      node.revert(bit);
    });
  }

  /// Returns log2 of the estimated probability of the current history.
  fn log_block_prob(&mut self) -> f64 {
    return self.root.log_weighted_prob();
  }

  /// Executes the provided closure on the current context, i.e., follows the
  /// current history down to a leaf and executes the closure on every node
  /// from the leaf to the root. Does nothing if the current history is not
  /// long enough to reach a leaf.
  fn for_node_in_context<F: Fn(&mut Node)>(&mut self, closure: &F) {
    if !self.has_context() {
      return;
    }
    self.root.execute(closure, &self.history, 0);
  }

  /// Inidicates whether this tree "has context", i.e., whether or not the
  /// history is long enough to be able to follow a path down to a leaf.
  fn has_context(&self) -> bool { self.history_size() >= self.depth }
}

/// One node in the context tree. Maintains a count of the number
/// of zeroes and ones seen and provides access to partial results
/// of the overall probability estimates computed in the context
/// tree.
struct Node {
  /// Children.
  zero_child: Option<Box<Node>>,
  one_child: Option<Box<Node>>,

  /// Frequencies.
  zeroes: u64,
  ones: u64,

  /// The log2 of the probability of a sequence with self.zeroes zeroes
  /// and self.ones one. Computed using the Krichevsky Trofimov
  /// estimator algorithm.
  log_kt_prob: f64,

  /// The log2 of the weighted block probability of the current sequence.
  /// Funtion of self.zeroes and self.ones and the values recursively
  /// provided by the children of this node (if any), lazily recomputed.
  log_weighted_prob: f64,
}

impl Node {
  /// Creates a node and all its children recursively until the provided
  /// depth is reached.
  pub fn create_root(target_depth: usize, current_depth: usize) -> Self {
    if target_depth == current_depth {
      return Node::new(None, None);  // Leaf.
    } else {
      return Node::new(
          Some(Box::new(Node::create_root(target_depth, current_depth + 1))),
          Some(Box::new(Node::create_root(target_depth, current_depth + 1))));
    }
  }

  fn new(zero_child: Option<Box<Node>>, one_child: Option<Box<Node>>) -> Self {
    // Context trees are complete trees (either both children or no children).
    assert!(zero_child.is_none() == one_child.is_none());

    return Node {
      zero_child: zero_child,
      one_child: one_child,
      zeroes: 0,
      ones: 0,
      log_kt_prob: 0.0,
      log_weighted_prob: 0.0,
    };
  }

  pub fn log_weighted_prob(&self) -> f64 {
    self.log_weighted_prob
  }

  /// Executes the given closure on the appropriate child (if any) and then on
  /// itself.
  pub fn execute<F: Fn(&mut Node)>(
      &mut self,
      closure: &F,
      history: &Bitstring,
      depth: usize) {
    if !self.is_leaf() {
      let bit = history.bit(history.len() - depth - 1);
      self.mut_child(bit).execute(closure, history, depth + 1);
    }
    closure(self);
  }

  /// Updates the state of the node based on the value of the new bit.
  /// Assumes that child nodes (if any) have already been updated.
  pub fn update(&mut self, bit: Bit) {
    *self.mut_frequency(bit) += 1;

    // Old value + 0.5 is the same as new value - 0.5.
    let first_summand = (*self.mut_frequency(bit) as f64 - 0.5).log2();
    let norm_summand = ((self.zeroes + self.ones) as f64).log2();

    self.log_kt_prob = first_summand - norm_summand + self.log_kt_prob;
    self.update_weighted_prob();
  }

  /// Corresponds to undoing one update for the specified bit value.
  /// Assumes that child nodes (if any) have already been reverted.
  pub fn revert(&mut self, bit: Bit) {
    *self.mut_frequency(bit) -= 1;

    let first_summand = (*self.mut_frequency(bit) as f64 + 0.5).log2();
    let norm_summand = ((self.zeroes + self.ones + 1) as f64).log2();

    self.log_kt_prob = self.log_kt_prob - first_summand + norm_summand;
    self.update_weighted_prob();
  }

  /// Returns the child corresponding to the supplied bit. Must not be called
  /// on leaf nodes.
  fn mut_child(&mut self, bit: Bit) -> &mut Node {
    assert!(!self.is_leaf());
    match bit {
      Bit::Zero => self.zero_child.as_mut().unwrap(),
      Bit::One => self.one_child.as_mut().unwrap(),
    }
  }

  /// Returns the frequency corresponding to the supplied bit.
  fn mut_frequency(&mut self, bit: Bit) -> &mut u64 {
    match bit {
      Bit::Zero => &mut self.zeroes,
      Bit::One => &mut self.ones,
    }
  }

  /// Updates the value of log_weighted_prob. Assumes that log_kt_prob
  /// as well as all values in all children (if any) are up-to-date.
  fn update_weighted_prob(&mut self) {
    if self.is_leaf() {
      self.log_weighted_prob = self.log_kt_prob;
      return;
    }

    let log_weighted_probs =
        self.zero_child.as_ref().unwrap().log_weighted_prob() +
        self.one_child.as_ref().unwrap().log_weighted_prob();
    let exponent = log_weighted_probs - self.log_kt_prob;
    let power = exponent.exp2();

    // The recursive term should be log2(1 + 2^exponent). If the exponent is
    // large enough, 2^exponent is inf. In this case, we approximate
    // log2(1 + 2^exponent) by log2(2^exponent) = exponent.
    let recursive_term =
        if power == f64::INFINITY { exponent } else { (power + 1.0).log2() };

    self.log_weighted_prob = self.log_kt_prob + recursive_term - 1.0;
  }

  /// Returns the number of nodes in the subtree rooted at this node.
  fn size(&self) -> usize {
    1 + match self.zero_child {
      None => 0,
      Some(ref node) => node.size(),
    } + match self.one_child {
      None => 0,
      Some(ref node) => node.size(),
    }
  }

  fn is_leaf(&self) -> bool {
    // By construction we have either both children or none.
    return self.zero_child.is_none();
  }
}
<commit_msg>Fix a comment.<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2015 dinowernli
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

use std::f64;

use bitstring::Bit;
use bitstring::Bitstring;

// Open questions:
// - How to declare some fields final (apparently rejected by rust community).
// - How to declare some fields mutable (can be changed for non-mut self).
// - Have the tree implement iter and simplify some methods with "fold".

/// An object capable of predicting observations and rewards based on
/// experience. Predictors have an abstract notion of history which
/// grows over time and represents the experience.
pub trait Predictor {
  /// Returns the size of the currently tracked history.
  fn history_size(&self) -> usize;

  /// Appends the provided bit string to the tracked history.
  fn update(&mut self, bitstring: &Bitstring);

  /// Reverts the context tree to a previous state by undoing update
  /// operations. The specified size must be at most the current size.
  fn revert_to_history_size(&mut self, target_size: usize);

  /// Returns the probability, given the current history, that "bits" are the
  /// next observed symbols.
  fn predict(&mut self, bits: &Bitstring) -> f64;
}

/// Context tree which computes a probability estimate for a sequence of
/// bits based on a mixture model of Markov chains. Lower-degree Markov
/// chains have higher priors in the mixture model, thus formalizing the
/// notion of Occam's Razor.
///
/// In order to get a probability for a sequence, all bits in the sequence
/// must be processed in-order.
pub struct ContextTree {
  root: Node,
  history: Bitstring,
  depth: usize,
}

impl Predictor for ContextTree {
  fn history_size(&self) -> usize {
    self.history.len()
  }

  fn revert_to_history_size(&mut self, target_size: usize) {
    assert!(self.history_size() >= target_size);
    while self.history_size() > target_size {
      self.revert_last();
    }
  }

  fn update(&mut self, bitstring: &Bitstring) {
    for bit in bitstring.bits() {
      self.update_bit(*bit);
    }
  }

  fn predict(&mut self, bits: &Bitstring) -> f64 {
    // If we don't have enough history, assume a uniform distribution.
    if !self.has_context() {
      let neg_len = -(bits.len() as i64);
      return (neg_len as f64).exp2();
    }

    // Add all the symbols and store the block probability before and after.
    let initial_size = self.history_size();
    let initial_log_prob = self.log_block_prob();

    self.update(bits);
    let final_log_prob = self.log_block_prob();
    self.revert_to_history_size(initial_size);

    // Let "b" be bits and "h" be our history. We use the equation:
    //      P[hb] = P[h] * P[b | h]
    // <==> log(P[hb]) = log(P[h]) + log(P[b | h])
    // <==> log(P[b | h]) = log(P[hb]) - log(P[h])
    let log_conditional_prob = final_log_prob - initial_log_prob;
    return log_conditional_prob.exp2();
  }
}

impl ContextTree {
  /// The depth of the tree is the distance between leaves and the root.
  /// For a given depth d, the tree will compute a mixture model of all
  /// Markov chains of degree at most d.
  pub fn create(depth: usize) -> Self {
    ContextTree {
      root: Node::create_root(depth, 0),
      history: Bitstring::new(),
      depth: depth,
    }
  }

  /// Returns the total number of nodes in the tree.
  pub fn size(&self) -> usize {
    self.root.size()
  }

  /// Adds a single bit to the history seen by this tree, updating all the
  /// involved nodes if necessary.
  fn update_bit(&mut self, bit: Bit) {
    self.for_node_in_context(&|node: &mut Node| {
      node.update(bit);
    });
    self.history.push(bit);
  }

  /// Removes a single bit from the history seen by this tree, updating all the
  /// involved nodes if necessary.
  fn revert_last(&mut self) {
    let bit = self.history.pop();
    self.for_node_in_context(&|node: &mut Node| {
      node.revert(bit);
    });
  }

  /// Returns log2 of the estimated probability of the current history.
  fn log_block_prob(&mut self) -> f64 {
    return self.root.log_weighted_prob();
  }

  /// Executes the provided closure on the current context, i.e., follows the
  /// current history down to a leaf and executes the closure on every node
  /// from the leaf to the root. Does nothing if the current history is not
  /// long enough to reach a leaf.
  fn for_node_in_context<F: Fn(&mut Node)>(&mut self, closure: &F) {
    if !self.has_context() {
      return;
    }
    self.root.execute(closure, &self.history, 0);
  }

  /// Inidicates whether this tree "has context", i.e., whether or not the
  /// history is long enough to be able to follow a path down to a leaf.
  fn has_context(&self) -> bool { self.history_size() >= self.depth }
}

/// One node in the context tree. Maintains a count of the number
/// of zeroes and ones seen and provides access to partial results
/// of the overall probability estimates computed in the context
/// tree.
struct Node {
  /// Children.
  zero_child: Option<Box<Node>>,
  one_child: Option<Box<Node>>,

  /// Frequencies.
  zeroes: u64,
  ones: u64,

  /// The log2 of the probability of a sequence with self.zeroes zeroes
  /// and self.ones one. Computed using the Krichevsky Trofimov
  /// estimator algorithm.
  log_kt_prob: f64,

  /// The log2 of the weighted block probability of the current sequence.
  /// Funtion of self.zeroes and self.ones and the values recursively
  /// provided by the children of this node (if any), lazily recomputed.
  log_weighted_prob: f64,
}

impl Node {
  /// Creates a node and all its children recursively until the provided
  /// depth is reached.
  pub fn create_root(target_depth: usize, current_depth: usize) -> Self {
    if target_depth == current_depth {
      return Node::new(None, None);  // Leaf.
    } else {
      return Node::new(
          Some(Box::new(Node::create_root(target_depth, current_depth + 1))),
          Some(Box::new(Node::create_root(target_depth, current_depth + 1))));
    }
  }

  fn new(zero_child: Option<Box<Node>>, one_child: Option<Box<Node>>) -> Self {
    // Context trees are complete trees (either both children or no children).
    assert!(zero_child.is_none() == one_child.is_none());

    return Node {
      zero_child: zero_child,
      one_child: one_child,
      zeroes: 0,
      ones: 0,
      log_kt_prob: 0.0,
      log_weighted_prob: 0.0,
    };
  }

  pub fn log_weighted_prob(&self) -> f64 {
    self.log_weighted_prob
  }

  /// Executes the given closure on the appropriate child (if any) and then on
  /// itself.
  pub fn execute<F: Fn(&mut Node)>(
      &mut self,
      closure: &F,
      history: &Bitstring,
      depth: usize) {
    if !self.is_leaf() {
      let bit = history.bit(history.len() - depth - 1);
      self.mut_child(bit).execute(closure, history, depth + 1);
    }
    closure(self);
  }

  /// Updates the state of the node based on the value of the new bit.
  /// Assumes that child nodes (if any) have already been updated.
  pub fn update(&mut self, bit: Bit) {
    *self.mut_frequency(bit) += 1;

    // Old value + 0.5 is the same as new value - 0.5.
    let first_summand = (*self.mut_frequency(bit) as f64 - 0.5).log2();
    let norm_summand = ((self.zeroes + self.ones) as f64).log2();

    self.log_kt_prob = first_summand - norm_summand + self.log_kt_prob;
    self.update_weighted_prob();
  }

  /// Corresponds to undoing one update for the specified bit value.
  /// Assumes that child nodes (if any) have already been reverted.
  pub fn revert(&mut self, bit: Bit) {
    *self.mut_frequency(bit) -= 1;

    let first_summand = (*self.mut_frequency(bit) as f64 + 0.5).log2();
    let norm_summand = ((self.zeroes + self.ones + 1) as f64).log2();

    self.log_kt_prob = self.log_kt_prob - first_summand + norm_summand;
    self.update_weighted_prob();
  }

  /// Returns the child corresponding to the supplied bit. Must not be called
  /// on leaf nodes.
  fn mut_child(&mut self, bit: Bit) -> &mut Node {
    assert!(!self.is_leaf());
    match bit {
      Bit::Zero => self.zero_child.as_mut().unwrap(),
      Bit::One => self.one_child.as_mut().unwrap(),
    }
  }

  /// Returns the frequency corresponding to the supplied bit.
  fn mut_frequency(&mut self, bit: Bit) -> &mut u64 {
    match bit {
      Bit::Zero => &mut self.zeroes,
      Bit::One => &mut self.ones,
    }
  }

  /// Updates the value of log_weighted_prob. Assumes that log_kt_prob
  /// as well as all values in all children (if any) are up-to-date.
  fn update_weighted_prob(&mut self) {
    if self.is_leaf() {
      self.log_weighted_prob = self.log_kt_prob;
      return;
    }

    let log_weighted_probs =
        self.zero_child.as_ref().unwrap().log_weighted_prob() +
        self.one_child.as_ref().unwrap().log_weighted_prob();
    let exponent = log_weighted_probs - self.log_kt_prob;
    let power = exponent.exp2();

    // The recursive term should be log2(1 + 2^exponent). If the exponent is
    // large enough, 2^exponent is inf. In this case, we approximate
    // log2(1 + 2^exponent) by log2(2^exponent) = exponent.
    let recursive_term =
        if power == f64::INFINITY { exponent } else { (power + 1.0).log2() };

    self.log_weighted_prob = self.log_kt_prob + recursive_term - 1.0;
  }

  /// Returns the number of nodes in the subtree rooted at this node.
  fn size(&self) -> usize {
    1 + match self.zero_child {
      None => 0,
      Some(ref node) => node.size(),
    } + match self.one_child {
      None => 0,
      Some(ref node) => node.size(),
    }
  }

  fn is_leaf(&self) -> bool {
    // By construction we have either both children or none.
    return self.zero_child.is_none();
  }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

use crate::dom::bindings::cell::DomRefCell;
use crate::dom::bindings::codegen::Bindings::RTCDataChannelBinding::RTCDataChannelInit;
use crate::dom::bindings::codegen::Bindings::RTCDataChannelBinding::RTCDataChannelMethods;
use crate::dom::bindings::codegen::Bindings::RTCErrorBinding::{RTCErrorDetailType, RTCErrorInit};
use crate::dom::bindings::inheritance::Castable;
use crate::dom::bindings::refcounted::Trusted;
use crate::dom::bindings::reflector::{reflect_dom_object, DomObject};
use crate::dom::bindings::root::DomRoot;
use crate::dom::bindings::str::{DOMString, USVString};
use crate::dom::blob::Blob;
use crate::dom::event::{Event, EventBubbles, EventCancelable};
use crate::dom::eventtarget::EventTarget;
use crate::dom::globalscope::GlobalScope;
use crate::dom::messageevent::MessageEvent;
use crate::dom::rtcerror::RTCError;
use crate::dom::rtcerrorevent::RTCErrorEvent;
use crate::task_source::TaskSource;
use dom_struct::dom_struct;
use js::conversions::ToJSValConvertible;
use js::jsapi::JSAutoRealm;
use js::jsval::UndefinedValue;
use js::rust::CustomAutoRooterGuard;
use js::typedarray::{ArrayBuffer, ArrayBufferView};
use servo_media::webrtc::{
    WebRtcController, WebRtcDataChannelBackend, WebRtcDataChannelInit, WebRtcError,
};
use std::sync::mpsc;

#[dom_struct]
pub struct RTCDataChannel {
    eventtarget: EventTarget,
    #[ignore_malloc_size_of = "defined in servo-media"]
    channel: Box<dyn WebRtcDataChannelBackend>,
    label: USVString,
    ordered: bool,
    max_packet_life_time: Option<u16>,
    max_retransmits: Option<u16>,
    protocol: USVString,
    negotiated: bool,
    id: Option<u16>,
}

impl RTCDataChannel {
    #[allow(unrooted_must_root)]
    pub fn new_inherited(
        webrtc_controller: &DomRefCell<Option<WebRtcController>>,
        label: USVString,
        options: &RTCDataChannelInit,
    ) -> RTCDataChannel {
        let webrtc = webrtc_controller.borrow();
        let webrtc = webrtc.as_ref().unwrap();

        let (sender, receiver) = mpsc::channel::<Box<dyn WebRtcDataChannelBackend>>();

        let mut init: WebRtcDataChannelInit = options.into();
        init.label = label.to_string();

        webrtc.create_data_channel(init, sender);
        let channel = receiver.recv().unwrap();

        let rtc_data_channel = RTCDataChannel {
            eventtarget: EventTarget::new_inherited(),
            channel,
            label,
            ordered: options.ordered,
            max_packet_life_time: options.maxPacketLifeTime,
            max_retransmits: options.maxRetransmits,
            protocol: options.protocol.clone(),
            negotiated: options.negotiated,
            id: options.id,
        };

        let trusted = Trusted::new(&rtc_data_channel);

        let this = trusted.clone();
        rtc_data_channel.channel.set_on_open(Box::new(move || {
            let this_ = this.clone();
            let global = this.root().global();
            let task_source = global.networking_task_source();
            let _ = task_source.queue(
                task!(on_open: move || {
                    this_.root().on_open();
                }),
                global.upcast(),
            );
        }));

        let this = trusted.clone();
        rtc_data_channel.channel.set_on_close(Box::new(move || {
            let this_ = this.clone();
            let global = this.root().global();
            let task_source = global.networking_task_source();
            let _ = task_source.queue(
                task!(on_close: move || {
                    this_.root().on_close();
                }),
                global.upcast(),
            );
        }));

        let this = trusted.clone();
        rtc_data_channel
            .channel
            .set_on_error(Box::new(move |error| {
                let this_ = this.clone();
                let global = this.root().global();
                let task_source = global.networking_task_source();
                let _ = task_source.queue(
                    task!(on_error: move || {
                        this_.root().on_error(error);
                    }),
                    global.upcast(),
                );
            }));

        let this = trusted.clone();
        rtc_data_channel
            .channel
            .set_on_message(Box::new(move |message| {
                let this_ = this.clone();
                let global = this.root().global();
                let task_source = global.networking_task_source();
                let _ = task_source.queue(
                    task!(on_message: move || {
                        this_.root().on_message(message);
                    }),
                    global.upcast(),
                );
            }));

        rtc_data_channel
    }

    pub fn new(
        global: &GlobalScope,
        webrtc_controller: &DomRefCell<Option<WebRtcController>>,
        label: USVString,
        options: &RTCDataChannelInit,
    ) -> DomRoot<RTCDataChannel> {
        reflect_dom_object(
            Box::new(RTCDataChannel::new_inherited(
                webrtc_controller,
                label,
                options,
            )),
            global,
        )
    }

    fn on_open(&self) {
        let event = Event::new(
            &self.global(),
            atom!("open"),
            EventBubbles::DoesNotBubble,
            EventCancelable::NotCancelable,
        );
        event.upcast::<Event>().fire(self.upcast());
    }

    fn on_close(&self) {
        let event = Event::new(
            &self.global(),
            atom!("close"),
            EventBubbles::DoesNotBubble,
            EventCancelable::NotCancelable,
        );
        event.upcast::<Event>().fire(self.upcast());
    }

    fn on_error(&self, error: WebRtcError) {
        let init = RTCErrorInit {
            errorDetail: RTCErrorDetailType::Data_channel_failure,
            httpRequestStatusCode: None,
            receivedAlert: None,
            sctpCauseCode: None,
            sdpLineNumber: None,
            sentAlert: None,
        };
        let message = match error {
            WebRtcError::Backend(message) => DOMString::from(message),
        };
        let error = RTCError::new(&self.global(), &init, message);
        let event = RTCErrorEvent::new(&self.global(), atom!("error"), false, false, &error);
        event.upcast::<Event>().fire(self.upcast());
    }

    #[allow(unsafe_code)]
    fn on_message(&self, text: String) {
        // XXX(ferjm) Support binary messages
        unsafe {
            let global = self.global();
            let cx = global.get_cx();
            let _ac = JSAutoRealm::new(*cx, self.reflector().get_jsobject().get());
            rooted!(in(*cx) let mut message = UndefinedValue());
            text.to_jsval(*cx, message.handle_mut());

            MessageEvent::dispatch_jsval(
                self.upcast(),
                &global,
                message.handle(),
                Some(&global.origin().immutable().ascii_serialization()),
                None,
                vec![],
            );
        }
    }
}

impl RTCDataChannelMethods for RTCDataChannel {
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onopen
    event_handler!(open, GetOnopen, SetOnopen);
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onbufferedamountlow
    event_handler!(
        bufferedamountlow,
        GetOnbufferedamountlow,
        SetOnbufferedamountlow
    );
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onerror
    event_handler!(error, GetOnerror, SetOnerror);
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onclosing
    event_handler!(closing, GetOnclosing, SetOnclosing);
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onclose
    event_handler!(close, GetOnclose, SetOnclose);
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onmessage
    event_handler!(message, GetOnmessage, SetOnmessage);

    // https://www.w3.org/TR/webrtc/#dom-datachannel-label
    fn Label(&self) -> USVString {
        self.label.clone()
    }
    // https://www.w3.org/TR/webrtc/#dom-datachannel-ordered
    fn Ordered(&self) -> bool {
        self.ordered
    }

    // https://www.w3.org/TR/webrtc/#dom-datachannel-maxpacketlifetime
    fn GetMaxPacketLifeTime(&self) -> Option<u16> {
        self.max_packet_life_time
    }

    // https://www.w3.org/TR/webrtc/#dom-datachannel-maxretransmits
    fn GetMaxRetransmits(&self) -> Option<u16> {
        self.max_retransmits
    }

    // https://www.w3.org/TR/webrtc/#dom-datachannel-protocol
    fn Protocol(&self) -> USVString {
        self.protocol.clone()
    }

    // https://www.w3.org/TR/webrtc/#dom-datachannel-negotiated
    fn Negotiated(&self) -> bool {
        self.negotiated
    }

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-id
    fn GetId(&self) -> Option<u16> {
        self.id
    }

    //    fn ReadyState(&self) -> RTCDataChannelState;
    //    fn BufferedAmount(&self) -> u32;
    //    fn BufferedAmountLowThreshold(&self) -> u32;
    //    fn SetBufferedAmountLowThreshold(&self, value: u32) -> ();

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-close
    fn Close(&self) -> () {}

    //    fn BinaryType(&self) -> DOMString;
    //    fn SetBinaryType(&self, value: DOMString) -> ();

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-send
    fn Send(&self, data: USVString) -> () {}

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-send!overload-1
    fn Send_(&self, data: &Blob) -> () {}

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-send!overload-2
    fn Send__(&self, data: CustomAutoRooterGuard<ArrayBuffer>) -> () {}

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-send!overload-3
    fn Send___(&self, data: CustomAutoRooterGuard<ArrayBufferView>) -> () {}
}

impl From<&RTCDataChannelInit> for WebRtcDataChannelInit {
    fn from(init: &RTCDataChannelInit) -> WebRtcDataChannelInit {
        WebRtcDataChannelInit {
            label: String::new(),
            id: init.id,
            max_packet_life_time: init.maxPacketLifeTime,
            max_retransmits: init.maxRetransmits,
            negotiated: init.negotiated,
            ordered: init.ordered,
            protocol: init.protocol.to_string(),
        }
    }
}
<commit_msg>Allow sending strings through data channels<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

use crate::dom::bindings::cell::DomRefCell;
use crate::dom::bindings::codegen::Bindings::RTCDataChannelBinding::RTCDataChannelInit;
use crate::dom::bindings::codegen::Bindings::RTCDataChannelBinding::RTCDataChannelMethods;
use crate::dom::bindings::codegen::Bindings::RTCErrorBinding::{RTCErrorDetailType, RTCErrorInit};
use crate::dom::bindings::inheritance::Castable;
use crate::dom::bindings::refcounted::Trusted;
use crate::dom::bindings::reflector::{reflect_dom_object, DomObject};
use crate::dom::bindings::root::DomRoot;
use crate::dom::bindings::str::{DOMString, USVString};
use crate::dom::blob::Blob;
use crate::dom::event::{Event, EventBubbles, EventCancelable};
use crate::dom::eventtarget::EventTarget;
use crate::dom::globalscope::GlobalScope;
use crate::dom::messageevent::MessageEvent;
use crate::dom::rtcerror::RTCError;
use crate::dom::rtcerrorevent::RTCErrorEvent;
use crate::task_source::TaskSource;
use dom_struct::dom_struct;
use js::conversions::ToJSValConvertible;
use js::jsapi::JSAutoRealm;
use js::jsval::UndefinedValue;
use js::rust::CustomAutoRooterGuard;
use js::typedarray::{ArrayBuffer, ArrayBufferView};
use servo_media::webrtc::{
    WebRtcController, WebRtcDataChannelBackend, WebRtcDataChannelInit, WebRtcError,
};
use std::sync::mpsc;

#[dom_struct]
pub struct RTCDataChannel {
    eventtarget: EventTarget,
    #[ignore_malloc_size_of = "defined in servo-media"]
    channel: Box<dyn WebRtcDataChannelBackend>,
    label: USVString,
    ordered: bool,
    max_packet_life_time: Option<u16>,
    max_retransmits: Option<u16>,
    protocol: USVString,
    negotiated: bool,
    id: Option<u16>,
}

impl RTCDataChannel {
    #[allow(unrooted_must_root)]
    pub fn new_inherited(
        webrtc_controller: &DomRefCell<Option<WebRtcController>>,
        label: USVString,
        options: &RTCDataChannelInit,
    ) -> RTCDataChannel {
        let webrtc = webrtc_controller.borrow();
        let webrtc = webrtc.as_ref().unwrap();

        let (sender, receiver) = mpsc::channel::<Box<dyn WebRtcDataChannelBackend>>();

        let mut init: WebRtcDataChannelInit = options.into();
        init.label = label.to_string();

        webrtc.create_data_channel(init, sender);
        let channel = receiver.recv().unwrap();

        let rtc_data_channel = RTCDataChannel {
            eventtarget: EventTarget::new_inherited(),
            channel,
            label,
            ordered: options.ordered,
            max_packet_life_time: options.maxPacketLifeTime,
            max_retransmits: options.maxRetransmits,
            protocol: options.protocol.clone(),
            negotiated: options.negotiated,
            id: options.id,
        };

        let trusted = Trusted::new(&rtc_data_channel);

        let this = trusted.clone();
        rtc_data_channel.channel.set_on_open(Box::new(move || {
            let this_ = this.clone();
            let global = this.root().global();
            let task_source = global.networking_task_source();
            let _ = task_source.queue(
                task!(on_open: move || {
                    this_.root().on_open();
                }),
                global.upcast(),
            );
        }));

        let this = trusted.clone();
        rtc_data_channel.channel.set_on_close(Box::new(move || {
            let this_ = this.clone();
            let global = this.root().global();
            let task_source = global.networking_task_source();
            let _ = task_source.queue(
                task!(on_close: move || {
                    this_.root().on_close();
                }),
                global.upcast(),
            );
        }));

        let this = trusted.clone();
        rtc_data_channel
            .channel
            .set_on_error(Box::new(move |error| {
                let this_ = this.clone();
                let global = this.root().global();
                let task_source = global.networking_task_source();
                let _ = task_source.queue(
                    task!(on_error: move || {
                        this_.root().on_error(error);
                    }),
                    global.upcast(),
                );
            }));

        let this = trusted.clone();
        rtc_data_channel
            .channel
            .set_on_message(Box::new(move |message| {
                let this_ = this.clone();
                let global = this.root().global();
                let task_source = global.networking_task_source();
                let _ = task_source.queue(
                    task!(on_message: move || {
                        this_.root().on_message(message);
                    }),
                    global.upcast(),
                );
            }));

        rtc_data_channel
    }

    pub fn new(
        global: &GlobalScope,
        webrtc_controller: &DomRefCell<Option<WebRtcController>>,
        label: USVString,
        options: &RTCDataChannelInit,
    ) -> DomRoot<RTCDataChannel> {
        reflect_dom_object(
            Box::new(RTCDataChannel::new_inherited(
                webrtc_controller,
                label,
                options,
            )),
            global,
        )
    }

    fn on_open(&self) {
        let event = Event::new(
            &self.global(),
            atom!("open"),
            EventBubbles::DoesNotBubble,
            EventCancelable::NotCancelable,
        );
        event.upcast::<Event>().fire(self.upcast());
    }

    fn on_close(&self) {
        let event = Event::new(
            &self.global(),
            atom!("close"),
            EventBubbles::DoesNotBubble,
            EventCancelable::NotCancelable,
        );
        event.upcast::<Event>().fire(self.upcast());
    }

    fn on_error(&self, error: WebRtcError) {
        let init = RTCErrorInit {
            errorDetail: RTCErrorDetailType::Data_channel_failure,
            httpRequestStatusCode: None,
            receivedAlert: None,
            sctpCauseCode: None,
            sdpLineNumber: None,
            sentAlert: None,
        };
        let message = match error {
            WebRtcError::Backend(message) => DOMString::from(message),
        };
        let error = RTCError::new(&self.global(), &init, message);
        let event = RTCErrorEvent::new(&self.global(), atom!("error"), false, false, &error);
        event.upcast::<Event>().fire(self.upcast());
    }

    #[allow(unsafe_code)]
    fn on_message(&self, text: String) {
        // XXX(ferjm) Support binary messages
        unsafe {
            let global = self.global();
            let cx = global.get_cx();
            let _ac = JSAutoRealm::new(*cx, self.reflector().get_jsobject().get());
            rooted!(in(*cx) let mut message = UndefinedValue());
            text.to_jsval(*cx, message.handle_mut());

            MessageEvent::dispatch_jsval(
                self.upcast(),
                &global,
                message.handle(),
                Some(&global.origin().immutable().ascii_serialization()),
                None,
                vec![],
            );
        }
    }
}

impl RTCDataChannelMethods for RTCDataChannel {
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onopen
    event_handler!(open, GetOnopen, SetOnopen);
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onbufferedamountlow
    event_handler!(
        bufferedamountlow,
        GetOnbufferedamountlow,
        SetOnbufferedamountlow
    );
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onerror
    event_handler!(error, GetOnerror, SetOnerror);
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onclosing
    event_handler!(closing, GetOnclosing, SetOnclosing);
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onclose
    event_handler!(close, GetOnclose, SetOnclose);
    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-onmessage
    event_handler!(message, GetOnmessage, SetOnmessage);

    // https://www.w3.org/TR/webrtc/#dom-datachannel-label
    fn Label(&self) -> USVString {
        self.label.clone()
    }
    // https://www.w3.org/TR/webrtc/#dom-datachannel-ordered
    fn Ordered(&self) -> bool {
        self.ordered
    }

    // https://www.w3.org/TR/webrtc/#dom-datachannel-maxpacketlifetime
    fn GetMaxPacketLifeTime(&self) -> Option<u16> {
        self.max_packet_life_time
    }

    // https://www.w3.org/TR/webrtc/#dom-datachannel-maxretransmits
    fn GetMaxRetransmits(&self) -> Option<u16> {
        self.max_retransmits
    }

    // https://www.w3.org/TR/webrtc/#dom-datachannel-protocol
    fn Protocol(&self) -> USVString {
        self.protocol.clone()
    }

    // https://www.w3.org/TR/webrtc/#dom-datachannel-negotiated
    fn Negotiated(&self) -> bool {
        self.negotiated
    }

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-id
    fn GetId(&self) -> Option<u16> {
        self.id
    }

    //    fn ReadyState(&self) -> RTCDataChannelState;
    //    fn BufferedAmount(&self) -> u32;
    //    fn BufferedAmountLowThreshold(&self) -> u32;
    //    fn SetBufferedAmountLowThreshold(&self, value: u32) -> ();

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-close
    fn Close(&self) -> () {}

    //    fn BinaryType(&self) -> DOMString;
    //    fn SetBinaryType(&self, value: DOMString) -> ();

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-send
    fn Send(&self, data: USVString) -> () {
        if let Err(error) = self.channel.send(&data.0) {
            warn!("Could not send data channel message. Error: {:?}", error);
        }
    }

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-send!overload-1
    fn Send_(&self, data: &Blob) -> () {}

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-send!overload-2
    fn Send__(&self, data: CustomAutoRooterGuard<ArrayBuffer>) -> () {}

    // https://www.w3.org/TR/webrtc/#dom-rtcdatachannel-send!overload-3
    fn Send___(&self, data: CustomAutoRooterGuard<ArrayBufferView>) -> () {}
}

impl From<&RTCDataChannelInit> for WebRtcDataChannelInit {
    fn from(init: &RTCDataChannelInit) -> WebRtcDataChannelInit {
        WebRtcDataChannelInit {
            label: String::new(),
            id: init.id,
            max_packet_life_time: init.maxPacketLifeTime,
            max_retransmits: init.maxRetransmits,
            negotiated: init.negotiated,
            ordered: init.ordered,
            protocol: init.protocol.to_string(),
        }
    }
}
<|endoftext|>


<commit_before>use l2::ast::*;
use l2::codegen::*;
use l2::heapview::*;
use l2::stack::*;

use std::collections::HashMap;
use std::vec::Vec;

#[derive(Clone, Copy)]
enum MachineMode {
    Read,
    Write
}

struct MachineState {
    h: usize,
    s: usize,
    p: CodePtr,
    cp: CodePtr,
    fail: bool,
    heap: Heap,
    mode: MachineMode,
    stack: Stack,
    registers: Registers
}

type CodeDir = HashMap<(Atom, usize), usize>;

pub struct Machine {
    ms: MachineState,
    code: Code,
    code_dir: CodeDir
}

impl Machine {
    pub fn new() -> Self {
        Machine {
            ms: MachineState::new(),
            code: Vec::new(),
            code_dir: HashMap::new()
        }
    }

    fn failed(&self) -> bool {
        self.ms.fail
    }

    pub fn add_fact(&mut self, fact: &Term, mut code: Code) {
        let p = self.code.len();
        let name  = fact.name().clone();
        let arity = fact.arity();

        self.code.append(&mut code);
        self.code_dir.insert((name, arity), p);
    }

    pub fn add_rule(&mut self, rule: &Rule, mut code: Code) {
        let p = self.code.len();
        let name  = rule.head.0.name().clone();
        let arity = rule.head.0.arity();

        self.code.append(&mut code);
        self.code_dir.insert((name, arity), p);
    }

    fn execute_instr(&mut self, instr: &Line) -> bool {
        let mut instr = instr;

        loop {
            match instr {
                &Line::Fact(ref fact) => {
                    for fact_instr in fact {
                        self.ms.execute_fact_instr(&fact_instr);
                    }
                    self.ms.p += 1;
                },
                &Line::Query(ref query) => {
                    for query_instr in query {
                        self.ms.execute_query_instr(&query_instr);
                    }
                    self.ms.p += 1;
                },
                &Line::Control(ref control_instr) =>
                    self.ms.execute_ctrl_instr(&self.code_dir, control_instr),
            }

            if self.failed() {
                return false;
            }

            match self.ms.p {
                CodePtr::DirEntry(p) if p < self.code.len() =>
                    instr = &self.code[p],
                _ => break
            }
        }

        true
    }

    fn heap_view(&self, var_dir: HashMap<&Var, HeapCellRef>) -> String {
        let mut result = String::new();

        for (var, hcr) in var_dir {
            let mut arities = Vec::new();
            let viewer = HeapCellViewer::new(&self.ms.heap, hcr.heap_offset());

            if result != "" {
                result += "\n";
            }

            result += var.as_str();
            result += " = ";

            for view in viewer {
                match arities.pop() {
                    Some(n) => arities.push(n-1),
                    None => {}
                }

                if !(arities.is_empty() || result.ends_with("(")) {
                    result += ", ";
                }

                match view {
                    HeapCellView::Str(arity, ref name) => {
                        result += name.as_str();

                        if arity > 0 {
                            arities.push(arity);
                            result += "(";
                        }
                    },
                    HeapCellView::Var(cell_num) => {
                        result += "_";
                        result += cell_num.to_string().as_str();
                    }
                }

                while let Some(&0) = arities.last() {
                    result += ")";
                    arities.pop();
                }
            }
        }

        result
    }

    pub fn run_query(&mut self, code: Code, cg: &CodeGenerator) -> Option<String>
    {
        let mut succeeded = true;

        for instr in code.iter().take(1) {
            succeeded = self.execute_instr(&instr);
        }

        if succeeded {
            let mut heap_locs = HashMap::new();

            for (var, vr) in cg.vars() {
                let hcr = self.ms.registers[vr.root_register()];
                heap_locs.insert(*var, hcr);
            }

            for instr in code.iter().skip(1) {
                succeeded = self.execute_instr(&instr);
                if !succeeded {
                    break;
                }
            }
        }

        if succeeded {
            let result = Some(self.heap_view(heap_locs));
            self.ms.reset();
            result
        } else {
            self.ms.reset();
            None
        }
    }
}

impl MachineState {
    fn new() -> MachineState {
        MachineState { h: 0,
                       s: 0,
                       p: CodePtr::TopLevel,
                       cp: CodePtr::TopLevel,
                       fail: false,
                       heap: Vec::with_capacity(256),
                       mode: MachineMode::Write,
                       stack: Stack::new(),
                       registers: vec![HeapCellRef::Ref(0); 32] }
    }

    fn register_mut(&mut self, r: RegType) -> &mut HeapCellRef {
        match r {
            RegType::Temp(r) => &mut self.registers[r],
            RegType::Perm(r) => &mut self.stack[r]
        }
    }

    fn lookup(&self, a: Addr) -> HeapCellRef {
        match a {
            Addr::HeapCell(r)  => self.heap[r].as_ref(r),
            Addr::RegNum(r)    => self.registers[r],
            Addr::StackCell(s) => self.stack[s]
        }
    }

    fn deref(&self, a: Addr) -> Addr {
        let mut a = a;

        loop {
            if let HeapCellRef::Ref(value) = self.lookup(a) {
                if let Addr::HeapCell(av) = a {
                    if value != av {
                        a = Addr::HeapCell(value);
                        continue;
                    }
                } else {
                    a = Addr::HeapCell(value);
                    continue;
                }
            }

            return a;
        };
    }

    fn is_unbound(hc: &HeapCellValue, index: usize) -> bool {
        match hc {
            &HeapCellValue::Ref(r) => r == index,
            _ => false
        }
    }

    fn bind(&mut self, a: Addr, val: usize) {
        let mut a = a;

        loop {
            match a {
                addr @ Addr::RegNum(_) | addr @ Addr::StackCell(_) => {
                    if let HeapCellRef::Ref(hc) = self.lookup(addr) {
                        a = Addr::HeapCell(hc);
                    } else if Self::is_unbound(&self.heap[val], val) {
                        self.heap[val] = HeapCellValue::from(self.lookup(addr));
                        break;
                    } else {
                        self.fail = true;
                        break;
                    }
                },
                Addr::HeapCell(hc) => {
                    if Self::is_unbound(&self.heap[hc], hc) {
                        self.heap[hc] = HeapCellValue::Ref(val);
                        break;
                    } else if Self::is_unbound(&self.heap[val], val) {
                        self.heap[val] = HeapCellValue::Ref(hc);
                        break;
                    } else {
                        self.fail = true;
                        break;
                    }
                }
            };
        }
    }

    fn unify(&mut self, a1: Addr, a2: Addr) {
        let mut pdl = vec![a1, a2];

        self.fail = false;

        while !(pdl.is_empty() || self.fail) {
            let d1 = self.deref(pdl.pop().unwrap());
            let d2 = self.deref(pdl.pop().unwrap());

            if d1 != d2 {
                match (self.lookup(d1), self.lookup(d2)) {
                    (HeapCellRef::Ref(hc), _) =>
                        self.bind(d2, hc),
                    (_, HeapCellRef::Ref(hc)) =>
                        self.bind(d1, hc),
                    (HeapCellRef::Str(a1), HeapCellRef::Str(a2)) => {
                        let r1 = &self.heap[a1];
                        let r2 = &self.heap[a2];

                        if let &HeapCellValue::NamedStr(n1, ref f1) = r1 {
                            if let &HeapCellValue::NamedStr(n2, ref f2) = r2 {
                                if n1 == n2 && *f1 == *f2 {
                                    for i in 1 .. n1 + 1 {
                                        pdl.push(Addr::HeapCell(a1 + i));
                                        pdl.push(Addr::HeapCell(a2 + i));
                                    }

                                    continue;
                                }
                            }
                        }

                        self.fail = true;
                    },
                };
            }
        }
    }

    fn execute_query_instr(&mut self, instr: &QueryInstruction) {
        match instr {
            &QueryInstruction::PutStructure(_, ref name, arity, reg) => {
                self.heap.push(HeapCellValue::Str(self.h + 1));
                self.heap.push(HeapCellValue::NamedStr(arity, name.clone()));

                *self.register_mut(reg) = HeapCellRef::Str(self.h + 1);

                self.h += 2;
            },
            &QueryInstruction::PutValue(norm, arg) =>
                self.registers[arg] = match norm {
                    RegType::Temp(reg) => self.registers[reg],
                    RegType::Perm(reg) => self.stack[reg]
                },
            &QueryInstruction::PutVariable(norm, arg) => {
                self.heap.push(HeapCellValue::Ref(self.h));

                *self.register_mut(norm) = HeapCellRef::Ref(self.h);
                self.registers[arg] = HeapCellRef::Ref(self.h);

                self.h += 1;
            },
            &QueryInstruction::SetVariable(reg) => {
                self.heap.push(HeapCellValue::Ref(self.h));
                *self.register_mut(reg) = HeapCellRef::Ref(self.h);

                self.h += 1;
            },
            &QueryInstruction::SetValue(reg) => {
                let heap_val = self.lookup(Addr::from(reg));
                self.heap.push(HeapCellValue::from(heap_val));

                self.h += 1;
            },
        }
    }

    fn execute_fact_instr(&mut self, instr: &FactInstruction) {
        match instr {
            &FactInstruction::GetStructure(_, ref name, arity, reg) => {
                let addr = self.deref(Addr::from(reg));

                match self.lookup(addr) {
                    HeapCellRef::Str(a) => {
                        let result = &self.heap[a];

                        if let &HeapCellValue::NamedStr(narity, ref str) = result {
                            if narity == arity && *name == *str {
                                self.s = a + 1;
                                self.mode = MachineMode::Read;
                            } else {
                                self.fail = true;
                            }
                        }
                    },
                    HeapCellRef::Ref(_) => {
                        self.heap.push(HeapCellValue::Str(self.h + 1));
                        self.heap.push(HeapCellValue::NamedStr(arity, name.clone()));

                        let h = self.h;

                        self.bind(addr, h);

                        self.h += 2;
                        self.mode = MachineMode::Write;
                    }
                };
            },
            &FactInstruction::GetVariable(norm, arg) =>
                *self.register_mut(norm) = self.registers[arg],
            &FactInstruction::GetValue(norm, arg) =>
                self.unify(Addr::from(norm), Addr::RegNum(arg)),
            &FactInstruction::UnifyVariable(reg) => {
                match self.mode {
                    MachineMode::Read  =>
                        *self.register_mut(reg) = self.heap[self.s].as_ref(self.s),
                    MachineMode::Write => {
                        self.heap.push(HeapCellValue::Ref(self.h));
                        *self.register_mut(reg) = HeapCellRef::Ref(self.h);
                        self.h += 1;
                    }
                };

                self.s += 1;
            },
            &FactInstruction::UnifyValue(reg) => {
                let s = self.s;

                match self.mode {
                    MachineMode::Read  =>
                        self.unify(Addr::from(reg), Addr::HeapCell(s)),
                    MachineMode::Write => {
                        let heap_val = self.lookup(Addr::from(reg));
                        self.heap.push(HeapCellValue::from(heap_val));
                        self.h += 1;
                    }
                };

                self.s += 1;
            }
        }
    }

    fn execute_ctrl_instr(&mut self, code_dir: &CodeDir, instr: &ControlInstruction)
    {
        match instr {
            &ControlInstruction::Allocate(num_cells) => {
                self.stack.push(self.cp, num_cells);
                self.p += 1;
            },
            &ControlInstruction::Call(ref name, arity) => {
                let compiled_tl_index = code_dir.get(&(name.clone(), arity))
                                                .map(|index| *index);

                match compiled_tl_index {
                    Some(compiled_tl_index) => {
                        self.cp = self.p + 1;
                        self.p  = CodePtr::DirEntry(compiled_tl_index);
                    },
                    None => self.fail = true
                };
            },
            &ControlInstruction::Deallocate => {
                self.p = self.stack.get_cp();
                self.stack.pop();
            },
            &ControlInstruction::Proceed =>
                self.p = self.cp,
        };
    }

    fn reset(&mut self) {
        self.h = 0;
        self.s = 0;
        self.p = CodePtr::TopLevel;
        self.cp = CodePtr::TopLevel;

        self.fail = false;
        self.heap.clear();
        self.mode = MachineMode::Write;
        self.stack = Stack::new();
        self.registers = vec![HeapCellRef::Ref(0); 32];
    }
}
<commit_msg>bug fix<commit_after>use l2::ast::*;
use l2::codegen::*;
use l2::heapview::*;
use l2::stack::*;

use std::collections::HashMap;
use std::vec::Vec;

#[derive(Clone, Copy)]
enum MachineMode {
    Read,
    Write
}

struct MachineState {
    h: usize,
    s: usize,
    p: CodePtr,
    cp: CodePtr,
    fail: bool,
    heap: Heap,
    mode: MachineMode,
    stack: Stack,
    registers: Registers
}

type CodeDir = HashMap<(Atom, usize), usize>;

pub struct Machine {
    ms: MachineState,
    code: Code,
    code_dir: CodeDir
}

impl Machine {
    pub fn new() -> Self {
        Machine {
            ms: MachineState::new(),
            code: Vec::new(),
            code_dir: HashMap::new()
        }
    }

    fn failed(&self) -> bool {
        self.ms.fail
    }

    pub fn add_fact(&mut self, fact: &Term, mut code: Code) {
        let p = self.code.len();
        let name  = fact.name().clone();
        let arity = fact.arity();

        self.code.append(&mut code);
        self.code_dir.insert((name, arity), p);
    }

    pub fn add_rule(&mut self, rule: &Rule, mut code: Code) {
        let p = self.code.len();
        let name  = rule.head.0.name().clone();
        let arity = rule.head.0.arity();

        self.code.append(&mut code);
        self.code_dir.insert((name, arity), p);
    }

    fn execute_instr(&mut self, instr: &Line) -> bool {
        let mut instr = instr;

        loop {
            match instr {
                &Line::Fact(ref fact) => {
                    for fact_instr in fact {
                        self.ms.execute_fact_instr(&fact_instr);
                    }
                    self.ms.p += 1;
                },
                &Line::Query(ref query) => {
                    for query_instr in query {
                        self.ms.execute_query_instr(&query_instr);
                    }
                    self.ms.p += 1;
                },
                &Line::Control(ref control_instr) =>
                    self.ms.execute_ctrl_instr(&self.code_dir, control_instr),
            }

            if self.failed() {
                return false;
            }

            match self.ms.p {
                CodePtr::DirEntry(p) if p < self.code.len() =>
                    instr = &self.code[p],
                _ => break
            }
        }

        true
    }

    fn heap_view(&self, var_dir: HashMap<&Var, HeapCellRef>) -> String {
        let mut result = String::new();

        for (var, hcr) in var_dir {
            let mut arities = Vec::new();
            let viewer = HeapCellViewer::new(&self.ms.heap, hcr.heap_offset());

            if result != "" {
                result += "\n";
            }

            result += var.as_str();
            result += " = ";

            for view in viewer {
                match arities.pop() {
                    Some(n) => arities.push(n-1),
                    None => {}
                }

                if !(arities.is_empty() || result.ends_with("(")) {
                    result += ", ";
                }

                match view {
                    HeapCellView::Str(arity, ref name) => {
                        result += name.as_str();

                        if arity > 0 {
                            arities.push(arity);
                            result += "(";
                        }
                    },
                    HeapCellView::Var(cell_num) => {
                        result += "_";
                        result += cell_num.to_string().as_str();
                    }
                }

                while let Some(&0) = arities.last() {
                    result += ")";
                    arities.pop();
                }
            }
        }

        result
    }

    pub fn run_query(&mut self, code: Code, cg: &CodeGenerator) -> Option<String>
    {
        let mut succeeded = true;

        for instr in code.iter().take(1) {
            succeeded = self.execute_instr(&instr);
        }

        let mut heap_locs = HashMap::new();

        if succeeded {
            for (var, vr) in cg.vars() {
                let hcr = self.ms.registers[vr.root_register()];
                heap_locs.insert(*var, hcr);
            }

            for instr in code.iter().skip(1) {
                succeeded = self.execute_instr(&instr);
                if !succeeded {
                    break;
                }
            }
        }

        if succeeded {
            let result = Some(self.heap_view(heap_locs));
            self.ms.reset();
            result
        } else {
            self.ms.reset();
            None
        }
    }
}

impl MachineState {
    fn new() -> MachineState {
        MachineState { h: 0,
                       s: 0,
                       p: CodePtr::TopLevel,
                       cp: CodePtr::TopLevel,
                       fail: false,
                       heap: Vec::with_capacity(256),
                       mode: MachineMode::Write,
                       stack: Stack::new(),
                       registers: vec![HeapCellRef::Ref(0); 32] }
    }

    fn register_mut(&mut self, r: RegType) -> &mut HeapCellRef {
        match r {
            RegType::Temp(r) => &mut self.registers[r],
            RegType::Perm(r) => &mut self.stack[r]
        }
    }

    fn lookup(&self, a: Addr) -> HeapCellRef {
        match a {
            Addr::HeapCell(r)  => self.heap[r].as_ref(r),
            Addr::RegNum(r)    => self.registers[r],
            Addr::StackCell(s) => self.stack[s]
        }
    }

    fn deref(&self, a: Addr) -> Addr {
        let mut a = a;

        loop {
            if let HeapCellRef::Ref(value) = self.lookup(a) {
                if let Addr::HeapCell(av) = a {
                    if value != av {
                        a = Addr::HeapCell(value);
                        continue;
                    }
                } else {
                    a = Addr::HeapCell(value);
                    continue;
                }
            }

            return a;
        };
    }

    fn is_unbound(hc: &HeapCellValue, index: usize) -> bool {
        match hc {
            &HeapCellValue::Ref(r) => r == index,
            _ => false
        }
    }

    fn bind(&mut self, a: Addr, val: usize) {
        let mut a = a;

        loop {
            match a {
                addr @ Addr::RegNum(_) | addr @ Addr::StackCell(_) => {
                    if let HeapCellRef::Ref(hc) = self.lookup(addr) {
                        a = Addr::HeapCell(hc);
                    } else if Self::is_unbound(&self.heap[val], val) {
                        self.heap[val] = HeapCellValue::from(self.lookup(addr));
                        break;
                    } else {
                        self.fail = true;
                        break;
                    }
                },
                Addr::HeapCell(hc) => {
                    if Self::is_unbound(&self.heap[hc], hc) {
                        self.heap[hc] = HeapCellValue::Ref(val);
                        break;
                    } else if Self::is_unbound(&self.heap[val], val) {
                        self.heap[val] = HeapCellValue::Ref(hc);
                        break;
                    } else {
                        self.fail = true;
                        break;
                    }
                }
            };
        }
    }

    fn unify(&mut self, a1: Addr, a2: Addr) {
        let mut pdl = vec![a1, a2];

        self.fail = false;

        while !(pdl.is_empty() || self.fail) {
            let d1 = self.deref(pdl.pop().unwrap());
            let d2 = self.deref(pdl.pop().unwrap());

            if d1 != d2 {
                match (self.lookup(d1), self.lookup(d2)) {
                    (HeapCellRef::Ref(hc), _) =>
                        self.bind(d2, hc),
                    (_, HeapCellRef::Ref(hc)) =>
                        self.bind(d1, hc),
                    (HeapCellRef::Str(a1), HeapCellRef::Str(a2)) => {
                        let r1 = &self.heap[a1];
                        let r2 = &self.heap[a2];

                        if let &HeapCellValue::NamedStr(n1, ref f1) = r1 {
                            if let &HeapCellValue::NamedStr(n2, ref f2) = r2 {
                                if n1 == n2 && *f1 == *f2 {
                                    for i in 1 .. n1 + 1 {
                                        pdl.push(Addr::HeapCell(a1 + i));
                                        pdl.push(Addr::HeapCell(a2 + i));
                                    }

                                    continue;
                                }
                            }
                        }

                        self.fail = true;
                    },
                };
            }
        }
    }

    fn execute_query_instr(&mut self, instr: &QueryInstruction) {
        match instr {
            &QueryInstruction::PutStructure(_, ref name, arity, reg) => {
                self.heap.push(HeapCellValue::Str(self.h + 1));
                self.heap.push(HeapCellValue::NamedStr(arity, name.clone()));

                *self.register_mut(reg) = HeapCellRef::Str(self.h + 1);

                self.h += 2;
            },
            &QueryInstruction::PutValue(norm, arg) =>
                self.registers[arg] = match norm {
                    RegType::Temp(reg) => self.registers[reg],
                    RegType::Perm(reg) => self.stack[reg]
                },
            &QueryInstruction::PutVariable(norm, arg) => {
                self.heap.push(HeapCellValue::Ref(self.h));

                *self.register_mut(norm) = HeapCellRef::Ref(self.h);
                self.registers[arg] = HeapCellRef::Ref(self.h);

                self.h += 1;
            },
            &QueryInstruction::SetVariable(reg) => {
                self.heap.push(HeapCellValue::Ref(self.h));
                *self.register_mut(reg) = HeapCellRef::Ref(self.h);

                self.h += 1;
            },
            &QueryInstruction::SetValue(reg) => {
                let heap_val = self.lookup(Addr::from(reg));
                self.heap.push(HeapCellValue::from(heap_val));

                self.h += 1;
            },
        }
    }

    fn execute_fact_instr(&mut self, instr: &FactInstruction) {
        match instr {
            &FactInstruction::GetStructure(_, ref name, arity, reg) => {
                let addr = self.deref(Addr::from(reg));

                match self.lookup(addr) {
                    HeapCellRef::Str(a) => {
                        let result = &self.heap[a];

                        if let &HeapCellValue::NamedStr(narity, ref str) = result {
                            if narity == arity && *name == *str {
                                self.s = a + 1;
                                self.mode = MachineMode::Read;
                            } else {
                                self.fail = true;
                            }
                        }
                    },
                    HeapCellRef::Ref(_) => {
                        self.heap.push(HeapCellValue::Str(self.h + 1));
                        self.heap.push(HeapCellValue::NamedStr(arity, name.clone()));

                        let h = self.h;

                        self.bind(addr, h);

                        self.h += 2;
                        self.mode = MachineMode::Write;
                    }
                };
            },
            &FactInstruction::GetVariable(norm, arg) =>
                *self.register_mut(norm) = self.registers[arg],
            &FactInstruction::GetValue(norm, arg) =>
                self.unify(Addr::from(norm), Addr::RegNum(arg)),
            &FactInstruction::UnifyVariable(reg) => {
                match self.mode {
                    MachineMode::Read  =>
                        *self.register_mut(reg) = self.heap[self.s].as_ref(self.s),
                    MachineMode::Write => {
                        self.heap.push(HeapCellValue::Ref(self.h));
                        *self.register_mut(reg) = HeapCellRef::Ref(self.h);
                        self.h += 1;
                    }
                };

                self.s += 1;
            },
            &FactInstruction::UnifyValue(reg) => {
                let s = self.s;

                match self.mode {
                    MachineMode::Read  =>
                        self.unify(Addr::from(reg), Addr::HeapCell(s)),
                    MachineMode::Write => {
                        let heap_val = self.lookup(Addr::from(reg));
                        self.heap.push(HeapCellValue::from(heap_val));
                        self.h += 1;
                    }
                };

                self.s += 1;
            }
        }
    }

    fn execute_ctrl_instr(&mut self, code_dir: &CodeDir, instr: &ControlInstruction)
    {
        match instr {
            &ControlInstruction::Allocate(num_cells) => {
                self.stack.push(self.cp, num_cells);
                self.p += 1;
            },
            &ControlInstruction::Call(ref name, arity) => {
                let compiled_tl_index = code_dir.get(&(name.clone(), arity))
                                                .map(|index| *index);

                match compiled_tl_index {
                    Some(compiled_tl_index) => {
                        self.cp = self.p + 1;
                        self.p  = CodePtr::DirEntry(compiled_tl_index);
                    },
                    None => self.fail = true
                };
            },
            &ControlInstruction::Deallocate => {
                self.p = self.stack.get_cp();
                self.stack.pop();
            },
            &ControlInstruction::Proceed =>
                self.p = self.cp,
        };
    }

    fn reset(&mut self) {
        self.h = 0;
        self.s = 0;
        self.p = CodePtr::TopLevel;
        self.cp = CodePtr::TopLevel;

        self.fail = false;
        self.heap.clear();
        self.mode = MachineMode::Write;
        self.stack = Stack::new();
        self.registers = vec![HeapCellRef::Ref(0); 32];
    }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Manages the SQL database that holds the list of registered domain names.
// Each records is made of the name, the private token and the Let's Encrypt
// challenge value.

use std::sync::mpsc::{channel, Receiver};
use std::thread;

use r2d2_sqlite::SqliteConnectionManager;
use r2d2;
use rusqlite::Row;

#[derive(Clone, Debug, PartialEq)]
pub struct DomainRecord {
    name: String,
    token: String,
    dns_challenge: Option<String>,
}

impl DomainRecord {
    fn from_sql(row: Row) -> Self {
        let challenge = row.get::<i32, String>(2);
        let dns_challenge = if challenge.is_empty() {
            None
        } else {
            Some(challenge)
        };
        DomainRecord {
            name: row.get(0),
            token: row.get(1),
            dns_challenge: dns_challenge,
        }
    }

    fn new(name: &str, token: &str, challenge: Option<&str>) -> Self {
        let dns_challenge = match challenge {
            Some(val) => Some(val.to_owned()),
            None => None,
        };

        DomainRecord {
            name: name.to_owned(),
            token: token.to_owned(),
            dns_challenge: dns_challenge,
        }
    }
}

unsafe impl Send for DomainRecord {}
unsafe impl Sync for DomainRecord {}

#[derive(Debug)]
pub enum DomainError {
    DbUnavailable,
    SQLError,
    NoRecord,
}

pub struct DomainDb {
    pool: r2d2::Pool<SqliteConnectionManager>,
}

// try! like macro that sends back the error of a SQLite operation
// over the mpsc channel if needed.
macro_rules! sqltry {
    ($sql:expr, $tx:ident, $err:expr) => (
        match $sql {
            Err(_) => {
                $tx.send(Err($err)).unwrap();
                return;
            }
            Ok(value) => value,
        }
    )
}

impl DomainDb {
    pub fn new(path: &str) -> Self {
        let config = r2d2::Config::default();
        let manager = SqliteConnectionManager::new(path);
        let pool = r2d2::Pool::new(config, manager).expect("Unable to open database!");

        // Create the database table if needed.
        let conn = pool.get().unwrap();
        conn.execute("CREATE TABLE IF NOT EXISTS domains (
                      name          TEXT NOT NULL PRIMARY KEY,
                      token         TEXT NOT NULL,
                      dns_challenge TEXT NOT NULL)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains database: {}", err);
                            });

        conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS domains_token ON domains(token)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains.token index: {}", err);
                            });

        DomainDb { pool: pool }
    }

    fn select_record(&self,
                     request: &str,
                     value: &str)
                     -> Receiver<Result<DomainRecord, DomainError>> {
        let (tx, rx) = channel();

        // Run the sql command on a pooled thread.
        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            let mut stmt = sqltry!(conn.prepare(&request), tx, DomainError::SQLError);
            let mut rows = sqltry!(stmt.query(&[&value]), tx, DomainError::SQLError);
            if let Some(result_row) = rows.next() {
                let row = sqltry!(result_row, tx, DomainError::SQLError);
                tx.send(Ok(DomainRecord::from_sql(row))).unwrap();
            } else {
                tx.send(Err(DomainError::NoRecord)).unwrap();
            }
        });

        rx
    }

    pub fn get_record_by_name(&self, name: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT name, token, dns_challenge FROM domains WHERE name=$1",
                           name)
    }

    pub fn get_record_by_token(&self, token: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT name, token, dns_challenge FROM domains WHERE token=$1",
                           token)
    }

    pub fn add_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            sqltry!(conn.execute("INSERT INTO domains VALUES ($1, $2, $3)",
                                 &[&record.name,
                                   &record.token,
                                   &record.dns_challenge.unwrap_or("".to_owned())]),
                    tx,
                    DomainError::SQLError);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    pub fn update_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);

            sqltry!(conn.execute("UPDATE domains SET dns_challenge=$1 WHERE name=$2 AND token=$3",
                                 &[&record.dns_challenge.unwrap_or("".to_owned()),
                                   &record.name,
                                   &record.token]),
                    tx,
                    DomainError::SQLError);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    pub fn delete_record_by_name(&self, name: &str) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let name = name.to_owned();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          sqltry!(conn.execute("DELETE FROM domains WHERE name=$1", &[&name]),
                                  tx,
                                  DomainError::SQLError);
                          tx.send(Ok(())).unwrap();
                      });

        rx
    }

    pub fn delete_record_by_token(&self, token: &str) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let token = token.to_owned();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          sqltry!(conn.execute("DELETE FROM domains WHERE token=$1", &[&token]),
                                  tx,
                                  DomainError::SQLError);
                          tx.send(Ok(())).unwrap();
                      });
        rx
    }

    pub fn flush(&self) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          sqltry!(conn.execute("DELETE FROM domains", &[]),
                                  tx,
                                  DomainError::SQLError);
                          tx.send(Ok(())).unwrap();
                      });
        rx
    }
}

#[test]
fn test_domain_store() {
    let db = DomainDb::new("domain_db_test.sqlite");

    // Start with an empty db.
    db.flush().recv().unwrap().expect("Flushing the db");

    // Check that we don't find any record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find a record by name in an empty db."),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find a record by token in an empty db."),
    }

    // Add a record without a dns challenge.
    let no_challenge_record = DomainRecord::new("test.example.org", "test-token", None);
    db.add_record(no_challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Adding the no_challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Update the record to have challenge.
    let challenge_record =
        DomainRecord::new("test.example.org", "test-token", Some("dns-challenge"));
    db.update_record(challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Updating the challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Remove by name.
    db.delete_record_by_name(&challenge_record.name)
        .recv()
        .unwrap()
        .expect("Should delete record by name");
    match db.get_record_by_name(&challenge_record.name)
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find this record anymore."),
    }

    // Add again, and remove by token.
    db.add_record(challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Adding the challenge record");
    db.delete_record_by_token(&challenge_record.token)
        .recv()
        .unwrap()
        .expect("Should delete record by token");
    match db.get_record_by_name(&challenge_record.name)
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find this record anymore."),
    }
}
<commit_msg>Refactor delete_XXX methods<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Manages the SQL database that holds the list of registered domain names.
// Each records is made of the name, the private token and the Let's Encrypt
// challenge value.

use std::sync::mpsc::{channel, Receiver};
use std::thread;

use r2d2_sqlite::SqliteConnectionManager;
use r2d2;
use rusqlite::Row;

#[derive(Clone, Debug, PartialEq)]
pub struct DomainRecord {
    name: String,
    token: String,
    dns_challenge: Option<String>,
}

impl DomainRecord {
    fn from_sql(row: Row) -> Self {
        let challenge = row.get::<i32, String>(2);
        let dns_challenge = if challenge.is_empty() {
            None
        } else {
            Some(challenge)
        };
        DomainRecord {
            name: row.get(0),
            token: row.get(1),
            dns_challenge: dns_challenge,
        }
    }

    fn new(name: &str, token: &str, challenge: Option<&str>) -> Self {
        let dns_challenge = match challenge {
            Some(val) => Some(val.to_owned()),
            None => None,
        };

        DomainRecord {
            name: name.to_owned(),
            token: token.to_owned(),
            dns_challenge: dns_challenge,
        }
    }
}

unsafe impl Send for DomainRecord {}
unsafe impl Sync for DomainRecord {}

#[derive(Debug)]
pub enum DomainError {
    DbUnavailable,
    SQLError,
    NoRecord,
}

pub struct DomainDb {
    pool: r2d2::Pool<SqliteConnectionManager>,
}

// try! like macro that sends back the error of a SQLite operation
// over the mpsc channel if needed.
macro_rules! sqltry {
    ($sql:expr, $tx:ident, $err:expr) => (
        match $sql {
            Err(_) => {
                $tx.send(Err($err)).unwrap();
                return;
            }
            Ok(value) => value,
        }
    )
}

impl DomainDb {
    pub fn new(path: &str) -> Self {
        let config = r2d2::Config::default();
        let manager = SqliteConnectionManager::new(path);
        let pool = r2d2::Pool::new(config, manager).expect("Unable to open database!");

        // Create the database table if needed.
        let conn = pool.get().unwrap();
        conn.execute("CREATE TABLE IF NOT EXISTS domains (
                      name          TEXT NOT NULL PRIMARY KEY,
                      token         TEXT NOT NULL,
                      dns_challenge TEXT NOT NULL)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains database: {}", err);
                            });

        conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS domains_token ON domains(token)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains.token index: {}", err);
                            });

        DomainDb { pool: pool }
    }

    fn select_record(&self,
                     request: &str,
                     value: &str)
                     -> Receiver<Result<DomainRecord, DomainError>> {
        let (tx, rx) = channel();

        // Run the sql command on a pooled thread.
        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            let mut stmt = sqltry!(conn.prepare(&request), tx, DomainError::SQLError);
            let mut rows = sqltry!(stmt.query(&[&value]), tx, DomainError::SQLError);
            if let Some(result_row) = rows.next() {
                let row = sqltry!(result_row, tx, DomainError::SQLError);
                tx.send(Ok(DomainRecord::from_sql(row))).unwrap();
            } else {
                tx.send(Err(DomainError::NoRecord)).unwrap();
            }
        });

        rx
    }

    pub fn get_record_by_name(&self, name: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT name, token, dns_challenge FROM domains WHERE name=$1",
                           name)
    }

    pub fn get_record_by_token(&self, token: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT name, token, dns_challenge FROM domains WHERE token=$1",
                           token)
    }

    pub fn add_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            sqltry!(conn.execute("INSERT INTO domains VALUES ($1, $2, $3)",
                                 &[&record.name,
                                   &record.token,
                                   &record.dns_challenge.unwrap_or("".to_owned())]),
                    tx,
                    DomainError::SQLError);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    pub fn update_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);

            sqltry!(conn.execute("UPDATE domains SET dns_challenge=$1 WHERE name=$2 AND token=$3",
                                 &[&record.dns_challenge.unwrap_or("".to_owned()),
                                   &record.name,
                                   &record.token]),
                    tx,
                    DomainError::SQLError);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    fn execute_1param_sql(&self, request: &str, value: &str) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          sqltry!(conn.execute(&request, &[&value]), tx, DomainError::SQLError);
                          tx.send(Ok(())).unwrap();
                      });

        rx
    }

    pub fn delete_record_by_name(&self, name: &str) -> Receiver<Result<(), DomainError>> {
        self.execute_1param_sql("DELETE FROM domains WHERE name=$1", name)
    }

    pub fn delete_record_by_token(&self, token: &str) -> Receiver<Result<(), DomainError>> {
        self.execute_1param_sql("DELETE FROM domains WHERE token=$1", token)
    }

    pub fn flush(&self) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          sqltry!(conn.execute("DELETE FROM domains", &[]),
                                  tx,
                                  DomainError::SQLError);
                          tx.send(Ok(())).unwrap();
                      });
        rx
    }
}

#[test]
fn test_domain_store() {
    let db = DomainDb::new("domain_db_test.sqlite");

    // Start with an empty db.
    db.flush().recv().unwrap().expect("Flushing the db");

    // Check that we don't find any record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find a record by name in an empty db."),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find a record by token in an empty db."),
    }

    // Add a record without a dns challenge.
    let no_challenge_record = DomainRecord::new("test.example.org", "test-token", None);
    db.add_record(no_challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Adding the no_challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Update the record to have challenge.
    let challenge_record =
        DomainRecord::new("test.example.org", "test-token", Some("dns-challenge"));
    db.update_record(challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Updating the challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Remove by name.
    db.delete_record_by_name(&challenge_record.name)
        .recv()
        .unwrap()
        .expect("Should delete record by name");
    match db.get_record_by_name(&challenge_record.name)
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find this record anymore."),
    }

    // Add again, and remove by token.
    db.add_record(challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Adding the challenge record");
    db.delete_record_by_token(&challenge_record.token)
        .recv()
        .unwrap()
        .expect("Should delete record by token");
    match db.get_record_by_name(&challenge_record.name)
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find this record anymore."),
    }
}
<|endoftext|>