<commit_before>//! GC heap allocator.

use std::cell::RefCell;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::{fmt, mem, ptr};
use bit_vec::BitVec;

// What does this do? You'll never guess!
type HeapId<'a> = PhantomData<::std::cell::Cell<&'a mut ()>>;

/// Total number of Pair objects that can be allocated in a Heap at once.
///
/// XXX BOGUS: At present, this is number is carefully selected (with insider
/// knowledge of both `size_of::<PairStorage>()` on my platform and the `Box`
/// allocator's behavior) so that a HeapStorage object fits in a page and thus
/// the assertions below about HEAP_STORAGE_ALIGN pass.
pub const HEAP_SIZE: usize = 125;

/// We rely on all bits to the right of this bit being 0 in addresses of
/// HeapStorage instances.
const HEAP_STORAGE_ALIGN: usize = 0x1000;

// The heap is (for now) just a big array of Pairs
struct HeapStorage<'a> {
    mark_bits: BitVec,
    objects: [PairStorage<'a>; HEAP_SIZE]
}

impl<'a> HeapStorage<'a> {
    unsafe fn new() -> HeapStorage<'a> {
        HeapStorage {
            mark_bits: BitVec::from_elem(HEAP_SIZE, false),
            objects: mem::uninitialized()
        }
    }
}

pub struct Heap<'a> {
    id: HeapId<'a>,
    storage: Box<HeapStorage<'a>>,
    freelist: *mut PairStorage<'a>,
    pins: RefCell<HashMap<*mut PairStorage<'a>, usize>>
}

/// Create a heap, pass it to a callback, then destroy the heap.
///
/// The heap's lifetime is directly tied to this function call, for safety. (So
/// the API is a little wonky --- but how many heaps were you planning on
/// creating?)
pub fn with_heap<F, O>(f: F) -> O
    where F: for<'a> FnOnce(&mut Heap<'a>) -> O
{
    f(&mut Heap::new())
}

/// Trait implemented by all types that can be stored directly in the GC heap:
/// the `Storage` types associated with any `HeapInline` or `HeapRef` type.
///
/// XXX maybe this should not be its own trait - fold into HeapInline?
///
pub unsafe trait Mark<'a> {
    unsafe fn mark(ptr: *mut Self);
}

/// Trait implemented by all types that can be stored in fields of structs (or,
/// eventually, elements of GCVecs) that are stored in the GC heap.
pub unsafe trait HeapInline<'a> {
    /// The type of the value when it is physically stored in the heap.
    type Storage;

    /// Extract the value from the heap. Do not under any circumstances call
    /// this.  It is for macro-generated code to call; it is impossible for
    /// ordinary users to call this safely, because `ptr` must be a direct,
    /// unwrapped reference to a value stored in the GC heap, which ordinary
    /// users cannot obtain.
    ///
    /// This turns any raw pointers in the `Storage` value into safe
    /// references, so while it's a dangerous function, the result of a correct
    /// call can be safely handed out to user code.
    ///
    unsafe fn from_heap(heap: &Heap<'a>, ptr: &Self::Storage) -> Self;

    /// Convert the value to the form it should have in the heap.
    /// This is for macro-generated code to call.
    fn to_heap(self) -> Self::Storage;
}

impl<'a> Heap<'a> {
    fn new() -> Heap<'a> {
        let mut h = Heap {
            id: PhantomData,
            storage: Box::new(unsafe { HeapStorage::new() }),
            freelist: ptr::null_mut(),
            pins: RefCell::new(HashMap::new())
        };
        assert!(&mut *h.storage as *mut HeapStorage<'a> as usize & (HEAP_STORAGE_ALIGN - 1) == 0);
        assert!(mem::size_of::<HeapStorage<'a>>() <= HEAP_STORAGE_ALIGN);
        for i in 0 .. HEAP_SIZE {
            let p = &mut h.storage.objects[i] as *mut PairStorage<'a>;
            unsafe {
                h.add_to_free_list(p);
            }
        }
        h
    }

    fn pin(&self, p: *mut PairStorage<'a>) {
        let mut pins = self.pins.borrow_mut();
        let entry = pins.entry(p).or_insert(0);
        *entry += 1;
    }

    fn unpin(&self, p: *mut PairStorage<'a>) {
        let mut pins = self.pins.borrow_mut();
        if {
            let entry = pins.entry(p).or_insert(0);
            assert!(*entry != 0);
            *entry -= 1;
            *entry == 0
        } {
            pins.remove(&p);
        }
    }

    unsafe fn add_to_free_list(&mut self, p: *mut PairStorage<'a>) {
        let listp = p as *mut *mut PairStorage<'a>;
        *listp = self.freelist;
        assert_eq!(*listp, self.freelist);
        self.freelist = p;
    }

    pub fn try_alloc(&mut self, pair: (Value<'a>, Value<'a>)) -> Option<Pair<'a>> {
        if self.freelist.is_null() {
            unsafe {
                self.gc();
            }
            if self.freelist.is_null() {
                return None;
            }
        }

        let p = self.freelist;
        unsafe {
            let listp = p as *mut *mut PairStorage<'a>;
            self.freelist = *listp;

            let (h, t) = pair;
            *p = PairStorage {
                head: h.to_heap(),
                tail: t.to_heap()
            };
            Some(Pair(PinnedRef::new(self, p)))
        }
    }

    unsafe fn find_mark_bit<T>(_ptr: *const T) -> (*mut BitVec, usize) {
        assert_eq!(mem::size_of::<T>(), mem::size_of::<PairStorage<'a>>());
        let storage_addr = _ptr as usize & !(HEAP_STORAGE_ALIGN - 1);
        let storage = storage_addr as *mut HeapStorage<'a>;
        let objects_addr = &mut (*storage).objects[0] as *mut PairStorage<'a> as usize;
        let index = (_ptr as usize - objects_addr) / mem::size_of::<PairStorage<'a>>();
        (&mut (*storage).mark_bits as *mut BitVec, index)
    }

    pub unsafe fn get_mark_bit<T>(_ptr: *const T) -> bool {
        let (bits, index) = Heap::find_mark_bit(_ptr);
        (*bits)[index]
    }

    pub unsafe fn set_mark_bit<T>(_ptr: *const T) {
        let (bits, index) = Heap::find_mark_bit(_ptr);
        (*bits).set(index, true);
    }

    pub fn alloc(&mut self, pair: (Value<'a>, Value<'a>)) -> Pair<'a> {
        self.try_alloc(pair).expect("out of memory (gc did not collect anything)")
    }

    pub fn alloc_null(&mut self) -> Pair<'a> {
        self.alloc((Value::Null, Value::Null))
    }

    unsafe fn gc(&mut self) {
        // mark phase
        self.storage.mark_bits.clear();
        for (&p, _) in self.pins.borrow().iter() {
            Mark::mark(p);
        }

        // sweep phase
        self.freelist = ptr::null_mut();
        for i in 0 .. HEAP_SIZE {
            let p = &mut self.storage.objects[i] as *mut PairStorage<'a>;
            if !self.storage.mark_bits[i] {
                self.add_to_free_list(p);
            }
        }
    }

    #[cfg(test)]
    pub fn force_gc(&mut self) {
        unsafe { self.gc(); }
    }
}

pub struct PinnedRef<'a, T> {
    ptr: *mut T,
    heap: *const Heap<'a>,
    heap_id: HeapId<'a>
}

impl<'a, T> PinnedRef<'a, T> {
    unsafe fn new(heap: &Heap<'a>, p: *mut T) -> PinnedRef<'a, T> {
        heap.pin(p as *mut PairStorage<'a>);  // XXX BOGUS
        PinnedRef {
            ptr: p,
            heap: heap as *const Heap<'a>,
            heap_id: heap.id
        }
    }
}

impl<'a, T> Drop for PinnedRef<'a, T> {
    fn drop(&mut self) {
        unsafe {
            (*self.heap).unpin(self.ptr as *mut PairStorage<'a>);  // XXX BOGUS
        }
    }
}

impl<'a, T> Clone for PinnedRef<'a, T> {
    fn clone(&self) -> PinnedRef<'a, T> {
        let &PinnedRef { ptr, heap, heap_id } = self;
        unsafe {
            (*heap).pin(ptr as *mut PairStorage<'a>);  // XXX BOGUS
        }
        PinnedRef {
            ptr: ptr,
            heap: heap,
            heap_id: heap_id
        }
    }
}

impl<'a, T> fmt::Debug for PinnedRef<'a, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "PinnedRef {{ ptr: {:p} }}", self.ptr)
    }
}

impl<'a, T> PartialEq for PinnedRef<'a, T> {
    fn eq(&self, other: &PinnedRef<'a, T>) -> bool {
        self.ptr == other.ptr
    }
}

impl<'a, T> Eq for PinnedRef<'a, T> {}

pub trait GCRef {
    #[cfg(test)]
    fn address(&self) -> usize;
}

// === Pair, the reference type

gc_ref_type! {
    pub struct Pair / PairStorage<'a> {
        head / set_head: Value<'a>,
        tail / set_tail: Value<'a>
    }
}

// === Values (a heap-inline enum)

use std::rc::Rc;

// Values inside the heap (GC heap-to-heap cross-references) (private)
enum ValueStorage<'a> {
    Null,
    Int(i32),
    Str(Rc<String>),
    Pair(*mut PairStorage<'a>)
}

unsafe impl<'a> Mark<'a> for ValueStorage<'a> {
    unsafe fn mark(ptr: *mut ValueStorage<'a>) {
        match *ptr {
            ValueStorage::Pair(p) => Mark::mark(p),
            _ => {}
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Value<'a> {
    Null,
    Int(i32),
    Str(Rc<String>),  // <-- equality is by value
    Pair(Pair<'a>)  // <-- equality is by pointer
}

unsafe impl<'a> HeapInline<'a> for Value<'a> {
    type Storage = ValueStorage<'a>;

    fn to_heap(self) -> ValueStorage<'a> {
        match self {
            Value::Null => ValueStorage::Null,
            Value::Int(n) => ValueStorage::Int(n),
            Value::Str(rcstr) => ValueStorage::Str(rcstr),
            Value::Pair(pair) => ValueStorage::Pair(HeapInline::<'a>::to_heap(pair))
        }
    }

    unsafe fn from_heap(heap: &Heap<'a>, v: &ValueStorage<'a>) -> Value<'a> {
        match v {
            &ValueStorage::Null => Value::Null,
            &ValueStorage::Int(n) => Value::Int(n),
            &ValueStorage::Str(ref rcstr) => Value::Str(rcstr.clone()),
            &ValueStorage::Pair(ref ptr) => Value::Pair(HeapInline::<'a>::from_heap(heap, ptr))
        }
    }
}

unsafe impl<'a, T: Copy + 'static> Mark<'a> for T {
    unsafe fn mark(_ptr: *mut T) {}
}

unsafe impl<'a, T: Copy + 'static> HeapInline<'a> for T {
    type Storage = Self;

    fn to_heap(self) -> T { self }

    unsafe fn from_heap(_heap: &Heap<'a>, v: &T) -> T { *v }
}
<commit_msg>Vaguely comment the alignment assertions.<commit_after>//! GC heap allocator.

use std::cell::RefCell;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::{fmt, mem, ptr};
use bit_vec::BitVec;

// What does this do? You'll never guess!
type HeapId<'a> = PhantomData<::std::cell::Cell<&'a mut ()>>;

/// Total number of Pair objects that can be allocated in a Heap at once.
///
/// XXX BOGUS: At present, this is number is carefully selected (with insider
/// knowledge of both `size_of::<PairStorage>()` on my platform and the `Box`
/// allocator's behavior) so that a HeapStorage object fits in a page and thus
/// the assertions below about HEAP_STORAGE_ALIGN pass.
pub const HEAP_SIZE: usize = 125;

/// We rely on all bits to the right of this bit being 0 in addresses of
/// HeapStorage instances.
const HEAP_STORAGE_ALIGN: usize = 0x1000;

// The heap is (for now) just a big array of Pairs
struct HeapStorage<'a> {
    mark_bits: BitVec,
    objects: [PairStorage<'a>; HEAP_SIZE]
}

impl<'a> HeapStorage<'a> {
    unsafe fn new() -> HeapStorage<'a> {
        HeapStorage {
            mark_bits: BitVec::from_elem(HEAP_SIZE, false),
            objects: mem::uninitialized()
        }
    }
}

pub struct Heap<'a> {
    id: HeapId<'a>,
    storage: Box<HeapStorage<'a>>,
    freelist: *mut PairStorage<'a>,
    pins: RefCell<HashMap<*mut PairStorage<'a>, usize>>
}

/// Create a heap, pass it to a callback, then destroy the heap.
///
/// The heap's lifetime is directly tied to this function call, for safety. (So
/// the API is a little wonky --- but how many heaps were you planning on
/// creating?)
pub fn with_heap<F, O>(f: F) -> O
    where F: for<'a> FnOnce(&mut Heap<'a>) -> O
{
    f(&mut Heap::new())
}

/// Trait implemented by all types that can be stored directly in the GC heap:
/// the `Storage` types associated with any `HeapInline` or `HeapRef` type.
///
/// XXX maybe this should not be its own trait - fold into HeapInline?
///
pub unsafe trait Mark<'a> {
    unsafe fn mark(ptr: *mut Self);
}

/// Trait implemented by all types that can be stored in fields of structs (or,
/// eventually, elements of GCVecs) that are stored in the GC heap.
pub unsafe trait HeapInline<'a> {
    /// The type of the value when it is physically stored in the heap.
    type Storage;

    /// Extract the value from the heap. Do not under any circumstances call
    /// this.  It is for macro-generated code to call; it is impossible for
    /// ordinary users to call this safely, because `ptr` must be a direct,
    /// unwrapped reference to a value stored in the GC heap, which ordinary
    /// users cannot obtain.
    ///
    /// This turns any raw pointers in the `Storage` value into safe
    /// references, so while it's a dangerous function, the result of a correct
    /// call can be safely handed out to user code.
    ///
    unsafe fn from_heap(heap: &Heap<'a>, ptr: &Self::Storage) -> Self;

    /// Convert the value to the form it should have in the heap.
    /// This is for macro-generated code to call.
    fn to_heap(self) -> Self::Storage;
}

impl<'a> Heap<'a> {
    fn new() -> Heap<'a> {
        let mut h = Heap {
            id: PhantomData,
            storage: Box::new(unsafe { HeapStorage::new() }),
            freelist: ptr::null_mut(),
            pins: RefCell::new(HashMap::new())
        };

        // These assertions will likely fail on 32-bit platforms or if someone
        // is somehow using a custom Box allocator. If either one fails, this
        // GC will not work.
        assert_eq!(&mut *h.storage as *mut HeapStorage<'a> as usize & (HEAP_STORAGE_ALIGN - 1), 0);
        assert!(mem::size_of::<HeapStorage<'a>>() <= HEAP_STORAGE_ALIGN);

        for i in 0 .. HEAP_SIZE {
            let p = &mut h.storage.objects[i] as *mut PairStorage<'a>;
            unsafe {
                h.add_to_free_list(p);
            }
        }
        h
    }

    fn pin(&self, p: *mut PairStorage<'a>) {
        let mut pins = self.pins.borrow_mut();
        let entry = pins.entry(p).or_insert(0);
        *entry += 1;
    }

    fn unpin(&self, p: *mut PairStorage<'a>) {
        let mut pins = self.pins.borrow_mut();
        if {
            let entry = pins.entry(p).or_insert(0);
            assert!(*entry != 0);
            *entry -= 1;
            *entry == 0
        } {
            pins.remove(&p);
        }
    }

    unsafe fn add_to_free_list(&mut self, p: *mut PairStorage<'a>) {
        let listp = p as *mut *mut PairStorage<'a>;
        *listp = self.freelist;
        assert_eq!(*listp, self.freelist);
        self.freelist = p;
    }

    pub fn try_alloc(&mut self, pair: (Value<'a>, Value<'a>)) -> Option<Pair<'a>> {
        if self.freelist.is_null() {
            unsafe {
                self.gc();
            }
            if self.freelist.is_null() {
                return None;
            }
        }

        let p = self.freelist;
        unsafe {
            let listp = p as *mut *mut PairStorage<'a>;
            self.freelist = *listp;

            let (h, t) = pair;
            *p = PairStorage {
                head: h.to_heap(),
                tail: t.to_heap()
            };
            Some(Pair(PinnedRef::new(self, p)))
        }
    }

    unsafe fn find_mark_bit<T>(_ptr: *const T) -> (*mut BitVec, usize) {
        assert_eq!(mem::size_of::<T>(), mem::size_of::<PairStorage<'a>>());
        let storage_addr = _ptr as usize & !(HEAP_STORAGE_ALIGN - 1);
        let storage = storage_addr as *mut HeapStorage<'a>;
        let objects_addr = &mut (*storage).objects[0] as *mut PairStorage<'a> as usize;
        let index = (_ptr as usize - objects_addr) / mem::size_of::<PairStorage<'a>>();
        (&mut (*storage).mark_bits as *mut BitVec, index)
    }

    pub unsafe fn get_mark_bit<T>(_ptr: *const T) -> bool {
        let (bits, index) = Heap::find_mark_bit(_ptr);
        (*bits)[index]
    }

    pub unsafe fn set_mark_bit<T>(_ptr: *const T) {
        let (bits, index) = Heap::find_mark_bit(_ptr);
        (*bits).set(index, true);
    }

    pub fn alloc(&mut self, pair: (Value<'a>, Value<'a>)) -> Pair<'a> {
        self.try_alloc(pair).expect("out of memory (gc did not collect anything)")
    }

    pub fn alloc_null(&mut self) -> Pair<'a> {
        self.alloc((Value::Null, Value::Null))
    }

    unsafe fn gc(&mut self) {
        // mark phase
        self.storage.mark_bits.clear();
        for (&p, _) in self.pins.borrow().iter() {
            Mark::mark(p);
        }

        // sweep phase
        self.freelist = ptr::null_mut();
        for i in 0 .. HEAP_SIZE {
            let p = &mut self.storage.objects[i] as *mut PairStorage<'a>;
            if !self.storage.mark_bits[i] {
                self.add_to_free_list(p);
            }
        }
    }

    #[cfg(test)]
    pub fn force_gc(&mut self) {
        unsafe { self.gc(); }
    }
}

pub struct PinnedRef<'a, T> {
    ptr: *mut T,
    heap: *const Heap<'a>,
    heap_id: HeapId<'a>
}

impl<'a, T> PinnedRef<'a, T> {
    unsafe fn new(heap: &Heap<'a>, p: *mut T) -> PinnedRef<'a, T> {
        heap.pin(p as *mut PairStorage<'a>);  // XXX BOGUS
        PinnedRef {
            ptr: p,
            heap: heap as *const Heap<'a>,
            heap_id: heap.id
        }
    }
}

impl<'a, T> Drop for PinnedRef<'a, T> {
    fn drop(&mut self) {
        unsafe {
            (*self.heap).unpin(self.ptr as *mut PairStorage<'a>);  // XXX BOGUS
        }
    }
}

impl<'a, T> Clone for PinnedRef<'a, T> {
    fn clone(&self) -> PinnedRef<'a, T> {
        let &PinnedRef { ptr, heap, heap_id } = self;
        unsafe {
            (*heap).pin(ptr as *mut PairStorage<'a>);  // XXX BOGUS
        }
        PinnedRef {
            ptr: ptr,
            heap: heap,
            heap_id: heap_id
        }
    }
}

impl<'a, T> fmt::Debug for PinnedRef<'a, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "PinnedRef {{ ptr: {:p} }}", self.ptr)
    }
}

impl<'a, T> PartialEq for PinnedRef<'a, T> {
    fn eq(&self, other: &PinnedRef<'a, T>) -> bool {
        self.ptr == other.ptr
    }
}

impl<'a, T> Eq for PinnedRef<'a, T> {}

pub trait GCRef {
    #[cfg(test)]
    fn address(&self) -> usize;
}

// === Pair, the reference type

gc_ref_type! {
    pub struct Pair / PairStorage<'a> {
        head / set_head: Value<'a>,
        tail / set_tail: Value<'a>
    }
}

// === Values (a heap-inline enum)

use std::rc::Rc;

// Values inside the heap (GC heap-to-heap cross-references) (private)
enum ValueStorage<'a> {
    Null,
    Int(i32),
    Str(Rc<String>),
    Pair(*mut PairStorage<'a>)
}

unsafe impl<'a> Mark<'a> for ValueStorage<'a> {
    unsafe fn mark(ptr: *mut ValueStorage<'a>) {
        match *ptr {
            ValueStorage::Pair(p) => Mark::mark(p),
            _ => {}
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Value<'a> {
    Null,
    Int(i32),
    Str(Rc<String>),  // <-- equality is by value
    Pair(Pair<'a>)  // <-- equality is by pointer
}

unsafe impl<'a> HeapInline<'a> for Value<'a> {
    type Storage = ValueStorage<'a>;

    fn to_heap(self) -> ValueStorage<'a> {
        match self {
            Value::Null => ValueStorage::Null,
            Value::Int(n) => ValueStorage::Int(n),
            Value::Str(rcstr) => ValueStorage::Str(rcstr),
            Value::Pair(pair) => ValueStorage::Pair(HeapInline::<'a>::to_heap(pair))
        }
    }

    unsafe fn from_heap(heap: &Heap<'a>, v: &ValueStorage<'a>) -> Value<'a> {
        match v {
            &ValueStorage::Null => Value::Null,
            &ValueStorage::Int(n) => Value::Int(n),
            &ValueStorage::Str(ref rcstr) => Value::Str(rcstr.clone()),
            &ValueStorage::Pair(ref ptr) => Value::Pair(HeapInline::<'a>::from_heap(heap, ptr))
        }
    }
}

unsafe impl<'a, T: Copy + 'static> Mark<'a> for T {
    unsafe fn mark(_ptr: *mut T) {}
}

unsafe impl<'a, T: Copy + 'static> HeapInline<'a> for T {
    type Storage = Self;

    fn to_heap(self) -> T { self }

    unsafe fn from_heap(_heap: &Heap<'a>, v: &T) -> T { *v }
}
<|endoftext|>


<commit_before>// # Deleting and Updating (?) documents
//
// This example explains how to delete and update documents.
// In fact there is actually no such thing as an update in tantivy.
//
// To update a document, you need to delete a document and then reinsert
// its new version.
//
// ---
// Importing tantivy...
use tantivy::collector::TopDocs;
use tantivy::query::TermQuery;
use tantivy::schema::*;
use tantivy::{doc, Index, IndexReader};

// A simple helper function to fetch a single document
// given its id from our index.
// It will be helpful to check our work.
fn extract_doc_given_isbn(
    reader: &IndexReader,
    isbn_term: &Term,
) -> tantivy::Result<Option<Document>> {
    let searcher = reader.searcher();

    // This is the simplest query you can think of.
    // It matches all of the documents containing a specific term.
    //
    // The second argument is here to tell we don't care about decoding positions,
    // or term frequencies.
    let term_query = TermQuery::new(isbn_term.clone(), IndexRecordOption::Basic);
    let top_docs = searcher.search(&term_query, &TopDocs::with_limit(1))?;

    if let Some((_score, doc_address)) = top_docs.first() {
        let doc = searcher.doc(*doc_address)?;
        Ok(Some(doc))
    } else {
        // no doc matching this ID.
        Ok(None)
    }
}

fn main() -> tantivy::Result<()> {
    // # Defining the schema
    //
    // Check out the *basic_search* example if this makes
    // small sense to you.
    let mut schema_builder = Schema::builder();

    // Tantivy does not really have a notion of primary id.
    // This may change in the future.
    //
    // Still, we can create a `isbn` field and use it as an id. This
    // field can be `u64` or a `text`, depending on your use case.
    // It just needs to be indexed.
    //
    // If it is `text`, let's make sure to keep it `raw` and let's avoid
    // running any text processing on it.
    // This is done by associating this field to the tokenizer named `raw`.
    // Rather than building our [`TextOptions`](//docs.rs/tantivy/~0/tantivy/schema/struct.TextOptions.html) manually,
    // We use the `STRING` shortcut. `STRING` stands for indexed (without term frequency or positions)
    // and untokenized.
    //
    // Because we also want to be able to see this `id` in our returned documents,
    // we also mark the field as stored.
    let isbn = schema_builder.add_text_field("isbn", STRING | STORED);
    let title = schema_builder.add_text_field("title", TEXT | STORED);
    let schema = schema_builder.build();

    let index = Index::create_in_ram(schema.clone());

    let mut index_writer = index.writer(50_000_000)?;

    // Let's add a couple of documents, for the sake of the example.
    let mut old_man_doc = Document::default();
    old_man_doc.add_text(title, "The Old Man and the Sea");
    index_writer.add_document(doc!(
        isbn => "978-0099908401",
        title => "The old Man and the see"
    ));
    index_writer.add_document(doc!(
        isbn => "978-0140177398",
        title => "Of Mice and Men",
    ));
    index_writer.add_document(doc!(
       title => "Frankentein", //< Oops there is a typo here.
       isbn => "978-9176370711",
    ));
    index_writer.commit()?;
    let reader = index.reader()?;

    let frankenstein_isbn = Term::from_field_text(isbn, "978-9176370711");

    // Oops our frankenstein doc seems mispelled
    let frankenstein_doc_misspelled = extract_doc_given_isbn(&reader, &frankenstein_isbn)?.unwrap();
    assert_eq!(
        schema.to_json(&frankenstein_doc_misspelled),
        r#"{"isbn":["978-9176370711"],"title":["Frankentein"]}"#,
    );

    // # Update = Delete + Insert
    //
    // Here we will want to update the typo in the `Frankenstein` book.
    //
    // Tantivy does not handle updates directly, we need to delete
    // and reinsert the document.
    //
    // This can be complicated as it means you need to have access
    // to the entire document. It is good practise to integrate tantivy
    // with a key value store for this reason.
    //
    // To remove one of the document, we just call `delete_term`
    // on its id.
    //
    // Note that `tantivy` does nothing to enforce the idea that
    // there is only one document associated to this id.
    //
    // Also you might have noticed that we apply the delete before
    // having committed. This does not matter really...
    index_writer.delete_term(frankenstein_isbn.clone());

    // We now need to reinsert our document without the typo.
    index_writer.add_document(doc!(
       title => "Frankenstein",
       isbn => "978-9176370711",
    ));

    // You are guaranteed that your clients will only observe your index in
    // the state it was in after a commit.
    // In this example, your search engine will at no point be missing the *Frankenstein* document.
    // Everything happened as if the document was updated.
    index_writer.commit()?;
    // We reload our searcher to make our change available to clients.
    reader.reload()?;

    // No more typo!
    let frankenstein_new_doc = extract_doc_given_isbn(&reader, &frankenstein_isbn)?.unwrap();
    assert_eq!(
        schema.to_json(&frankenstein_new_doc),
        r#"{"isbn":["978-9176370711"],"title":["Frankenstein"]}"#,
    );

    Ok(())
}
<commit_msg>Misspelling of misspelled was fixed (#1078)<commit_after>// # Deleting and Updating (?) documents
//
// This example explains how to delete and update documents.
// In fact there is actually no such thing as an update in tantivy.
//
// To update a document, you need to delete a document and then reinsert
// its new version.
//
// ---
// Importing tantivy...
use tantivy::collector::TopDocs;
use tantivy::query::TermQuery;
use tantivy::schema::*;
use tantivy::{doc, Index, IndexReader};

// A simple helper function to fetch a single document
// given its id from our index.
// It will be helpful to check our work.
fn extract_doc_given_isbn(
    reader: &IndexReader,
    isbn_term: &Term,
) -> tantivy::Result<Option<Document>> {
    let searcher = reader.searcher();

    // This is the simplest query you can think of.
    // It matches all of the documents containing a specific term.
    //
    // The second argument is here to tell we don't care about decoding positions,
    // or term frequencies.
    let term_query = TermQuery::new(isbn_term.clone(), IndexRecordOption::Basic);
    let top_docs = searcher.search(&term_query, &TopDocs::with_limit(1))?;

    if let Some((_score, doc_address)) = top_docs.first() {
        let doc = searcher.doc(*doc_address)?;
        Ok(Some(doc))
    } else {
        // no doc matching this ID.
        Ok(None)
    }
}

fn main() -> tantivy::Result<()> {
    // # Defining the schema
    //
    // Check out the *basic_search* example if this makes
    // small sense to you.
    let mut schema_builder = Schema::builder();

    // Tantivy does not really have a notion of primary id.
    // This may change in the future.
    //
    // Still, we can create a `isbn` field and use it as an id. This
    // field can be `u64` or a `text`, depending on your use case.
    // It just needs to be indexed.
    //
    // If it is `text`, let's make sure to keep it `raw` and let's avoid
    // running any text processing on it.
    // This is done by associating this field to the tokenizer named `raw`.
    // Rather than building our [`TextOptions`](//docs.rs/tantivy/~0/tantivy/schema/struct.TextOptions.html) manually,
    // We use the `STRING` shortcut. `STRING` stands for indexed (without term frequency or positions)
    // and untokenized.
    //
    // Because we also want to be able to see this `id` in our returned documents,
    // we also mark the field as stored.
    let isbn = schema_builder.add_text_field("isbn", STRING | STORED);
    let title = schema_builder.add_text_field("title", TEXT | STORED);
    let schema = schema_builder.build();

    let index = Index::create_in_ram(schema.clone());

    let mut index_writer = index.writer(50_000_000)?;

    // Let's add a couple of documents, for the sake of the example.
    let mut old_man_doc = Document::default();
    old_man_doc.add_text(title, "The Old Man and the Sea");
    index_writer.add_document(doc!(
        isbn => "978-0099908401",
        title => "The old Man and the see"
    ));
    index_writer.add_document(doc!(
        isbn => "978-0140177398",
        title => "Of Mice and Men",
    ));
    index_writer.add_document(doc!(
       title => "Frankentein", //< Oops there is a typo here.
       isbn => "978-9176370711",
    ));
    index_writer.commit()?;
    let reader = index.reader()?;

    let frankenstein_isbn = Term::from_field_text(isbn, "978-9176370711");

    // Oops our frankenstein doc seems misspelled
    let frankenstein_doc_misspelled = extract_doc_given_isbn(&reader, &frankenstein_isbn)?.unwrap();
    assert_eq!(
        schema.to_json(&frankenstein_doc_misspelled),
        r#"{"isbn":["978-9176370711"],"title":["Frankentein"]}"#,
    );

    // # Update = Delete + Insert
    //
    // Here we will want to update the typo in the `Frankenstein` book.
    //
    // Tantivy does not handle updates directly, we need to delete
    // and reinsert the document.
    //
    // This can be complicated as it means you need to have access
    // to the entire document. It is good practise to integrate tantivy
    // with a key value store for this reason.
    //
    // To remove one of the document, we just call `delete_term`
    // on its id.
    //
    // Note that `tantivy` does nothing to enforce the idea that
    // there is only one document associated to this id.
    //
    // Also you might have noticed that we apply the delete before
    // having committed. This does not matter really...
    index_writer.delete_term(frankenstein_isbn.clone());

    // We now need to reinsert our document without the typo.
    index_writer.add_document(doc!(
       title => "Frankenstein",
       isbn => "978-9176370711",
    ));

    // You are guaranteed that your clients will only observe your index in
    // the state it was in after a commit.
    // In this example, your search engine will at no point be missing the *Frankenstein* document.
    // Everything happened as if the document was updated.
    index_writer.commit()?;
    // We reload our searcher to make our change available to clients.
    reader.reload()?;

    // No more typo!
    let frankenstein_new_doc = extract_doc_given_isbn(&reader, &frankenstein_isbn)?.unwrap();
    assert_eq!(
        schema.to_json(&frankenstein_new_doc),
        r#"{"isbn":["978-9176370711"],"title":["Frankenstein"]}"#,
    );

    Ok(())
}
<|endoftext|>


<commit_before>extern crate http;
extern crate hyper;
extern crate native_tls;
extern crate tokio_tcp;
extern crate tokio_tls;

use std::{
    fmt, fs,
    marker::PhantomData,
    path::Path,
    sync::{Arc, Mutex},
};

use self::{http::StatusCode, hyper::service::Service, hyper::Chunk, hyper::Server};

use futures::{
    future::{self, Either},
    Async, Future, Stream,
};

use self::http::header::{HeaderMap, HeaderName, HeaderValue};

use base::types::{ArcType, Type};

use vm::{
    self,
    api::{Collect, Function, Getable, OpaqueValue, PushAsRef, Pushable, VmType, WithVM, IO},
    thread::{ActiveThread, RootedThread, Thread},
    ExternModule, Variants,
};

macro_rules! try_future {
    ($e:expr) => {
        try_future!($e, Box::new)
    };
    ($e:expr, $f:expr) => {
        match $e {
            Ok(x) => x,
            Err(err) => return $f(::futures::future::err(err.into())),
        }
    };
}

mod std {
    pub use http;
}

// `Handler` is a type defined in http.glu but since we need to refer to it in the signature of
// listen we define a phantom type which we can use with `OpaqueValue` to store a `Handler` in Rust
struct Handler<T>(PhantomData<T>);

impl<T: VmType + 'static> VmType for Handler<T> {
    type Type = Self;
    fn make_type(vm: &Thread) -> ArcType {
        vm.find_type_info("std.http.types.Handler")
            .map(|typ| Type::app(typ.into_type(), collect![T::make_type(vm)]))
            .unwrap_or_else(|_| Type::hole())
    }
}

struct Headers(HeaderMap);

impl VmType for Headers {
    type Type = Vec<(String, Vec<u8>)>;

    fn make_type(vm: &Thread) -> ArcType {
        Vec::<(String, Vec<u8>)>::make_type(vm)
    }
}

impl<'vm> Pushable<'vm> for Headers {
    fn push(self, context: &mut ActiveThread<'vm>) -> vm::Result<()> {
        Collect::new(
            self.0
                .iter()
                .map(|(name, value)| (name.as_str(), value.as_bytes())),
        )
        .push(context)
    }
}

impl<'vm, 'value> Getable<'vm, 'value> for Headers {
    fn from_value(vm: &'vm Thread, value: Variants<'value>) -> Self {
        Headers(
            Collect::from_value(vm, value)
                // TODO Error somehow on invalid headers
                .filter_map(|(name, value): (&str, &[u8])| {
                    match (HeaderName::from_bytes(name.as_bytes()), HeaderValue::from_bytes(value)) {
                        (Ok(name), Ok(value)) => Some((name, value)),
                        _ => None,
                    }
                })
                .collect(),
        )
    }
}

// By implementing `Userdata` on `Body` it can be automatically pushed and retrieved from gluon
// threads
#[derive(Userdata)]
#[gluon(crate_name = "::vm")]
// Representation of a http body that is in the prograss of being read
pub struct Body(Arc<Mutex<Box<Stream<Item = PushAsRef<Chunk, [u8]>, Error = vm::Error> + Send>>>);

// Types implementing `Userdata` requires a `std::fmt::Debug` implementation so it can be displayed
impl fmt::Debug for Body {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "hyper::Body")
    }
}

// Since `Body` implements `Userdata` gluon will automatically marshal the gluon representation
// into `&Body` argument
fn read_chunk(
    body: &Body,
) -> impl Future<Item = IO<Option<PushAsRef<Chunk, [u8]>>>, Error = vm::Error> {
    use futures::future::poll_fn;

    let body = body.0.clone();
    poll_fn(move || {
        let mut stream = body.lock().unwrap();
        stream.poll().map(|async| async.map(IO::Value))
    })
}

// A http body that is being written
#[derive(Userdata)]
#[gluon(crate_name = "::vm")]
pub struct ResponseBody(Arc<Mutex<Option<hyper::body::Sender>>>);

impl fmt::Debug for ResponseBody {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "hyper::Response")
    }
}

fn write_response(
    response: &ResponseBody,
    bytes: &[u8],
) -> impl Future<Item = IO<()>, Error = vm::Error> {
    use futures::future::poll_fn;

    // Turn `bytes´ into a `Chunk` which can be sent to the http body
    let mut unsent_chunk = Some(bytes.to_owned().into());
    let response = response.0.clone();
    poll_fn(move || {
        info!("Starting response send");
        let mut sender = response.lock().unwrap();
        let sender = sender
            .as_mut()
            .expect("Sender has been dropped while still in use");
        let chunk = unsent_chunk
            .take()
            .expect("Attempt to poll after chunk is sent");
        match sender.poll_ready() {
            Ok(Async::NotReady) => {
                unsent_chunk = Some(chunk);
                return Ok(Async::NotReady);
            }
            Ok(Async::Ready(_)) => (),
            Err(_) => {
                info!("Could not send http response");
                return Ok(Async::Ready(IO::Value(())));
            }
        }
        match sender.send_data(chunk) {
            Ok(()) => Ok(Async::Ready(IO::Value(()))),
            Err(chunk) => {
                info!("Could not send http response");
                unsent_chunk = Some(chunk);
                Ok(Async::NotReady)
            }
        }
    })
}

#[derive(Debug, Userdata)]
#[gluon(crate_name = "::vm")]
struct Uri(http::Uri);

// Next we define some record types which are marshalled to and from gluon. These have equivalent
// definitions in http_types.glu
field_decl! { method, uri, status, body, request, response, headers }

type Request = record_type! {
    method => String,
    uri => Uri,
    body => Body
};

type Response = record_type! {
    status => u16,
    headers => Headers
};

type HttpState = record_type! {
    request => Request,
    response => ResponseBody
};

#[derive(Getable, VmType)]
#[gluon(crate_name = "::vm")]
struct Settings<'a> {
    port: u16,
    tls_cert: Option<&'a Path>,
}

fn listen(
    settings: Settings,
    value: WithVM<OpaqueValue<RootedThread, Handler<Response>>>,
) -> impl Future<Item = IO<()>, Error = vm::Error> + Send + 'static {
    let WithVM {
        value: handler,
        vm: thread,
    } = value;

    let thread = match thread.new_thread() {
        Ok(thread) => thread,
        Err(err) => return Either::A(future::err(err)),
    };

    // Retrieve the `handle` function from the http module which we use to evaluate values of type
    // `Handler Response`
    type ListenFn = fn(OpaqueValue<RootedThread, Handler<Response>>, HttpState) -> IO<Response>;
    let handle: Function<RootedThread, ListenFn> = thread
        .get_global("std.http.http.handle")
        .unwrap_or_else(|err| panic!("{}", err));

    struct Listen {
        handle: Function<RootedThread, ListenFn>,
        handler: OpaqueValue<RootedThread, Handler<Response>>,
    }

    impl Service for Listen {
        type ReqBody = hyper::Body;
        type ResBody = hyper::Body;
        type Error = vm::Error;
        type Future =
            Box<Future<Item = http::Response<hyper::Body>, Error = Self::Error> + Send + 'static>;

        fn call(&mut self, request: http::Request<hyper::Body>) -> Self::Future {
            let (parts, body) = request.into_parts();
            let gluon_request = record_no_decl! {
                method => parts.method.as_str().to_owned(),
                uri => Uri(parts.uri),
                // Since `Body` implements `Userdata` it can be directly pushed to gluon
                body => Body(Arc::new(Mutex::new(Box::new(body
                    .map_err(|err| vm::Error::Message(format!("{}", err)))
                    // `PushAsRef` makes the `body` parameter act as a `&[u8]` which means it is
                    // marshalled to `Array Byte` in gluon
                    .map(PushAsRef::<_, [u8]>::new)))))
            };
            let (response_sender, response_body) = hyper::Body::channel();
            let response_sender = Arc::new(Mutex::new(Some(response_sender)));
            let http_state = record_no_decl! {
                request => gluon_request,
                response => ResponseBody(response_sender.clone())
            };

            let child_thread = try_future!(self.handle.vm().new_thread());
            let mut handle = try_future!(self.handle.re_root(child_thread));

            Box::new(
                handle.call_async(self.handler.clone(), http_state).then(
                    move |result| match result {
                        Ok(value) => {
                            match value {
                                IO::Value(record_p! { status, headers }) => {
                                    // Drop the sender to so that it the receiver stops waiting for
                                    // more chunks
                                    *response_sender.lock().unwrap() = None;

                                    let status = StatusCode::from_u16(status)
                                        .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                                    let mut response = http::Response::builder()
                                        .status(status)
                                        .body(response_body)
                                        .unwrap();
                                    *response.headers_mut() = headers.0;
                                    Ok(response)
                                }
                                IO::Exception(err) => {
                                    error!("{}", err);
                                    Ok(http::Response::builder()
                                        .status(StatusCode::INTERNAL_SERVER_ERROR)
                                        .body("".into())
                                        .unwrap())
                                }
                            }
                        }
                        Err(err) => {
                            error!("{}", err);
                            Ok(http::Response::builder()
                                .status(StatusCode::INTERNAL_SERVER_ERROR)
                                .body("".into())
                                .unwrap())
                        }
                    },
                ),
            )
        }
    }

    let addr = format!("0.0.0.0:{}", settings.port).parse().unwrap();

    if let Some(cert_path) = settings.tls_cert {
        let identity = try_future!(
            fs::read(cert_path).map_err(|err| vm::Error::Message(format!(
                "Unable to open certificate `{}`: {}",
                cert_path.display(),
                err
            ))),
            Either::A
        );

        let identity = try_future!(
            native_tls::Identity::from_pkcs12(&identity, "")
                .map_err(|err| vm::Error::Message(err.to_string())),
            Either::A
        );
        let acceptor = tokio_tls::TlsAcceptor::from(try_future!(
            native_tls::TlsAcceptor::new(identity)
                .map_err(|err| vm::Error::Message(err.to_string())),
            Either::A
        ));

        let http = hyper::server::conn::Http::new();
        let listener = try_future!(
            tokio_tcp::TcpListener::bind(&addr).map_err(|err| vm::Error::Message(err.to_string())),
            Either::A
        );
        let incoming = listener
            .incoming()
            .and_then(move |stream| {
                acceptor.accept(stream).map(Some).or_else(|err| {
                    info!("Unable to accept TLS connection: {}", err);
                    Ok(None)
                })
            })
            .filter_map(|opt_tls_stream| opt_tls_stream);

        let future = http
            .serve_incoming(incoming, move || -> Result<_, hyper::Error> {
                Ok(Listen {
                    handle: handle.clone(),
                    handler: handler.clone(),
                })
            })
            .and_then(|connecting| connecting)
            .for_each(|connection| {
                hyper::rt::spawn(connection.map_err(|err| error!("{}", err)));
                Ok(())
            })
            .map(|_| IO::Value(()))
            .map_err(|err| vm::Error::from(format!("Server error: {}", err)));
        return Either::B(Either::A(future));
    }

    Either::B(Either::B(
        Server::bind(&addr)
            .serve(move || -> Result<_, hyper::Error> {
                Ok(Listen {
                    handle: handle.clone(),
                    handler: handler.clone(),
                })
            })
            .map_err(|err| vm::Error::from(format!("Server error: {}", err)))
            .map(|_| IO::Value(())),
    ))
}

// To let the `http_types` module refer to `Body` and `ResponseBody` we register these types in a
// separate function which is called before loading `http_types`
pub fn load_types(vm: &Thread) -> vm::Result<ExternModule> {
    vm.register_type::<Body>("Body", &[])?;
    vm.register_type::<ResponseBody>("ResponseBody", &[])?;
    vm.register_type::<Uri>("Uri", &[])?;

    ExternModule::new(
        vm,
        record! {
            // Define the types so that they can be used from gluon
            type Body => Body,
            type ResponseBody => ResponseBody,
            type Uri => Uri,
            type Method => String,
            type StatusCode => u16,
            type Request => Request,
            type Response => Response,
            type Headers => Headers,
            type HttpState => HttpState
        },
    )
}

macro_rules! uri_binds {
    ($($id: ident)*) => {
        record!{
            $(
                $id => primitive!(1, concat!("std.http.prim.uri.", stringify!($id)), |u: &Uri| (u.0).$id())
            ),*
        }
    }
}

pub fn load(vm: &Thread) -> vm::Result<ExternModule> {
    ExternModule::new(
        vm,
        record! {
            listen => primitive!(2, async fn listen),
            read_chunk => primitive!(1, async fn read_chunk),
            write_response => primitive!(2, async fn write_response),
            uri => uri_binds!(path host port query to_string)
        },
    )
}
<commit_msg>Don't use deprecated port function<commit_after>extern crate http;
extern crate hyper;
extern crate native_tls;
extern crate tokio_tcp;
extern crate tokio_tls;

use std::{
    fmt, fs,
    marker::PhantomData,
    path::Path,
    sync::{Arc, Mutex},
};

use self::{http::StatusCode, hyper::service::Service, hyper::Chunk, hyper::Server};

use futures::{
    future::{self, Either},
    Async, Future, Stream,
};

use self::http::header::{HeaderMap, HeaderName, HeaderValue};

use base::types::{ArcType, Type};

use vm::{
    self,
    api::{Collect, Function, Getable, OpaqueValue, PushAsRef, Pushable, VmType, WithVM, IO},
    thread::{ActiveThread, RootedThread, Thread},
    ExternModule, Variants,
};

macro_rules! try_future {
    ($e:expr) => {
        try_future!($e, Box::new)
    };
    ($e:expr, $f:expr) => {
        match $e {
            Ok(x) => x,
            Err(err) => return $f(::futures::future::err(err.into())),
        }
    };
}

mod std {
    pub use http;
}

// `Handler` is a type defined in http.glu but since we need to refer to it in the signature of
// listen we define a phantom type which we can use with `OpaqueValue` to store a `Handler` in Rust
struct Handler<T>(PhantomData<T>);

impl<T: VmType + 'static> VmType for Handler<T> {
    type Type = Self;
    fn make_type(vm: &Thread) -> ArcType {
        vm.find_type_info("std.http.types.Handler")
            .map(|typ| Type::app(typ.into_type(), collect![T::make_type(vm)]))
            .unwrap_or_else(|_| Type::hole())
    }
}

struct Headers(HeaderMap);

impl VmType for Headers {
    type Type = Vec<(String, Vec<u8>)>;

    fn make_type(vm: &Thread) -> ArcType {
        Vec::<(String, Vec<u8>)>::make_type(vm)
    }
}

impl<'vm> Pushable<'vm> for Headers {
    fn push(self, context: &mut ActiveThread<'vm>) -> vm::Result<()> {
        Collect::new(
            self.0
                .iter()
                .map(|(name, value)| (name.as_str(), value.as_bytes())),
        )
        .push(context)
    }
}

impl<'vm, 'value> Getable<'vm, 'value> for Headers {
    fn from_value(vm: &'vm Thread, value: Variants<'value>) -> Self {
        Headers(
            Collect::from_value(vm, value)
                // TODO Error somehow on invalid headers
                .filter_map(|(name, value): (&str, &[u8])| {
                    match (HeaderName::from_bytes(name.as_bytes()), HeaderValue::from_bytes(value)) {
                        (Ok(name), Ok(value)) => Some((name, value)),
                        _ => None,
                    }
                })
                .collect(),
        )
    }
}

// By implementing `Userdata` on `Body` it can be automatically pushed and retrieved from gluon
// threads
#[derive(Userdata)]
#[gluon(crate_name = "::vm")]
// Representation of a http body that is in the prograss of being read
pub struct Body(Arc<Mutex<Box<Stream<Item = PushAsRef<Chunk, [u8]>, Error = vm::Error> + Send>>>);

// Types implementing `Userdata` requires a `std::fmt::Debug` implementation so it can be displayed
impl fmt::Debug for Body {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "hyper::Body")
    }
}

// Since `Body` implements `Userdata` gluon will automatically marshal the gluon representation
// into `&Body` argument
fn read_chunk(
    body: &Body,
) -> impl Future<Item = IO<Option<PushAsRef<Chunk, [u8]>>>, Error = vm::Error> {
    use futures::future::poll_fn;

    let body = body.0.clone();
    poll_fn(move || {
        let mut stream = body.lock().unwrap();
        stream.poll().map(|async| async.map(IO::Value))
    })
}

// A http body that is being written
#[derive(Userdata)]
#[gluon(crate_name = "::vm")]
pub struct ResponseBody(Arc<Mutex<Option<hyper::body::Sender>>>);

impl fmt::Debug for ResponseBody {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "hyper::Response")
    }
}

fn write_response(
    response: &ResponseBody,
    bytes: &[u8],
) -> impl Future<Item = IO<()>, Error = vm::Error> {
    use futures::future::poll_fn;

    // Turn `bytes´ into a `Chunk` which can be sent to the http body
    let mut unsent_chunk = Some(bytes.to_owned().into());
    let response = response.0.clone();
    poll_fn(move || {
        info!("Starting response send");
        let mut sender = response.lock().unwrap();
        let sender = sender
            .as_mut()
            .expect("Sender has been dropped while still in use");
        let chunk = unsent_chunk
            .take()
            .expect("Attempt to poll after chunk is sent");
        match sender.poll_ready() {
            Ok(Async::NotReady) => {
                unsent_chunk = Some(chunk);
                return Ok(Async::NotReady);
            }
            Ok(Async::Ready(_)) => (),
            Err(_) => {
                info!("Could not send http response");
                return Ok(Async::Ready(IO::Value(())));
            }
        }
        match sender.send_data(chunk) {
            Ok(()) => Ok(Async::Ready(IO::Value(()))),
            Err(chunk) => {
                info!("Could not send http response");
                unsent_chunk = Some(chunk);
                Ok(Async::NotReady)
            }
        }
    })
}

#[derive(Debug, Userdata)]
#[gluon(crate_name = "::vm")]
struct Uri(http::Uri);

// Next we define some record types which are marshalled to and from gluon. These have equivalent
// definitions in http_types.glu
field_decl! { method, uri, status, body, request, response, headers }

type Request = record_type! {
    method => String,
    uri => Uri,
    body => Body
};

type Response = record_type! {
    status => u16,
    headers => Headers
};

type HttpState = record_type! {
    request => Request,
    response => ResponseBody
};

#[derive(Getable, VmType)]
#[gluon(crate_name = "::vm")]
struct Settings<'a> {
    port: u16,
    tls_cert: Option<&'a Path>,
}

fn listen(
    settings: Settings,
    value: WithVM<OpaqueValue<RootedThread, Handler<Response>>>,
) -> impl Future<Item = IO<()>, Error = vm::Error> + Send + 'static {
    let WithVM {
        value: handler,
        vm: thread,
    } = value;

    let thread = match thread.new_thread() {
        Ok(thread) => thread,
        Err(err) => return Either::A(future::err(err)),
    };

    // Retrieve the `handle` function from the http module which we use to evaluate values of type
    // `Handler Response`
    type ListenFn = fn(OpaqueValue<RootedThread, Handler<Response>>, HttpState) -> IO<Response>;
    let handle: Function<RootedThread, ListenFn> = thread
        .get_global("std.http.http.handle")
        .unwrap_or_else(|err| panic!("{}", err));

    struct Listen {
        handle: Function<RootedThread, ListenFn>,
        handler: OpaqueValue<RootedThread, Handler<Response>>,
    }

    impl Service for Listen {
        type ReqBody = hyper::Body;
        type ResBody = hyper::Body;
        type Error = vm::Error;
        type Future =
            Box<Future<Item = http::Response<hyper::Body>, Error = Self::Error> + Send + 'static>;

        fn call(&mut self, request: http::Request<hyper::Body>) -> Self::Future {
            let (parts, body) = request.into_parts();
            let gluon_request = record_no_decl! {
                method => parts.method.as_str().to_owned(),
                uri => Uri(parts.uri),
                // Since `Body` implements `Userdata` it can be directly pushed to gluon
                body => Body(Arc::new(Mutex::new(Box::new(body
                    .map_err(|err| vm::Error::Message(format!("{}", err)))
                    // `PushAsRef` makes the `body` parameter act as a `&[u8]` which means it is
                    // marshalled to `Array Byte` in gluon
                    .map(PushAsRef::<_, [u8]>::new)))))
            };
            let (response_sender, response_body) = hyper::Body::channel();
            let response_sender = Arc::new(Mutex::new(Some(response_sender)));
            let http_state = record_no_decl! {
                request => gluon_request,
                response => ResponseBody(response_sender.clone())
            };

            let child_thread = try_future!(self.handle.vm().new_thread());
            let mut handle = try_future!(self.handle.re_root(child_thread));

            Box::new(
                handle.call_async(self.handler.clone(), http_state).then(
                    move |result| match result {
                        Ok(value) => {
                            match value {
                                IO::Value(record_p! { status, headers }) => {
                                    // Drop the sender to so that it the receiver stops waiting for
                                    // more chunks
                                    *response_sender.lock().unwrap() = None;

                                    let status = StatusCode::from_u16(status)
                                        .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                                    let mut response = http::Response::builder()
                                        .status(status)
                                        .body(response_body)
                                        .unwrap();
                                    *response.headers_mut() = headers.0;
                                    Ok(response)
                                }
                                IO::Exception(err) => {
                                    error!("{}", err);
                                    Ok(http::Response::builder()
                                        .status(StatusCode::INTERNAL_SERVER_ERROR)
                                        .body("".into())
                                        .unwrap())
                                }
                            }
                        }
                        Err(err) => {
                            error!("{}", err);
                            Ok(http::Response::builder()
                                .status(StatusCode::INTERNAL_SERVER_ERROR)
                                .body("".into())
                                .unwrap())
                        }
                    },
                ),
            )
        }
    }

    let addr = format!("0.0.0.0:{}", settings.port).parse().unwrap();

    if let Some(cert_path) = settings.tls_cert {
        let identity = try_future!(
            fs::read(cert_path).map_err(|err| vm::Error::Message(format!(
                "Unable to open certificate `{}`: {}",
                cert_path.display(),
                err
            ))),
            Either::A
        );

        let identity = try_future!(
            native_tls::Identity::from_pkcs12(&identity, "")
                .map_err(|err| vm::Error::Message(err.to_string())),
            Either::A
        );
        let acceptor = tokio_tls::TlsAcceptor::from(try_future!(
            native_tls::TlsAcceptor::new(identity)
                .map_err(|err| vm::Error::Message(err.to_string())),
            Either::A
        ));

        let http = hyper::server::conn::Http::new();
        let listener = try_future!(
            tokio_tcp::TcpListener::bind(&addr).map_err(|err| vm::Error::Message(err.to_string())),
            Either::A
        );
        let incoming = listener
            .incoming()
            .and_then(move |stream| {
                acceptor.accept(stream).map(Some).or_else(|err| {
                    info!("Unable to accept TLS connection: {}", err);
                    Ok(None)
                })
            })
            .filter_map(|opt_tls_stream| opt_tls_stream);

        let future = http
            .serve_incoming(incoming, move || -> Result<_, hyper::Error> {
                Ok(Listen {
                    handle: handle.clone(),
                    handler: handler.clone(),
                })
            })
            .and_then(|connecting| connecting)
            .for_each(|connection| {
                hyper::rt::spawn(connection.map_err(|err| error!("{}", err)));
                Ok(())
            })
            .map(|_| IO::Value(()))
            .map_err(|err| vm::Error::from(format!("Server error: {}", err)));
        return Either::B(Either::A(future));
    }

    Either::B(Either::B(
        Server::bind(&addr)
            .serve(move || -> Result<_, hyper::Error> {
                Ok(Listen {
                    handle: handle.clone(),
                    handler: handler.clone(),
                })
            })
            .map_err(|err| vm::Error::from(format!("Server error: {}", err)))
            .map(|_| IO::Value(())),
    ))
}

// To let the `http_types` module refer to `Body` and `ResponseBody` we register these types in a
// separate function which is called before loading `http_types`
pub fn load_types(vm: &Thread) -> vm::Result<ExternModule> {
    vm.register_type::<Body>("Body", &[])?;
    vm.register_type::<ResponseBody>("ResponseBody", &[])?;
    vm.register_type::<Uri>("Uri", &[])?;

    ExternModule::new(
        vm,
        record! {
            // Define the types so that they can be used from gluon
            type Body => Body,
            type ResponseBody => ResponseBody,
            type Uri => Uri,
            type Method => String,
            type StatusCode => u16,
            type Request => Request,
            type Response => Response,
            type Headers => Headers,
            type HttpState => HttpState
        },
    )
}

macro_rules! uri_binds {
    ($($id: ident)*) => {
        record!{
            $(
                $id => primitive!(1, concat!("std.http.prim.uri.", stringify!($id)), |u: &Uri| (u.0).$id())
            ),*
        }
    }
}

pub fn load(vm: &Thread) -> vm::Result<ExternModule> {
    ExternModule::new(
        vm,
        record! {
            listen => primitive!(2, async fn listen),
            read_chunk => primitive!(1, async fn read_chunk),
            write_response => primitive!(2, async fn write_response),
            port => primitive!(1, "std.http.prim.uri.port", |u: &Uri| (u.0).port_part().map(|p| p.as_u16())),
            uri => uri_binds!(path host query to_string)
        },
    )
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2014-2015 Mathias Hällman
 *
 * 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/.
 */

/*
 * A key can be of any of the following types:
 *   - Function key such as F1 or F24
 *   - Symbolic key such as enter or escape
 *   - Unicode codepoint such as a, ä, or あ
 * The key also records key modifier information.
 */
#[derive(Clone, Copy, Eq, Hash, PartialEq)]
#[cfg_attr(test, derive(Debug))]
pub enum Key {
  Fn{num: isize, mods: KeyMod},
  Sym{sym: KeySym, mods: KeyMod},
  Unicode{codepoint: char, mods: KeyMod},
}

#[derive(Clone, Copy, Eq, Hash, PartialEq)]
#[cfg_attr(test, derive(Debug))]
pub enum KeySym
{
  Unknown = -1,
  None = 0,

  /* Special names in C0 */
  Backspace,
  Tab,
  Enter,
  Escape,

  /* Special names in G0 */
  Space,
  Del,

  /* Special keys */
  Up,
  Down,
  Left,
  Right,
  Begin,
  Find,
  Insert,
  Delete,
  Select,
  Pageup,
  Pagedown,
  Home,
  End,

  /* Special keys from terminfo */
  Cancel,
  Clear,
  Close,
  Command,
  Copy,
  Exit,
  Help,
  Mark,
  Message,
  Move,
  Open,
  Options,
  Print,
  Redo,
  Reference,
  Refresh,
  Replace,
  Restart,
  Resume,
  Save,
  Suspend,
  Undo,

  /* Numeric keypad special keys */
  KP0,
  KP1,
  KP2,
  KP3,
  KP4,
  KP5,
  KP6,
  KP7,
  KP8,
  KP9,
  KPEnter,
  KPPlus,
  KPMinus,
  KPMult,
  KPDiv,
  KPComma,
  KPPeriod,
  KPEquals,

  /* et cetera ad nauseum */
  NSyms,
}

bitflags! {
  pub flags KeyMod: u8 {
    const MOD_NONE  = 0,
    const MOD_SHIFT = 1 << 0,
    const MOD_ALT   = 1 << 1,
    const MOD_CTRL  = 1 << 2,
  }
}

#[allow(dead_code)]  // as the name suggests, this function should not be called
fn totally_useless_function_just_to_suppress_warnings_about_dead_code() {
  let mut a = MOD_NONE;
  a.toggle(MOD_NONE); a.remove(MOD_NONE); a.intersects(MOD_NONE);
  a.is_all(); KeyMod::from_bits_truncate(0); KeyMod::from_bits(0);
  a.insert(MOD_NONE); a.contains(MOD_NONE);
}
<commit_msg>Remove totally useless function<commit_after>/*
 * Copyright (c) 2014-2015 Mathias Hällman
 *
 * 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/.
 */

/*
 * A key can be of any of the following types:
 *   - Function key such as F1 or F24
 *   - Symbolic key such as enter or escape
 *   - Unicode codepoint such as a, ä, or あ
 * The key also records key modifier information.
 */
#[derive(Clone, Copy, Eq, Hash, PartialEq)]
#[cfg_attr(test, derive(Debug))]
pub enum Key {
  Fn{num: isize, mods: KeyMod},
  Sym{sym: KeySym, mods: KeyMod},
  Unicode{codepoint: char, mods: KeyMod},
}

#[derive(Clone, Copy, Eq, Hash, PartialEq)]
#[cfg_attr(test, derive(Debug))]
pub enum KeySym
{
  Unknown = -1,
  None = 0,

  /* Special names in C0 */
  Backspace,
  Tab,
  Enter,
  Escape,

  /* Special names in G0 */
  Space,
  Del,

  /* Special keys */
  Up,
  Down,
  Left,
  Right,
  Begin,
  Find,
  Insert,
  Delete,
  Select,
  Pageup,
  Pagedown,
  Home,
  End,

  /* Special keys from terminfo */
  Cancel,
  Clear,
  Close,
  Command,
  Copy,
  Exit,
  Help,
  Mark,
  Message,
  Move,
  Open,
  Options,
  Print,
  Redo,
  Reference,
  Refresh,
  Replace,
  Restart,
  Resume,
  Save,
  Suspend,
  Undo,

  /* Numeric keypad special keys */
  KP0,
  KP1,
  KP2,
  KP3,
  KP4,
  KP5,
  KP6,
  KP7,
  KP8,
  KP9,
  KPEnter,
  KPPlus,
  KPMinus,
  KPMult,
  KPDiv,
  KPComma,
  KPPeriod,
  KPEquals,

  /* et cetera ad nauseum */
  NSyms,
}

bitflags! {
  pub flags KeyMod: u8 {
    const MOD_NONE  = 0,
    const MOD_SHIFT = 1 << 0,
    const MOD_ALT   = 1 << 1,
    const MOD_CTRL  = 1 << 2,
  }
}
<|endoftext|>


<commit_before>//! # io
//!
//! IO helper functions
//!

#[cfg(test)]
#[path = "io_test.rs"]
mod io_test;

use fsio::file::modify_file;
use fsio::path::from_path::FromPath;
use fsio::path::get_temporary_file_path;
use glob::glob;
use ignore::WalkBuilder;
use std::collections::HashSet;
use std::fs::File;
use std::io;
use std::path::PathBuf;

pub(crate) fn create_text_file(text: &str, extension: &str) -> String {
    let file_path = get_temporary_file_path(extension);

    match fsio::file::write_text_file(&file_path, text) {
        Ok(_) => file_path,
        Err(error) => {
            error!("Unable to create file: {} {:#?}", &file_path, &error);
            panic!("Unable to create file, error: {}", error);
        }
    }
}

pub(crate) fn create_file(
    write_content: &Fn(&mut File) -> io::Result<()>,
    extension: &str,
) -> String {
    let file_path = get_temporary_file_path(extension);

    match modify_file(&file_path, write_content, false) {
        Ok(_) => file_path,
        Err(error) => {
            error!("Unable to write to file: {} {:#?}", &file_path, &error);
            panic!("Unable to write to file, error: {}", error);
        }
    }
}

pub(crate) fn delete_file(file: &str) {
    match fsio::file::delete(file) {
        Ok(_) => debug!("File deleted: {}", &file),
        Err(error) => debug!("Unable to delete file: {} {:#?}", &file, error),
    }
}

pub(crate) fn write_text_file(file_path: &str, text: &str) -> bool {
    match fsio::file::write_text_file(file_path, text) {
        Ok(_) => true,
        Err(_) => false,
    }
}

pub(crate) fn read_text_file(file_path: &PathBuf) -> String {
    debug!("Opening file: {:#?}", &file_path);

    match fsio::file::read_text_file(file_path) {
        Ok(content) => content,
        Err(error) => panic!("Unable to read file: {:?} error: {:#?}", file_path, error),
    }
}

pub(crate) fn get_path_list(
    glob_pattern: &str,
    include_files: bool,
    include_dirs: bool,
    ignore_type: Option<String>,
) -> Vec<String> {
    let mut path_list = vec![];
    match glob(glob_pattern) {
        Ok(paths) => {
            for entry in paths {
                match entry {
                    Ok(path) => {
                        let add =
                            (include_dirs && path.is_dir()) || (include_files && path.is_file());

                        if add {
                            let mut value_string: String = FromPath::from_path(&path);
                            value_string = value_string.replace("\\", "/");
                            if !value_string.starts_with("./") {
                                value_string.insert_str(0, "./");
                            }

                            path_list.push(value_string);
                        }
                    }
                    Err(error) => {
                        error!(
                            "Error while iterating over path entries of glob: {}, error: {:#?}",
                            glob_pattern, error
                        );
                        return vec![];
                    }
                }
            }
        }
        Err(error) => {
            error!(
                "Error while running glob: {}, error: {:#?}",
                glob_pattern, error
            );
            return vec![];
        }
    }

    if !path_list.is_empty() {
        if let Some(ignore_type_value) = ignore_type {
            let mut included_paths = HashSet::new();

            match ignore_type_value.as_str() {
                "git" => {
                    for entry in WalkBuilder::new("./")
                        .hidden(true)
                        .parents(true)
                        .git_ignore(true)
                        .git_exclude(true)
                        .build()
                    {
                        match entry {
                            Ok(path) => {
                                let mut value_string: String = FromPath::from_path(&path.path());
                                value_string = value_string.replace("\\", "/");
                                included_paths.insert(value_string);
                            }
                            Err(error) => error!(
                                "Error while running git ignore path checks, error: {:#?}",
                                error
                            ),
                        }
                    }
                }
                _ => error!("Unsupported ignore type: {}", &ignore_type_value),
            };

            println!("path_list: {:#?}", &path_list); //todo re
            println!("included_paths: {:#?}", &included_paths); //todo re
            if included_paths.is_empty() {
                path_list.clear();
            } else {
                let org_path_list = path_list;
                path_list = vec![];
                for path in org_path_list {
                    if included_paths.contains(&path) {
                        path_list.push(path);
                    }
                }
            }
        }
    }

    path_list
}
<commit_msg>Support path glob (with optional gitignore support) environment variables #542<commit_after>//! # io
//!
//! IO helper functions
//!

#[cfg(test)]
#[path = "io_test.rs"]
mod io_test;

use fsio::file::modify_file;
use fsio::path::from_path::FromPath;
use fsio::path::get_temporary_file_path;
use glob::glob;
use ignore::WalkBuilder;
use std::collections::HashSet;
use std::fs::File;
use std::io;
use std::path::PathBuf;

pub(crate) fn create_text_file(text: &str, extension: &str) -> String {
    let file_path = get_temporary_file_path(extension);

    match fsio::file::write_text_file(&file_path, text) {
        Ok(_) => file_path,
        Err(error) => {
            error!("Unable to create file: {} {:#?}", &file_path, &error);
            panic!("Unable to create file, error: {}", error);
        }
    }
}

pub(crate) fn create_file(
    write_content: &Fn(&mut File) -> io::Result<()>,
    extension: &str,
) -> String {
    let file_path = get_temporary_file_path(extension);

    match modify_file(&file_path, write_content, false) {
        Ok(_) => file_path,
        Err(error) => {
            error!("Unable to write to file: {} {:#?}", &file_path, &error);
            panic!("Unable to write to file, error: {}", error);
        }
    }
}

pub(crate) fn delete_file(file: &str) {
    match fsio::file::delete(file) {
        Ok(_) => debug!("File deleted: {}", &file),
        Err(error) => debug!("Unable to delete file: {} {:#?}", &file, error),
    }
}

pub(crate) fn write_text_file(file_path: &str, text: &str) -> bool {
    match fsio::file::write_text_file(file_path, text) {
        Ok(_) => true,
        Err(_) => false,
    }
}

pub(crate) fn read_text_file(file_path: &PathBuf) -> String {
    debug!("Opening file: {:#?}", &file_path);

    match fsio::file::read_text_file(file_path) {
        Ok(content) => content,
        Err(error) => panic!("Unable to read file: {:?} error: {:#?}", file_path, error),
    }
}

pub(crate) fn get_path_list(
    glob_pattern: &str,
    include_files: bool,
    include_dirs: bool,
    ignore_type: Option<String>,
) -> Vec<String> {
    let mut path_list = vec![];
    match glob(glob_pattern) {
        Ok(paths) => {
            for entry in paths {
                match entry {
                    Ok(path) => {
                        if (include_dirs && path.is_dir()) || (include_files && path.is_file()) {
                            let mut value_string: String = FromPath::from_path(&path);
                            value_string = value_string.replace("\\", "/");
                            if !value_string.starts_with("./") {
                                value_string.insert_str(0, "./");
                            }

                            path_list.push(value_string);
                        }
                    }
                    Err(error) => {
                        error!(
                            "Error while iterating over path entries of glob: {}, error: {:#?}",
                            glob_pattern, error
                        );
                        return vec![];
                    }
                }
            }
        }
        Err(error) => {
            error!(
                "Error while running glob: {}, error: {:#?}",
                glob_pattern, error
            );
            return vec![];
        }
    }

    if !path_list.is_empty() {
        if let Some(ignore_type_value) = ignore_type {
            let mut included_paths = HashSet::new();

            match ignore_type_value.as_str() {
                "git" => {
                    for entry in WalkBuilder::new("./")
                        .hidden(true)
                        .parents(true)
                        .git_ignore(true)
                        .git_exclude(true)
                        .build()
                    {
                        match entry {
                            Ok(path) => {
                                let mut value_string: String = FromPath::from_path(&path.path());
                                value_string = value_string.replace("\\", "/");
                                included_paths.insert(value_string);
                            }
                            Err(error) => error!(
                                "Error while running git ignore path checks, error: {:#?}",
                                error
                            ),
                        }
                    }
                }
                _ => error!("Unsupported ignore type: {}", &ignore_type_value),
            };

            println!("path_list: {:#?}", &path_list); //todo re
            println!("included_paths: {:#?}", &included_paths); //todo re
            if included_paths.is_empty() {
                path_list.clear();
            } else {
                let org_path_list = path_list;
                path_list = vec![];
                for path in org_path_list {
                    if included_paths.contains(&path) {
                        path_list.push(path);
                    }
                }
            }
        }
    }

    path_list
}
<|endoftext|>


<commit_before>use std::cell::Ref;

use crate::dev::{AppConfig, RequestHead};
use crate::http::header::{self, HeaderName};

const X_FORWARDED_FOR: &[u8] = b"x-forwarded-for";
const X_FORWARDED_HOST: &[u8] = b"x-forwarded-host";
const X_FORWARDED_PROTO: &[u8] = b"x-forwarded-proto";

/// `HttpRequest` connection information
#[derive(Debug, Clone, Default)]
pub struct ConnectionInfo {
    scheme: String,
    host: String,
    remote: Option<String>,
    peer: Option<String>,
}

impl ConnectionInfo {
    /// Create *ConnectionInfo* instance for a request.
    pub fn get<'a>(req: &'a RequestHead, cfg: &AppConfig) -> Ref<'a, Self> {
        if !req.extensions().contains::<ConnectionInfo>() {
            req.extensions_mut().insert(ConnectionInfo::new(req, cfg));
        }
        Ref::map(req.extensions(), |e| e.get().unwrap())
    }

    #[allow(clippy::cognitive_complexity)]
    fn new(req: &RequestHead, cfg: &AppConfig) -> ConnectionInfo {
        let mut host = None;
        let mut scheme = None;
        let mut remote = None;
        let mut peer = None;

        // load forwarded header
        for hdr in req.headers.get_all(&header::FORWARDED) {
            if let Ok(val) = hdr.to_str() {
                for pair in val.split(';') {
                    for el in pair.split(',') {
                        let mut items = el.trim().splitn(2, '=');
                        if let Some(name) = items.next() {
                            if let Some(val) = items.next() {
                                match &name.to_lowercase() as &str {
                                    "for" => {
                                        if remote.is_none() {
                                            remote = Some(val.trim());
                                        }
                                    }
                                    "proto" => {
                                        if scheme.is_none() {
                                            scheme = Some(val.trim());
                                        }
                                    }
                                    "host" => {
                                        if host.is_none() {
                                            host = Some(val.trim());
                                        }
                                    }
                                    _ => (),
                                }
                            }
                        }
                    }
                }
            }
        }

        // scheme
        if scheme.is_none() {
            if let Some(h) = req
                .headers
                .get(&HeaderName::from_lowercase(X_FORWARDED_PROTO).unwrap())
            {
                if let Ok(h) = h.to_str() {
                    scheme = h.split(',').next().map(|v| v.trim());
                }
            }
            if scheme.is_none() {
                scheme = req.uri.scheme_part().map(|a| a.as_str());
                if scheme.is_none() && cfg.secure() {
                    scheme = Some("https")
                }
            }
        }

        // host
        if host.is_none() {
            if let Some(h) = req
                .headers
                .get(&HeaderName::from_lowercase(X_FORWARDED_HOST).unwrap())
            {
                if let Ok(h) = h.to_str() {
                    host = h.split(',').next().map(|v| v.trim());
                }
            }
            if host.is_none() {
                if let Some(h) = req.headers.get(&header::HOST) {
                    host = h.to_str().ok();
                }
                if host.is_none() {
                    host = req.uri.authority_part().map(|a| a.as_str());
                    if host.is_none() {
                        host = Some(cfg.host());
                    }
                }
            }
        }

        // remote addr
        if remote.is_none() {
            if let Some(h) = req
                .headers
                .get(&HeaderName::from_lowercase(X_FORWARDED_FOR).unwrap())
            {
                if let Ok(h) = h.to_str() {
                    remote = h.split(',').next().map(|v| v.trim());
                }
            }
            if remote.is_none() {
                // get peeraddr from socketaddr
                peer = req.peer_addr.map(|addr| format!("{}", addr));
            }
        }

        ConnectionInfo {
            peer,
            scheme: scheme.unwrap_or("http").to_owned(),
            host: host.unwrap_or("localhost").to_owned(),
            remote: remote.map(|s| s.to_owned()),
        }
    }

    /// Scheme of the request.
    ///
    /// Scheme is resolved through the following headers, in this order:
    ///
    /// - Forwarded
    /// - X-Forwarded-Proto
    /// - Uri
    #[inline]
    pub fn scheme(&self) -> &str {
        &self.scheme
    }

    /// Hostname of the request.
    ///
    /// Hostname is resolved through the following headers, in this order:
    ///
    /// - Forwarded
    /// - X-Forwarded-Host
    /// - Host
    /// - Uri
    /// - Server hostname
    pub fn host(&self) -> &str {
        &self.host
    }

    /// Remote socket addr of client initiated HTTP request.
    ///
    /// The addr is resolved through the following headers, in this order:
    ///
    /// - Forwarded
    /// - X-Forwarded-For
    /// - peer name of opened socket
    #[inline]
    pub fn remote(&self) -> Option<&str> {
        if let Some(ref r) = self.remote {
            Some(r)
        } else if let Some(ref peer) = self.peer {
            Some(peer)
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test::TestRequest;

    #[test]
    fn test_forwarded() {
        let req = TestRequest::default().to_http_request();
        let info = req.connection_info();
        assert_eq!(info.scheme(), "http");
        assert_eq!(info.host(), "localhost:8080");

        let req = TestRequest::default()
            .header(
                header::FORWARDED,
                "for=192.0.2.60; proto=https; by=203.0.113.43; host=rust-lang.org",
            )
            .to_http_request();

        let info = req.connection_info();
        assert_eq!(info.scheme(), "https");
        assert_eq!(info.host(), "rust-lang.org");
        assert_eq!(info.remote(), Some("192.0.2.60"));

        let req = TestRequest::default()
            .header(header::HOST, "rust-lang.org")
            .to_http_request();

        let info = req.connection_info();
        assert_eq!(info.scheme(), "http");
        assert_eq!(info.host(), "rust-lang.org");
        assert_eq!(info.remote(), None);

        let req = TestRequest::default()
            .header(X_FORWARDED_FOR, "192.0.2.60")
            .to_http_request();
        let info = req.connection_info();
        assert_eq!(info.remote(), Some("192.0.2.60"));

        let req = TestRequest::default()
            .header(X_FORWARDED_HOST, "192.0.2.60")
            .to_http_request();
        let info = req.connection_info();
        assert_eq!(info.host(), "192.0.2.60");
        assert_eq!(info.remote(), None);

        let req = TestRequest::default()
            .header(X_FORWARDED_PROTO, "https")
            .to_http_request();
        let info = req.connection_info();
        assert_eq!(info.scheme(), "https");
    }
}
<commit_msg>Add security note to ConnectionInfo::remote() (#1158)<commit_after>use std::cell::Ref;

use crate::dev::{AppConfig, RequestHead};
use crate::http::header::{self, HeaderName};

const X_FORWARDED_FOR: &[u8] = b"x-forwarded-for";
const X_FORWARDED_HOST: &[u8] = b"x-forwarded-host";
const X_FORWARDED_PROTO: &[u8] = b"x-forwarded-proto";

/// `HttpRequest` connection information
#[derive(Debug, Clone, Default)]
pub struct ConnectionInfo {
    scheme: String,
    host: String,
    remote: Option<String>,
    peer: Option<String>,
}

impl ConnectionInfo {
    /// Create *ConnectionInfo* instance for a request.
    pub fn get<'a>(req: &'a RequestHead, cfg: &AppConfig) -> Ref<'a, Self> {
        if !req.extensions().contains::<ConnectionInfo>() {
            req.extensions_mut().insert(ConnectionInfo::new(req, cfg));
        }
        Ref::map(req.extensions(), |e| e.get().unwrap())
    }

    #[allow(clippy::cognitive_complexity)]
    fn new(req: &RequestHead, cfg: &AppConfig) -> ConnectionInfo {
        let mut host = None;
        let mut scheme = None;
        let mut remote = None;
        let mut peer = None;

        // load forwarded header
        for hdr in req.headers.get_all(&header::FORWARDED) {
            if let Ok(val) = hdr.to_str() {
                for pair in val.split(';') {
                    for el in pair.split(',') {
                        let mut items = el.trim().splitn(2, '=');
                        if let Some(name) = items.next() {
                            if let Some(val) = items.next() {
                                match &name.to_lowercase() as &str {
                                    "for" => {
                                        if remote.is_none() {
                                            remote = Some(val.trim());
                                        }
                                    }
                                    "proto" => {
                                        if scheme.is_none() {
                                            scheme = Some(val.trim());
                                        }
                                    }
                                    "host" => {
                                        if host.is_none() {
                                            host = Some(val.trim());
                                        }
                                    }
                                    _ => (),
                                }
                            }
                        }
                    }
                }
            }
        }

        // scheme
        if scheme.is_none() {
            if let Some(h) = req
                .headers
                .get(&HeaderName::from_lowercase(X_FORWARDED_PROTO).unwrap())
            {
                if let Ok(h) = h.to_str() {
                    scheme = h.split(',').next().map(|v| v.trim());
                }
            }
            if scheme.is_none() {
                scheme = req.uri.scheme_part().map(|a| a.as_str());
                if scheme.is_none() && cfg.secure() {
                    scheme = Some("https")
                }
            }
        }

        // host
        if host.is_none() {
            if let Some(h) = req
                .headers
                .get(&HeaderName::from_lowercase(X_FORWARDED_HOST).unwrap())
            {
                if let Ok(h) = h.to_str() {
                    host = h.split(',').next().map(|v| v.trim());
                }
            }
            if host.is_none() {
                if let Some(h) = req.headers.get(&header::HOST) {
                    host = h.to_str().ok();
                }
                if host.is_none() {
                    host = req.uri.authority_part().map(|a| a.as_str());
                    if host.is_none() {
                        host = Some(cfg.host());
                    }
                }
            }
        }

        // remote addr
        if remote.is_none() {
            if let Some(h) = req
                .headers
                .get(&HeaderName::from_lowercase(X_FORWARDED_FOR).unwrap())
            {
                if let Ok(h) = h.to_str() {
                    remote = h.split(',').next().map(|v| v.trim());
                }
            }
            if remote.is_none() {
                // get peeraddr from socketaddr
                peer = req.peer_addr.map(|addr| format!("{}", addr));
            }
        }

        ConnectionInfo {
            peer,
            scheme: scheme.unwrap_or("http").to_owned(),
            host: host.unwrap_or("localhost").to_owned(),
            remote: remote.map(|s| s.to_owned()),
        }
    }

    /// Scheme of the request.
    ///
    /// Scheme is resolved through the following headers, in this order:
    ///
    /// - Forwarded
    /// - X-Forwarded-Proto
    /// - Uri
    #[inline]
    pub fn scheme(&self) -> &str {
        &self.scheme
    }

    /// Hostname of the request.
    ///
    /// Hostname is resolved through the following headers, in this order:
    ///
    /// - Forwarded
    /// - X-Forwarded-Host
    /// - Host
    /// - Uri
    /// - Server hostname
    pub fn host(&self) -> &str {
        &self.host
    }

    /// Remote socket addr of client initiated HTTP request.
    ///
    /// The addr is resolved through the following headers, in this order:
    ///
    /// - Forwarded
    /// - X-Forwarded-For
    /// - peer name of opened socket
    ///
    /// # Security
    /// Do not use this function for security purposes, unless you can ensure the Forwarded and
    /// X-Forwarded-For headers cannot be spoofed by the client. If you want the client's socket
    /// address explicitly, use
    /// [`HttpRequest::peer_addr()`](../web/struct.HttpRequest.html#method.peer_addr) instead.
    #[inline]
    pub fn remote(&self) -> Option<&str> {
        if let Some(ref r) = self.remote {
            Some(r)
        } else if let Some(ref peer) = self.peer {
            Some(peer)
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test::TestRequest;

    #[test]
    fn test_forwarded() {
        let req = TestRequest::default().to_http_request();
        let info = req.connection_info();
        assert_eq!(info.scheme(), "http");
        assert_eq!(info.host(), "localhost:8080");

        let req = TestRequest::default()
            .header(
                header::FORWARDED,
                "for=192.0.2.60; proto=https; by=203.0.113.43; host=rust-lang.org",
            )
            .to_http_request();

        let info = req.connection_info();
        assert_eq!(info.scheme(), "https");
        assert_eq!(info.host(), "rust-lang.org");
        assert_eq!(info.remote(), Some("192.0.2.60"));

        let req = TestRequest::default()
            .header(header::HOST, "rust-lang.org")
            .to_http_request();

        let info = req.connection_info();
        assert_eq!(info.scheme(), "http");
        assert_eq!(info.host(), "rust-lang.org");
        assert_eq!(info.remote(), None);

        let req = TestRequest::default()
            .header(X_FORWARDED_FOR, "192.0.2.60")
            .to_http_request();
        let info = req.connection_info();
        assert_eq!(info.remote(), Some("192.0.2.60"));

        let req = TestRequest::default()
            .header(X_FORWARDED_HOST, "192.0.2.60")
            .to_http_request();
        let info = req.connection_info();
        assert_eq!(info.host(), "192.0.2.60");
        assert_eq!(info.remote(), None);

        let req = TestRequest::default()
            .header(X_FORWARDED_PROTO, "https")
            .to_http_request();
        let info = req.connection_info();
        assert_eq!(info.scheme(), "https");
    }
}
<|endoftext|>