<commit_before>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! MemoryManager service shell test commands

extern crate alloc;
use crate::mstats;
use crate::CmdFn;
use crate::CommandError;
use crate::HashMap;
use alloc::vec;
use core::fmt::Write;

use kata_io as io;
use kata_memory_interface::*;
use kata_os_common::sel4_sys;

use sel4_sys::seL4_CPtr;
use sel4_sys::seL4_MinSchedContextBits;
use sel4_sys::seL4_ObjectType::*;
use sel4_sys::seL4_WordBits;

pub fn add_cmds(cmds: &mut HashMap<&str, CmdFn>) {
    cmds.extend([
        ("test_malloc", malloc_command as CmdFn),
        ("test_mfree", mfree_command as CmdFn),
        ("test_obj_alloc", obj_alloc_command as CmdFn),
    ]);
}

fn malloc_command(
    args: &mut dyn Iterator<Item = &str>,
    _input: &mut dyn io::BufRead,
    output: &mut dyn io::Write,
    _builtin_cpio: &[u8],
) -> Result<(), CommandError> {
    let space_str = args.next().ok_or(CommandError::BadArgs)?;
    let space_bytes = space_str.parse::<usize>()?;
    match kata_frame_alloc(space_bytes) {
        Ok(frames) => {
            writeln!(output, "Allocated {:?}", frames)?;
        }
        Err(status) => {
            writeln!(output, "malloc failed: {:?}", status)?;
        }
    }
    Ok(())
}

fn mfree_command(
    args: &mut dyn Iterator<Item = &str>,
    _input: &mut dyn io::BufRead,
    output: &mut dyn io::Write,
    _builtin_cpio: &[u8],
) -> Result<(), CommandError> {
    extern "C" {
        static SELF_CNODE: seL4_CPtr;
    }
    let cptr_str = args.next().ok_or(CommandError::BadArgs)?;
    let count_str = args.next().ok_or(CommandError::BadArgs)?;
    let frames = ObjDescBundle::new(
        unsafe { SELF_CNODE },
        seL4_WordBits as u8,
        vec![ObjDesc::new(
            sel4_sys::seL4_RISCV_4K_Page,
            count_str.parse::<usize>()?,
            cptr_str.parse::<usize>()? as seL4_CPtr,
        )],
    );
    match kata_object_free_toplevel(&frames) {
        Ok(_) => {
            writeln!(output, "Free'd {:?}", frames)?;
        }
        Err(status) => {
            writeln!(output, "mfree failed: {:?}", status)?;
        }
    }
    Ok(())
}

fn obj_alloc_command(
    _args: &mut dyn Iterator<Item = &str>,
    _input: &mut dyn io::BufRead,
    output: &mut dyn io::Write,
    _builtin_cpio: &[u8],
) -> Result<(), CommandError> {
    let before_stats = kata_memory_stats().expect("before stats");
    mstats(output, &before_stats)?;

    fn check_alloc(
        output: &mut dyn io::Write,
        name: &str,
        res: Result<ObjDescBundle, MemoryManagerError>,
    ) {
        match res {
            Ok(obj) => {
                if let Err(e) = kata_object_free_toplevel(&obj) {
                    let _ = writeln!(output, "free {} {:?} failed: {:?}", name, obj, e);
                }
            }
            Err(e) => {
                let _ = writeln!(output, "alloc {} failed: {:?}", name, e);
            }
        }
    }

    // NB: alloc+free immediately so we don't run out of top-level CNode slots
    check_alloc(output, "untyped", kata_untyped_alloc(12)); // NB: 4KB
    check_alloc(output, "tcb", kata_tcb_alloc());
    check_alloc(output, "endpoint", kata_endpoint_alloc());
    check_alloc(output, "notification", kata_notification_alloc());
    check_alloc(output, "cnode", kata_cnode_alloc(5)); // NB: 32 slots
    check_alloc(output, "frame", kata_frame_alloc(4096));
    //    check_alloc(output, "large frame",  kata_frame_alloc(1024*1024));
    check_alloc(output, "page table", kata_page_table_alloc());

    #[cfg(feature = "CONFIG_KERNEL_MCS")]
    check_alloc(
        output,
        "sched context",
        kata_sched_context_alloc(seL4_MinSchedContextBits),
    );

    #[cfg(feature = "CONFIG_KERNEL_MCS")]
    check_alloc(output, "reply", kata_reply_alloc());

    let after_stats = kata_memory_stats().expect("after stats");
    mstats(output, &after_stats)?;
    assert_eq!(before_stats.allocated_bytes, after_stats.allocated_bytes);
    assert_eq!(before_stats.free_bytes, after_stats.free_bytes);

    // Batch allocate into a private CNode as we might to build a process.
    const CNODE_DEPTH: usize = 7; // 128 slots
    let cnode = kata_cnode_alloc(CNODE_DEPTH).unwrap(); // XXX handle error
    let objs = ObjDescBundle::new(
        cnode.objs[0].cptr,
        CNODE_DEPTH as u8,
        vec![
            ObjDesc::new(seL4_TCBObject, 1, 0),      // 1 tcb
            ObjDesc::new(seL4_EndpointObject, 2, 1), // 2 endpoiints
            ObjDesc::new(seL4_ReplyObject, 2, 3),    // 2 replys
            ObjDesc::new(
                seL4_SchedContextObject, // 1 sched context
                seL4_MinSchedContextBits,
                5,
            ),
            ObjDesc::new(seL4_RISCV_4K_Page, 10, 6), // 10 4K pages
        ],
    );
    match kata_object_alloc(&objs) {
        Ok(_) => {
            writeln!(output, "Batch alloc ok: {:?}", objs)?;
            if let Err(e) = kata_object_free(&objs) {
                writeln!(output, "Batch free err: {:?}", e)?;
            }
        }
        Err(e) => {
            writeln!(output, "Batch alloc err: {:?} {:?}", objs, e)?;
        }
    }
    if let Err(e) = kata_object_free_toplevel(&cnode) {
        writeln!(output, "Cnode free err: {:?} {:?}", cnode, e)?;
    }

    // Batch allocate using the newer api that constructs a CNode based
    // on the batch of objects specified.
    match kata_object_alloc_in_cnode(vec![
        ObjDesc::new(seL4_TCBObject, 1, 0),      // 1 tcb
        ObjDesc::new(seL4_EndpointObject, 1, 1), // 1 endpoiints
        ObjDesc::new(seL4_ReplyObject, 1, 2),    // 1 replys
        ObjDesc::new(
            seL4_SchedContextObject, // 1 sched context
            seL4_MinSchedContextBits,
            3,
        ),
        ObjDesc::new(seL4_RISCV_4K_Page, 2, 4), // 2 4K pages
    ]) {
        Ok(objs) => {
            writeln!(output, "kata_object_alloc_in_cnode ok: {:?}", objs)?;
            if let Err(e) = kata_object_free_in_cnode(&objs) {
                writeln!(output, "kata_object_free_in_cnode failed: {:?}", e)?;
            }
        }
        Err(e) => {
            writeln!(output, "kata_object_alloc_in_cnode failed: {:?}", e)?;
        }
    }

    Ok(writeln!(output, "All tests passed!")?)
}
<commit_msg>kata-shell: improve portability of memory manager test cmds<commit_after>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! MemoryManager service shell test commands

extern crate alloc;
use crate::mstats;
use crate::CmdFn;
use crate::CommandError;
use crate::HashMap;
use alloc::vec;
use core::fmt::Write;

use kata_io as io;
use kata_memory_interface::*;
use kata_os_common::sel4_sys;

use sel4_sys::seL4_CPtr;
use sel4_sys::seL4_MinSchedContextBits;
use sel4_sys::seL4_ObjectType::*;
use sel4_sys::seL4_SmallPageObject;
use sel4_sys::seL4_WordBits;

pub fn add_cmds(cmds: &mut HashMap<&str, CmdFn>) {
    cmds.extend([
        ("test_malloc", malloc_command as CmdFn),
        ("test_mfree", mfree_command as CmdFn),
        ("test_obj_alloc", obj_alloc_command as CmdFn),
    ]);
}

fn malloc_command(
    args: &mut dyn Iterator<Item = &str>,
    _input: &mut dyn io::BufRead,
    output: &mut dyn io::Write,
    _builtin_cpio: &[u8],
) -> Result<(), CommandError> {
    let space_str = args.next().ok_or(CommandError::BadArgs)?;
    let space_bytes = space_str.parse::<usize>()?;
    match kata_frame_alloc(space_bytes) {
        Ok(frames) => {
            writeln!(output, "Allocated {:?}", frames)?;
        }
        Err(status) => {
            writeln!(output, "malloc failed: {:?}", status)?;
        }
    }
    Ok(())
}

fn mfree_command(
    args: &mut dyn Iterator<Item = &str>,
    _input: &mut dyn io::BufRead,
    output: &mut dyn io::Write,
    _builtin_cpio: &[u8],
) -> Result<(), CommandError> {
    extern "C" {
        static SELF_CNODE: seL4_CPtr;
    }
    let cptr_str = args.next().ok_or(CommandError::BadArgs)?;
    let count_str = args.next().ok_or(CommandError::BadArgs)?;
    let frames = ObjDescBundle::new(
        unsafe { SELF_CNODE },
        seL4_WordBits as u8,
        vec![ObjDesc::new(
            seL4_SmallPageObject,
            count_str.parse::<usize>()?,
            cptr_str.parse::<usize>()? as seL4_CPtr,
        )],
    );
    match kata_object_free_toplevel(&frames) {
        Ok(_) => {
            writeln!(output, "Free'd {:?}", frames)?;
        }
        Err(status) => {
            writeln!(output, "mfree failed: {:?}", status)?;
        }
    }
    Ok(())
}

fn obj_alloc_command(
    _args: &mut dyn Iterator<Item = &str>,
    _input: &mut dyn io::BufRead,
    output: &mut dyn io::Write,
    _builtin_cpio: &[u8],
) -> Result<(), CommandError> {
    let before_stats = kata_memory_stats().expect("before stats");
    mstats(output, &before_stats)?;

    fn check_alloc(
        output: &mut dyn io::Write,
        name: &str,
        res: Result<ObjDescBundle, MemoryManagerError>,
    ) {
        match res {
            Ok(obj) => {
                if let Err(e) = kata_object_free_toplevel(&obj) {
                    let _ = writeln!(output, "free {} {:?} failed: {:?}", name, obj, e);
                }
            }
            Err(e) => {
                let _ = writeln!(output, "alloc {} failed: {:?}", name, e);
            }
        }
    }

    // NB: alloc+free immediately so we don't run out of top-level CNode slots
    check_alloc(output, "untyped", kata_untyped_alloc(12)); // NB: 4KB
    check_alloc(output, "tcb", kata_tcb_alloc());
    check_alloc(output, "endpoint", kata_endpoint_alloc());
    check_alloc(output, "notification", kata_notification_alloc());
    check_alloc(output, "cnode", kata_cnode_alloc(5)); // NB: 32 slots
    check_alloc(output, "frame", kata_frame_alloc(4096));
    //    check_alloc(output, "large frame",  kata_frame_alloc(1024*1024));
    check_alloc(output, "page table", kata_page_table_alloc());

    #[cfg(feature = "CONFIG_KERNEL_MCS")]
    check_alloc(
        output,
        "sched context",
        kata_sched_context_alloc(seL4_MinSchedContextBits),
    );

    #[cfg(feature = "CONFIG_KERNEL_MCS")]
    check_alloc(output, "reply", kata_reply_alloc());

    let after_stats = kata_memory_stats().expect("after stats");
    mstats(output, &after_stats)?;
    assert_eq!(before_stats.allocated_bytes, after_stats.allocated_bytes);
    assert_eq!(before_stats.free_bytes, after_stats.free_bytes);

    // Batch allocate into a private CNode as we might to build a process.
    const CNODE_DEPTH: usize = 7; // 128 slots
    let cnode = kata_cnode_alloc(CNODE_DEPTH).unwrap(); // XXX handle error
    let objs = ObjDescBundle::new(
        cnode.objs[0].cptr,
        CNODE_DEPTH as u8,
        vec![
            ObjDesc::new(seL4_TCBObject, 1, 0),      // 1 tcb
            ObjDesc::new(seL4_EndpointObject, 2, 1), // 2 endpoiints
            ObjDesc::new(seL4_ReplyObject, 2, 3),    // 2 replys
            ObjDesc::new(
                seL4_SchedContextObject, // 1 sched context
                seL4_MinSchedContextBits,
                5,
            ),
            ObjDesc::new(seL4_SmallPageObject, 10, 6), // 10 4K pages
        ],
    );
    match kata_object_alloc(&objs) {
        Ok(_) => {
            writeln!(output, "Batch alloc ok: {:?}", objs)?;
            if let Err(e) = kata_object_free(&objs) {
                writeln!(output, "Batch free err: {:?}", e)?;
            }
        }
        Err(e) => {
            writeln!(output, "Batch alloc err: {:?} {:?}", objs, e)?;
        }
    }
    if let Err(e) = kata_object_free_toplevel(&cnode) {
        writeln!(output, "Cnode free err: {:?} {:?}", cnode, e)?;
    }

    // Batch allocate using the newer api that constructs a CNode based
    // on the batch of objects specified.
    match kata_object_alloc_in_cnode(vec![
        ObjDesc::new(seL4_TCBObject, 1, 0),      // 1 tcb
        ObjDesc::new(seL4_EndpointObject, 1, 1), // 1 endpoiints
        ObjDesc::new(seL4_ReplyObject, 1, 2),    // 1 replys
        ObjDesc::new(
            seL4_SchedContextObject, // 1 sched context
            seL4_MinSchedContextBits,
            3,
        ),
        ObjDesc::new(seL4_SmallPageObject, 2, 4), // 2 4K pages
    ]) {
        Ok(objs) => {
            writeln!(output, "kata_object_alloc_in_cnode ok: {:?}", objs)?;
            if let Err(e) = kata_object_free_in_cnode(&objs) {
                writeln!(output, "kata_object_free_in_cnode failed: {:?}", e)?;
            }
        }
        Err(e) => {
            writeln!(output, "kata_object_alloc_in_cnode failed: {:?}", e)?;
        }
    }

    Ok(writeln!(output, "All tests passed!")?)
}
<|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/. */

use dom::bindings::codegen::Bindings::FormDataBinding::FormDataMethods;
use dom::bindings::error::{Error, Fallible};
use dom::bindings::reflector::DomObject;
use dom::bindings::root::DomRoot;
use dom::bindings::str::USVString;
use dom::bindings::trace::RootedTraceableBox;
use dom::blob::{Blob, BlobImpl};
use dom::formdata::FormData;
use dom::globalscope::GlobalScope;
use dom::promise::Promise;
use js::jsapi::Heap;
use js::jsval::JSVal;
use js::jsapi::JSContext;
use js::jsapi::JSObject;
use js::jsapi::JS_ClearPendingException;
use js::jsapi::JS_ParseJSON;
use js::jsapi::Value as JSValue;
use js::jsval::UndefinedValue;
use js::typedarray::{ArrayBuffer, CreateWith};
use js::jsapi::JS_GetPendingException;
use mime::{Mime, TopLevel, SubLevel};
use std::cell::Ref;
use std::ptr;
use std::rc::Rc;
use std::str;
use url::form_urlencoded;

#[derive(Clone, Copy, JSTraceable, MallocSizeOf)]
pub enum BodyType {
    Blob,
    FormData,
    Json,
    Text,
    ArrayBuffer
}

pub enum FetchedData {
    Text(String),
    Json(RootedTraceableBox<Heap<JSValue>>),
    BlobData(DomRoot<Blob>),
    FormData(DomRoot<FormData>),
    ArrayBuffer(RootedTraceableBox<Heap<*mut JSObject>>),
    JSException(RootedTraceableBox<Heap<JSVal>>)
}

// https://fetch.spec.whatwg.org/#concept-body-consume-body
#[allow(unrooted_must_root)]
pub fn consume_body<T: BodyOperations + DomObject>(object: &T, body_type: BodyType) -> Rc<Promise> {
    let promise = Promise::new(&object.global());

    // Step 1
    if object.get_body_used() || object.is_locked() {
        promise.reject_error(Error::Type(
            "The response's stream is disturbed or locked".to_string(),
        ));
        return promise;
    }

    object.set_body_promise(&promise, body_type);

    // Steps 2-4
    // TODO: Body does not yet have a stream.

    consume_body_with_promise(object, body_type, &promise);

    promise
}

// https://fetch.spec.whatwg.org/#concept-body-consume-body
#[allow(unrooted_must_root)]
pub fn consume_body_with_promise<T: BodyOperations + DomObject>(object: &T,
                                                                body_type: BodyType,
                                                                promise: &Promise) {
    // Step 5
    let body = match object.take_body() {
        Some(body) => body,
        None => return,
    };

    let pkg_data_results = run_package_data_algorithm(object,
                                                      body,
                                                      body_type,
                                                      object.get_mime_type());

    match pkg_data_results {
        Ok(results) => {
            match results {
                FetchedData::Text(s) => promise.resolve_native(&USVString(s)),
                FetchedData::Json(j) => promise.resolve_native(&j),
                FetchedData::BlobData(b) => promise.resolve_native(&b),
                FetchedData::FormData(f) => promise.resolve_native(&f),
                FetchedData::ArrayBuffer(a) => promise.resolve_native(&a),
                FetchedData::JSException(e) => promise.reject_native(&e.handle()),
            };
        },
        Err(err) => promise.reject_error(err),
    }
}

// https://fetch.spec.whatwg.org/#concept-body-package-data
#[allow(unsafe_code)]
fn run_package_data_algorithm<T: BodyOperations + DomObject>(object: &T,
                                                             bytes: Vec<u8>,
                                                             body_type: BodyType,
                                                             mime_type: Ref<Vec<u8>>)
                                                             -> Fallible<FetchedData> {
    let global = object.global();
    let cx = global.get_cx();
    let mime = &*mime_type;
    match body_type {
        BodyType::Text => run_text_data_algorithm(bytes),
        BodyType::Json => run_json_data_algorithm(cx, bytes),
        BodyType::Blob => run_blob_data_algorithm(&global, bytes, mime),
        BodyType::FormData => run_form_data_algorithm(&global, bytes, mime),
        BodyType::ArrayBuffer => unsafe {
            run_array_buffer_data_algorithm(cx, bytes)
        }
    }
}

fn run_text_data_algorithm(bytes: Vec<u8>) -> Fallible<FetchedData> {
    Ok(FetchedData::Text(String::from_utf8_lossy(&bytes).into_owned()))
}

#[allow(unsafe_code)]
fn run_json_data_algorithm(cx: *mut JSContext,
                           bytes: Vec<u8>) -> Fallible<FetchedData> {
    let json_text = String::from_utf8_lossy(&bytes);
    let json_text: Vec<u16> = json_text.encode_utf16().collect();
    rooted!(in(cx) let mut rval = UndefinedValue());
    unsafe {
        if !JS_ParseJSON(cx,
                         json_text.as_ptr(),
                         json_text.len() as u32,
                         rval.handle_mut()) {
            rooted!(in(cx) let mut exception = UndefinedValue());
            JS_GetPendingException(cx, exception.handle_mut());
            JS_ClearPendingException(cx);
            return Ok(FetchedData::JSException(RootedTraceableBox::from_box(Heap::boxed(exception.get()))));
        }
        let rooted_heap = RootedTraceableBox::from_box(Heap::boxed(rval.get()));
        Ok(FetchedData::Json(rooted_heap))
    }
}

fn run_blob_data_algorithm(root: &GlobalScope,
                           bytes: Vec<u8>,
                           mime: &[u8]) -> Fallible<FetchedData> {
    let mime_string = if let Ok(s) = String::from_utf8(mime.to_vec()) {
        s
    } else {
        "".to_string()
    };
    let blob = Blob::new(root, BlobImpl::new_from_bytes(bytes), mime_string);
    Ok(FetchedData::BlobData(blob))
}

fn run_form_data_algorithm(root: &GlobalScope, bytes: Vec<u8>, mime: &[u8]) -> Fallible<FetchedData> {
    let mime_str = if let Ok(s) = str::from_utf8(mime) {
        s
    } else {
        ""
    };
    let mime: Mime = mime_str.parse().map_err(
        |_| Error::Type("Inappropriate MIME-type for Body".to_string()))?;
    match mime {
        // TODO
        // ... Parser for Mime(TopLevel::Multipart, SubLevel::FormData, _)
        // ... is not fully determined yet.
        Mime(TopLevel::Application, SubLevel::WwwFormUrlEncoded, _) => {
            let entries = form_urlencoded::parse(&bytes);
            let formdata = FormData::new(None, root);
            for (k, e) in entries {
                formdata.Append(USVString(k.into_owned()), USVString(e.into_owned()));
            }
            return Ok(FetchedData::FormData(formdata));
        },
        _ => return Err(Error::Type("Inappropriate MIME-type for Body".to_string())),
    }
}

#[allow(unsafe_code)]
unsafe fn run_array_buffer_data_algorithm(cx: *mut JSContext, bytes: Vec<u8>) -> Fallible<FetchedData> {
    rooted!(in(cx) let mut array_buffer_ptr = ptr::null_mut::<JSObject>());
    let arraybuffer = ArrayBuffer::create(cx, CreateWith::Slice(&bytes), array_buffer_ptr.handle_mut());
    if arraybuffer.is_err() {
        return Err(Error::JSFailed);
    }
    let rooted_heap = RootedTraceableBox::from_box(Heap::boxed(array_buffer_ptr.get()));
    Ok(FetchedData::ArrayBuffer(rooted_heap))
}

pub trait BodyOperations {
    fn get_body_used(&self) -> bool;
    fn set_body_promise(&self, p: &Rc<Promise>, body_type: BodyType);
    /// Returns `Some(_)` if the body is complete, `None` if there is more to
    /// come.
    fn take_body(&self) -> Option<Vec<u8>>;
    fn is_locked(&self) -> bool;
    fn get_mime_type(&self) -> Ref<Vec<u8>>;
}
<commit_msg>Reordered imports to pass tidy<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/. */

use dom::bindings::codegen::Bindings::FormDataBinding::FormDataMethods;
use dom::bindings::error::{Error, Fallible};
use dom::bindings::reflector::DomObject;
use dom::bindings::root::DomRoot;
use dom::bindings::str::USVString;
use dom::bindings::trace::RootedTraceableBox;
use dom::blob::{Blob, BlobImpl};
use dom::formdata::FormData;
use dom::globalscope::GlobalScope;
use dom::promise::Promise;
use js::jsapi::Heap;
use js::jsapi::JSContext;
use js::jsapi::JSObject;
use js::jsapi::JS_ClearPendingException;
use js::jsapi::JS_GetPendingException;
use js::jsapi::JS_ParseJSON;
use js::jsapi::Value as JSValue;
use js::jsval::JSVal;
use js::jsval::UndefinedValue;
use js::typedarray::{ArrayBuffer, CreateWith};
use mime::{Mime, TopLevel, SubLevel};
use std::cell::Ref;
use std::ptr;
use std::rc::Rc;
use std::str;
use url::form_urlencoded;

#[derive(Clone, Copy, JSTraceable, MallocSizeOf)]
pub enum BodyType {
    Blob,
    FormData,
    Json,
    Text,
    ArrayBuffer
}

pub enum FetchedData {
    Text(String),
    Json(RootedTraceableBox<Heap<JSValue>>),
    BlobData(DomRoot<Blob>),
    FormData(DomRoot<FormData>),
    ArrayBuffer(RootedTraceableBox<Heap<*mut JSObject>>),
    JSException(RootedTraceableBox<Heap<JSVal>>)
}

// https://fetch.spec.whatwg.org/#concept-body-consume-body
#[allow(unrooted_must_root)]
pub fn consume_body<T: BodyOperations + DomObject>(object: &T, body_type: BodyType) -> Rc<Promise> {
    let promise = Promise::new(&object.global());

    // Step 1
    if object.get_body_used() || object.is_locked() {
        promise.reject_error(Error::Type(
            "The response's stream is disturbed or locked".to_string(),
        ));
        return promise;
    }

    object.set_body_promise(&promise, body_type);

    // Steps 2-4
    // TODO: Body does not yet have a stream.

    consume_body_with_promise(object, body_type, &promise);

    promise
}

// https://fetch.spec.whatwg.org/#concept-body-consume-body
#[allow(unrooted_must_root)]
pub fn consume_body_with_promise<T: BodyOperations + DomObject>(object: &T,
                                                                body_type: BodyType,
                                                                promise: &Promise) {
    // Step 5
    let body = match object.take_body() {
        Some(body) => body,
        None => return,
    };

    let pkg_data_results = run_package_data_algorithm(object,
                                                      body,
                                                      body_type,
                                                      object.get_mime_type());

    match pkg_data_results {
        Ok(results) => {
            match results {
                FetchedData::Text(s) => promise.resolve_native(&USVString(s)),
                FetchedData::Json(j) => promise.resolve_native(&j),
                FetchedData::BlobData(b) => promise.resolve_native(&b),
                FetchedData::FormData(f) => promise.resolve_native(&f),
                FetchedData::ArrayBuffer(a) => promise.resolve_native(&a),
                FetchedData::JSException(e) => promise.reject_native(&e.handle()),
            };
        },
        Err(err) => promise.reject_error(err),
    }
}

// https://fetch.spec.whatwg.org/#concept-body-package-data
#[allow(unsafe_code)]
fn run_package_data_algorithm<T: BodyOperations + DomObject>(object: &T,
                                                             bytes: Vec<u8>,
                                                             body_type: BodyType,
                                                             mime_type: Ref<Vec<u8>>)
                                                             -> Fallible<FetchedData> {
    let global = object.global();
    let cx = global.get_cx();
    let mime = &*mime_type;
    match body_type {
        BodyType::Text => run_text_data_algorithm(bytes),
        BodyType::Json => run_json_data_algorithm(cx, bytes),
        BodyType::Blob => run_blob_data_algorithm(&global, bytes, mime),
        BodyType::FormData => run_form_data_algorithm(&global, bytes, mime),
        BodyType::ArrayBuffer => unsafe {
            run_array_buffer_data_algorithm(cx, bytes)
        }
    }
}

fn run_text_data_algorithm(bytes: Vec<u8>) -> Fallible<FetchedData> {
    Ok(FetchedData::Text(String::from_utf8_lossy(&bytes).into_owned()))
}

#[allow(unsafe_code)]
fn run_json_data_algorithm(cx: *mut JSContext,
                           bytes: Vec<u8>) -> Fallible<FetchedData> {
    let json_text = String::from_utf8_lossy(&bytes);
    let json_text: Vec<u16> = json_text.encode_utf16().collect();
    rooted!(in(cx) let mut rval = UndefinedValue());
    unsafe {
        if !JS_ParseJSON(cx,
                         json_text.as_ptr(),
                         json_text.len() as u32,
                         rval.handle_mut()) {
            rooted!(in(cx) let mut exception = UndefinedValue());
            JS_GetPendingException(cx, exception.handle_mut());
            JS_ClearPendingException(cx);
            return Ok(FetchedData::JSException(RootedTraceableBox::from_box(Heap::boxed(exception.get()))));
        }
        let rooted_heap = RootedTraceableBox::from_box(Heap::boxed(rval.get()));
        Ok(FetchedData::Json(rooted_heap))
    }
}

fn run_blob_data_algorithm(root: &GlobalScope,
                           bytes: Vec<u8>,
                           mime: &[u8]) -> Fallible<FetchedData> {
    let mime_string = if let Ok(s) = String::from_utf8(mime.to_vec()) {
        s
    } else {
        "".to_string()
    };
    let blob = Blob::new(root, BlobImpl::new_from_bytes(bytes), mime_string);
    Ok(FetchedData::BlobData(blob))
}

fn run_form_data_algorithm(root: &GlobalScope, bytes: Vec<u8>, mime: &[u8]) -> Fallible<FetchedData> {
    let mime_str = if let Ok(s) = str::from_utf8(mime) {
        s
    } else {
        ""
    };
    let mime: Mime = mime_str.parse().map_err(
        |_| Error::Type("Inappropriate MIME-type for Body".to_string()))?;
    match mime {
        // TODO
        // ... Parser for Mime(TopLevel::Multipart, SubLevel::FormData, _)
        // ... is not fully determined yet.
        Mime(TopLevel::Application, SubLevel::WwwFormUrlEncoded, _) => {
            let entries = form_urlencoded::parse(&bytes);
            let formdata = FormData::new(None, root);
            for (k, e) in entries {
                formdata.Append(USVString(k.into_owned()), USVString(e.into_owned()));
            }
            return Ok(FetchedData::FormData(formdata));
        },
        _ => return Err(Error::Type("Inappropriate MIME-type for Body".to_string())),
    }
}

#[allow(unsafe_code)]
unsafe fn run_array_buffer_data_algorithm(cx: *mut JSContext, bytes: Vec<u8>) -> Fallible<FetchedData> {
    rooted!(in(cx) let mut array_buffer_ptr = ptr::null_mut::<JSObject>());
    let arraybuffer = ArrayBuffer::create(cx, CreateWith::Slice(&bytes), array_buffer_ptr.handle_mut());
    if arraybuffer.is_err() {
        return Err(Error::JSFailed);
    }
    let rooted_heap = RootedTraceableBox::from_box(Heap::boxed(array_buffer_ptr.get()));
    Ok(FetchedData::ArrayBuffer(rooted_heap))
}

pub trait BodyOperations {
    fn get_body_used(&self) -> bool;
    fn set_body_promise(&self, p: &Rc<Promise>, body_type: BodyType);
    /// Returns `Some(_)` if the body is complete, `None` if there is more to
    /// come.
    fn take_body(&self) -> Option<Vec<u8>>;
    fn is_locked(&self) -> bool;
    fn get_mime_type(&self) -> Ref<Vec<u8>>;
}
<|endoftext|>


<commit_before>// Copyright (C) 2017, 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0
use argh::FromArgs;
use slog_scope::info;
use std::path::PathBuf;

#[derive(FromArgs)]
/// Top-level command.
struct TopLevel {
    #[argh(subcommand)]
    entry_point: EntryPoints,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum EntryPoints {
    Client(ClientOptions),
    Server(ServerOptions),
}

#[derive(FromArgs)]
/// Client subcommand
#[argh(subcommand, name = "client")]
struct ClientOptions {
    #[argh(subcommand)]
    commands: ClientCommands,

    /// change the client socket to listen
    #[argh(option, default = "String::from(\"localhost:8080\")")]
    listen_socket: String,

    /// print the output in json format
    #[argh(switch)]
    json: bool,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum ClientCommands {
    Info(Info),
    Log(Log),
    Probe(Probe),
    AbortDownload(AbortDownload),
    InstallPackage(InstallPackage),
}

#[derive(FromArgs)]
/// Fetches information about the current state of the agent
#[argh(subcommand, name = "info")]
struct Info {}

#[derive(FromArgs)]
/// Fetches the available log entries for the last update cycle
#[argh(subcommand, name = "log")]
struct Log {}

#[derive(FromArgs)]
/// Checks if the server has a new update for this device.
///
/// A custom server for the update cycle can be specified via the ´--server´
#[argh(subcommand, name = "probe")]
struct Probe {
    /// custom address to try probe
    #[argh(option)]
    server: Option<String>,
}

#[derive(FromArgs)]
/// Abort current running download
#[argh(subcommand, name = "abort-download")]
struct AbortDownload {}

#[derive(FromArgs)]
/// Request agent to install a package from a direct URL or a local
/// path
#[argh(subcommand, name = "install-package")]
struct InstallPackage {
    /// the URL or path to the update package
    #[argh(positional)]
    arg: String,
}

#[derive(FromArgs)]
/// Server subcommand
#[argh(subcommand, name = "server")]
struct ServerOptions {
    /// increase the verboseness level
    #[argh(option, short = 'v', from_str_fn(verbosity_level), default = "slog::Level::Info")]
    verbosity: slog::Level,

    /// configuration file to use (defaults to "/etc/updatehub.conf")
    #[argh(option, short = 'c', default = "PathBuf::from(\"/etc/updatehub.conf\")")]
    config: PathBuf,
}

fn verbosity_level(value: &str) -> Result<slog::Level, String> {
    use std::str::FromStr;
    slog::Level::from_str(value).map_err(|_| format!("failed to parse verbosity level: {}", value))
}

async fn server_main(cmd: ServerOptions) -> updatehub::Result<()> {
    let _guard = updatehub::logger::init(cmd.verbosity);
    info!("starting UpdateHub Agent {}", updatehub::version());

    updatehub::run(&cmd.config).await?;

    Ok(())
}

async fn client_main(client_options: ClientOptions) -> updatehub::Result<()> {
    let client = sdk::Client::new(&client_options.listen_socket);

    match client_options.commands {
        ClientCommands::Info(_) => {
            let response = client.info().await?;

            if client_options.json {
                println!("{}", serde_json::to_string(&response)?);
            } else {
                println!("{:#?}", response);
            }
        }
        ClientCommands::Log(_) => {
            let response = client.log().await?;

            if client_options.json {
                println!("{}", serde_json::to_string(&response)?);
            } else {
                println!("{}", response);
            }
        }
        ClientCommands::Probe(Probe { server }) => {
            let response = client.probe(server).await?;

            if client_options.json {
                println!("{}", serde_json::to_string(&response)?);
            } else {
                match response {
                    sdk::api::probe::Response::Updating => {
                        println!("Update available. The update is running in background.")
                    }
                    sdk::api::probe::Response::NoUpdate => {
                        println!("There are no updates available.")
                    }
                    sdk::api::probe::Response::TryAgain(t) => {
                        println!("Server replied asking us to try again in {} seconds", t);
                    }
                }
            }
        }
        ClientCommands::AbortDownload(_) => {
            let response = client.abort_download().await?;

            if client_options.json {
                println!("{}", serde_json::to_string(&response)?);
            } else {
                println!("{:#?}", response);
            }
        }
        ClientCommands::InstallPackage(InstallPackage { arg }) => {
            let response = if arg.starts_with("http://") {
                client.remote_install(&arg).await?
            } else {
                let file = PathBuf::from(&arg);
                let file = if file.is_absolute() {
                    file
                } else {
                    std::env::current_dir().unwrap().join(arg)
                };

                client.local_install(&file).await?
            };

            if client_options.json {
                println!("{}", serde_json::to_string(&response)?);
            } else {
                println!("{:#?}", response);
            }
        }
    }

    Ok(())
}

#[async_std::main]
async fn main() {
    let cmd: TopLevel = argh::from_env();

    let res = match cmd.entry_point {
        EntryPoints::Client(client) => client_main(client).await,
        EntryPoints::Server(cmd) => server_main(cmd).await,
    };

    if let Err(e) = res {
        eprintln!("{}", e);
        std::process::exit(1);
    }
}
<commit_msg>updatehub: Rework the cli interface so it is easier to use<commit_after>// Copyright (C) 2017, 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0
use argh::FromArgs;
use slog_scope::info;
use std::path::PathBuf;

#[derive(FromArgs)]
/// Top-level command.
struct TopLevel {
    #[argh(subcommand)]
    entry_point: EntryPoints,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum EntryPoints {
    Client(ClientOptions),
    Daemon(DaemonOptions),
}

#[derive(FromArgs)]
/// Query or send control commands to the UpdateHub Agent daemon
#[argh(subcommand, name = "client")]
struct ClientOptions {
    #[argh(subcommand)]
    commands: ClientCommands,

    /// address where the UpdateHub Agent daemon is running
    #[argh(option, default = "String::from(\"localhost:8080\")")]
    daemon_address: String,

    /// set the output format to JSON
    #[argh(switch)]
    json_output: bool,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum ClientCommands {
    Info(Info),
    Log(Log),
    Probe(Probe),
    AbortDownload(AbortDownload),
    InstallPackage(InstallPackage),
}

#[derive(FromArgs)]
/// Query the current state of the UpdateHub Agent
#[argh(subcommand, name = "info")]
struct Info {}

#[derive(FromArgs)]
/// Show the UpdateHub Agent last update/probe log
#[argh(subcommand, name = "log")]
struct Log {}

#[derive(FromArgs)]
/// Probe the UpdateHub server if there is an update available
#[argh(subcommand, name = "probe")]
struct Probe {
    /// override the UpdateHub daemon to use for inquiry.
    #[argh(option)]
    server: Option<String>,
}

#[derive(FromArgs)]
/// Ask UpdateHub Agent to abort any currently running download
#[argh(subcommand, name = "abort-download")]
struct AbortDownload {}

#[derive(FromArgs)]
/// Install a package from a direct URL or a local path
#[argh(subcommand, name = "install-package")]
struct InstallPackage {
    /// the URL or path to the update package
    #[argh(positional)]
    arg: String,
}

#[derive(FromArgs)]
/// Starts the UpdateHub Agent daemon
#[argh(subcommand, name = "daemon")]
struct DaemonOptions {
    /// increase the log level verboseness
    #[argh(option, short = 'v', from_str_fn(verbosity_level), default = "slog::Level::Info")]
    verbosity: slog::Level,

    /// configuration file to use (defaults to "/etc/updatehub.conf")
    #[argh(option, short = 'c', default = "PathBuf::from(\"/etc/updatehub.conf\")")]
    config: PathBuf,
}

fn verbosity_level(value: &str) -> Result<slog::Level, String> {
    use std::str::FromStr;
    slog::Level::from_str(value).map_err(|_| format!("failed to parse verbosity level: {}", value))
}

async fn daemon_main(cmd: DaemonOptions) -> updatehub::Result<()> {
    let _guard = updatehub::logger::init(cmd.verbosity);
    info!("starting UpdateHub Agent {}", updatehub::version());

    updatehub::run(&cmd.config).await?;

    Ok(())
}

async fn client_main(client_options: ClientOptions) -> updatehub::Result<()> {
    let client = sdk::Client::new(&client_options.daemon_address);

    match client_options.commands {
        ClientCommands::Info(_) => {
            let response = client.info().await?;

            if client_options.json_output {
                println!("{}", serde_json::to_string(&response)?);
            } else {
                println!("{:#?}", response);
            }
        }
        ClientCommands::Log(_) => {
            let response = client.log().await?;

            if client_options.json_output {
                println!("{}", serde_json::to_string(&response)?);
            } else {
                println!("{}", response);
            }
        }
        ClientCommands::Probe(Probe { server }) => {
            let response = client.probe(server).await?;

            if client_options.json_output {
                println!("{}", serde_json::to_string(&response)?);
            } else {
                match response {
                    sdk::api::probe::Response::Updating => {
                        println!("Update available. The update is running in background.")
                    }
                    sdk::api::probe::Response::NoUpdate => {
                        println!("There are no updates available.")
                    }
                    sdk::api::probe::Response::TryAgain(t) => {
                        println!("Server replied asking us to try again in {} seconds", t);
                    }
                }
            }
        }
        ClientCommands::AbortDownload(_) => {
            let response = client.abort_download().await?;

            if client_options.json_output {
                println!("{}", serde_json::to_string(&response)?);
            } else {
                println!("{:#?}", response);
            }
        }
        ClientCommands::InstallPackage(InstallPackage { arg }) => {
            let is_remote_install = arg.starts_with("http://") || arg.starts_with("https://");

            let response = if is_remote_install {
                client.remote_install(&arg).await?
            } else {
                let file = PathBuf::from(&arg);
                let file = if file.is_absolute() {
                    file
                } else {
                    std::env::current_dir().unwrap().join(arg)
                };

                client.local_install(&file).await?
            };

            if client_options.json_output {
                println!("{}", serde_json::to_string(&response)?);
            } else {
                println!("{:#?}", response);
            }
        }
    }

    Ok(())
}

#[async_std::main]
async fn main() {
    let cmd: TopLevel = argh::from_env();

    let res = match cmd.entry_point {
        EntryPoints::Client(client) => client_main(client).await,
        EntryPoints::Daemon(cmd) => daemon_main(cmd).await,
    };

    if let Err(e) = res {
        eprintln!("{}", e);
        std::process::exit(1);
    }
}
<|endoftext|>


<commit_before>//! Handle network connections for a varlink service

use {ErrorKind, Result};
//#![feature(getpid)]
//use std::process;
// FIXME
use libc::getpid;
use std::env;
use std::fs;
use std::io::{Read, Write};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::os::unix::io::{FromRawFd, IntoRawFd};
use std::os::unix::net::{UnixListener, UnixStream};
use std::sync::RwLock;
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
// FIXME: abstract unix domains sockets still not in std
// FIXME: https://github.com/rust-lang/rust/issues/14194
use unix_socket::UnixListener as AbstractUnixListener;

#[derive(Debug)]
enum VarlinkListener {
    TCP(Option<TcpListener>, bool),
    UNIX(Option<UnixListener>, bool),
}

#[derive(Debug)]
enum VarlinkStream {
    TCP(TcpStream),
    UNIX(UnixStream),
}

impl<'a> VarlinkStream {
    pub fn split(&mut self) -> Result<(Box<Read + Send + Sync>, Box<Write + Send + Sync>)> {
        match *self {
            VarlinkStream::TCP(ref mut s) => {
                Ok((Box::new(s.try_clone()?), Box::new(s.try_clone()?)))
            }
            VarlinkStream::UNIX(ref mut s) => {
                Ok((Box::new(s.try_clone()?), Box::new(s.try_clone()?)))
            }
        }
    }
    pub fn shutdown(&mut self) -> Result<()> {
        match *self {
            VarlinkStream::TCP(ref mut s) => s.shutdown(Shutdown::Both)?,
            VarlinkStream::UNIX(ref mut s) => s.shutdown(Shutdown::Both)?,
        }
        Ok(())
    }
}

fn activation_listener() -> Result<Option<i32>> {
    let nfds: u32;

    match env::var("LISTEN_FDS") {
        Ok(ref n) => match n.parse::<u32>() {
            Ok(n) if n >= 1 => nfds = n,
            _ => return Ok(None),
        },
        _ => return Ok(None),
    }

    unsafe {
        match env::var("LISTEN_PID") {
            //FIXME: replace Ok(getpid()) with Ok(process::id())
            Ok(ref pid) if pid.parse::<i32>() == Ok(getpid()) => {}
            _ => return Ok(None),
        }
    }

    if nfds == 1 {
        return Ok(Some(3));
    }

    let fdnames: String;

    match env::var("LISTEN_FDNAMES") {
        Ok(n) => {
            fdnames = n;
        }
        _ => return Ok(None),
    }

    for (i, v) in fdnames.split(":").enumerate() {
        if v == "varlink" {
            return Ok(Some(3 + i as i32));
        }
    }

    Ok(None)
}

//FIXME: add Drop with shutdown() and unix file removal
impl VarlinkListener {
    pub fn new<S: ?Sized + AsRef<str>>(address: &S, timeout: u64) -> Result<Self> {
        let address = address.as_ref();
        if let Some(l) = activation_listener()? {
            if address.starts_with("tcp:") {
                unsafe {
                    let s = TcpStream::from_raw_fd(l);
                    if timeout != 0 {
                        s.set_read_timeout(Some(Duration::from_secs(timeout)))?;
                    }
                    return Ok(VarlinkListener::TCP(
                        Some(TcpListener::from_raw_fd(s.into_raw_fd())),
                        true,
                    ));
                }
            } else if address.starts_with("unix:") {
                unsafe {
                    let s = UnixStream::from_raw_fd(l);
                    if timeout != 0 {
                        s.set_read_timeout(Some(Duration::from_secs(timeout)))?;
                    }
                    return Ok(VarlinkListener::UNIX(
                        Some(UnixListener::from_raw_fd(s.into_raw_fd())),
                        true,
                    ));
                }
            } else {
                return Err(ErrorKind::InvalidAddress.into());
            }
        }

        if address.starts_with("tcp:") {
            let l = TcpListener::bind(&address[4..])?;
            unsafe {
                let s = TcpStream::from_raw_fd(l.into_raw_fd());
                if timeout != 0 {
                    s.set_read_timeout(Some(Duration::from_secs(timeout)))?;
                }
                Ok(VarlinkListener::TCP(
                    Some(TcpListener::from_raw_fd(s.into_raw_fd())),
                    false,
                ))
            }
        } else if address.starts_with("unix:") {
            let mut addr = String::from(address[5..].split(";").next().unwrap());
            if addr.starts_with("@") {
                addr = addr.replacen("@", "\0", 1);
                let l = AbstractUnixListener::bind(addr)?;
                unsafe {
                    let s = UnixStream::from_raw_fd(l.into_raw_fd());
                    if timeout != 0 {
                        s.set_read_timeout(Some(Duration::from_secs(timeout)))?;
                    }
                    return Ok(VarlinkListener::UNIX(
                        Some(UnixListener::from_raw_fd(s.into_raw_fd())),
                        false,
                    ));
                }
            }
            // ignore error on non-existant file
            let _ = fs::remove_file(&*addr);
            let l = UnixListener::bind(addr)?;
            unsafe {
                let s = UnixStream::from_raw_fd(l.into_raw_fd());
                if timeout != 0 {
                    s.set_read_timeout(Some(Duration::from_secs(timeout)))?;
                }
                Ok(VarlinkListener::UNIX(
                    Some(UnixListener::from_raw_fd(s.into_raw_fd())),
                    false,
                ))
            }
        } else {
            Err(ErrorKind::InvalidAddress.into())
        }
    }

    pub fn accept(&self) -> Result<VarlinkStream> {
        match self {
            &VarlinkListener::TCP(Some(ref l), _) => {
                let (mut s, _addr) = l.accept()?;
                Ok(VarlinkStream::TCP(s))
            }
            &VarlinkListener::UNIX(Some(ref l), _) => {
                let (mut s, _addr) = l.accept()?;
                Ok(VarlinkStream::UNIX(s))
            }
            _ => Err(ErrorKind::ConnectionClosed.into()),
        }
    }
    pub fn set_nonblocking(&self, b: bool) -> Result<()> {
        match self {
            &VarlinkListener::TCP(Some(ref l), _) => l.set_nonblocking(b)?,
            &VarlinkListener::UNIX(Some(ref l), _) => l.set_nonblocking(b)?,
            _ => Err(ErrorKind::ConnectionClosed)?,
        }
        Ok(())
    }
}

impl Drop for VarlinkListener {
    fn drop(&mut self) {
        match *self {
            VarlinkListener::UNIX(Some(ref listener), false) => {
                if let Ok(local_addr) = listener.local_addr() {
                    if let Some(path) = local_addr.as_pathname() {
                        let _ = fs::remove_file(path);
                    }
                }
            }
            VarlinkListener::UNIX(ref mut listener, true) => {
                if let Some(l) = listener.take() {
                    unsafe {
                        let s = UnixStream::from_raw_fd(l.into_raw_fd());
                        let _ = s.set_read_timeout(None);
                    }
                }
            }
            _ => {}
        }
    }
}

enum Message {
    NewJob(Job),
    Terminate,
}

struct ThreadPool {
    workers: Vec<Worker>,
    num_busy: Arc<RwLock<i64>>,
    sender: mpsc::Sender<Message>,
}

trait FnBox {
    fn call_box(self: Box<Self>);
}

impl<F: FnOnce()> FnBox for F {
    fn call_box(self: Box<F>) {
        (*self)()
    }
}

type Job = Box<FnBox + Send + 'static>;

impl ThreadPool {
    /// Create a new ThreadPool.
    ///
    /// The size is the number of threads in the pool.
    ///
    /// # Panics
    ///
    /// The `new` function will panic if the size is zero.
    pub fn new(size: usize) -> ThreadPool {
        assert!(size > 0);

        let (sender, receiver) = mpsc::channel();

        let receiver = Arc::new(Mutex::new(receiver));

        let mut workers = Vec::with_capacity(size);

        let num_busy = Arc::new(RwLock::new(0 as i64));

        for _ in 0..size {
            workers.push(Worker::new(Arc::clone(&receiver), Arc::clone(&num_busy)));
        }

        ThreadPool {
            workers,
            sender,
            num_busy,
        }
    }

    pub fn execute<F>(&self, f: F)
    where
        F: FnOnce() + Send + 'static,
    {
        let job = Box::new(f);

        self.sender.send(Message::NewJob(job)).unwrap();
    }

    pub fn num_busy(&self) -> i64 {
        let num_busy = self.num_busy.read().unwrap();
        *num_busy
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        for _ in &mut self.workers {
            self.sender.send(Message::Terminate).unwrap();
        }

        for worker in &mut self.workers {
            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
    }
}

struct Worker {
    thread: Option<thread::JoinHandle<()>>,
}

impl Worker {
    fn new(receiver: Arc<Mutex<mpsc::Receiver<Message>>>, num_busy: Arc<RwLock<i64>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let message = receiver.lock().unwrap().recv().unwrap();

            match message {
                Message::NewJob(job) => {
                    {
                        let mut num_busy = num_busy.write().unwrap();
                        *num_busy += 1;
                    }
                    job.call_box();
                    {
                        let mut num_busy = num_busy.write().unwrap();
                        *num_busy -= 1;
                    }
                }
                Message::Terminate => {
                    break;
                }
            }
        });

        Worker {
            thread: Some(thread),
        }
    }
}

/// `listen` creates a server, with `num_worker` threads listening on `varlink_uri`.
///
/// If an `accept_timeout` != 0 is specified, this function returns after the specified
/// amount of seconds, if no new connection is made in that time frame. It still waits for
/// all pending connections to finish.
///
///# Examples
///
///```
/// extern crate failure;
/// extern crate varlink;
/// use failure::Fail;
///
/// let service = varlink::VarlinkService::new(
///     "org.varlink",
///     "test service",
///     "0.1",
///     "http://varlink.org",
///     vec![/* Your varlink interfaces go here */],
/// );
///
/// if let Err(e) = varlink::listen(service, "unix:/tmp/test_listen_timeout", 10, 1) {
///     if e.kind() != varlink::ErrorKind::Timeout {
///         panic!("Error listen: {:?}", e.cause());
///     }
/// }
///```
///# Note
/// You don't have to use this simple server. With the `VarlinkService::handle()` method you
/// can implement your own server model using whatever framework you prefer.
pub fn listen<S: ?Sized + AsRef<str>>(
    service: ::VarlinkService,
    address: &S,
    workers: usize,
    accept_timeout: u64,
) -> Result<()> {
    let service = Arc::new(service);
    let listener = Arc::new(VarlinkListener::new(address, accept_timeout)?);
    listener.set_nonblocking(false)?;
    let pool = ThreadPool::new(workers);

    loop {
        let mut stream: VarlinkStream = match listener.accept() {
            Err(ref e) if e.kind() == ErrorKind::Io(::std::io::ErrorKind::WouldBlock) => {
                if pool.num_busy() == 0 {
                    return Err(ErrorKind::Timeout.into());
                }
                continue;
            }
            r => r?,
        };
        let service = service.clone();

        pool.execute(move || {
            let (mut r, mut w) = stream.split().expect("Could not split stream");

            if let Err(_e) = service.handle(&mut r, &mut w) {
                //println!("Handle Error: {}", e);
                let _ = stream.shutdown();
            }
        });
    }
}
<commit_msg>server: also unset timeout for TCP listening socket<commit_after>//! Handle network connections for a varlink service

use {ErrorKind, Result};
//#![feature(getpid)]
//use std::process;
// FIXME
use libc::getpid;
use std::env;
use std::fs;
use std::io::{Read, Write};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::os::unix::io::{FromRawFd, IntoRawFd};
use std::os::unix::net::{UnixListener, UnixStream};
use std::sync::RwLock;
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
// FIXME: abstract unix domains sockets still not in std
// FIXME: https://github.com/rust-lang/rust/issues/14194
use unix_socket::UnixListener as AbstractUnixListener;

#[derive(Debug)]
enum VarlinkListener {
    TCP(Option<TcpListener>, bool),
    UNIX(Option<UnixListener>, bool),
}

#[derive(Debug)]
enum VarlinkStream {
    TCP(TcpStream),
    UNIX(UnixStream),
}

impl<'a> VarlinkStream {
    pub fn split(&mut self) -> Result<(Box<Read + Send + Sync>, Box<Write + Send + Sync>)> {
        match *self {
            VarlinkStream::TCP(ref mut s) => {
                Ok((Box::new(s.try_clone()?), Box::new(s.try_clone()?)))
            }
            VarlinkStream::UNIX(ref mut s) => {
                Ok((Box::new(s.try_clone()?), Box::new(s.try_clone()?)))
            }
        }
    }
    pub fn shutdown(&mut self) -> Result<()> {
        match *self {
            VarlinkStream::TCP(ref mut s) => s.shutdown(Shutdown::Both)?,
            VarlinkStream::UNIX(ref mut s) => s.shutdown(Shutdown::Both)?,
        }
        Ok(())
    }
}

fn activation_listener() -> Result<Option<i32>> {
    let nfds: u32;

    match env::var("LISTEN_FDS") {
        Ok(ref n) => match n.parse::<u32>() {
            Ok(n) if n >= 1 => nfds = n,
            _ => return Ok(None),
        },
        _ => return Ok(None),
    }

    unsafe {
        match env::var("LISTEN_PID") {
            //FIXME: replace Ok(getpid()) with Ok(process::id())
            Ok(ref pid) if pid.parse::<i32>() == Ok(getpid()) => {}
            _ => return Ok(None),
        }
    }

    if nfds == 1 {
        return Ok(Some(3));
    }

    let fdnames: String;

    match env::var("LISTEN_FDNAMES") {
        Ok(n) => {
            fdnames = n;
        }
        _ => return Ok(None),
    }

    for (i, v) in fdnames.split(":").enumerate() {
        if v == "varlink" {
            return Ok(Some(3 + i as i32));
        }
    }

    Ok(None)
}

impl VarlinkListener {
    pub fn new<S: ?Sized + AsRef<str>>(address: &S, timeout: u64) -> Result<Self> {
        let address = address.as_ref();
        if let Some(l) = activation_listener()? {
            if address.starts_with("tcp:") {
                unsafe {
                    let s = TcpStream::from_raw_fd(l);
                    if timeout != 0 {
                        s.set_read_timeout(Some(Duration::from_secs(timeout)))?;
                    }
                    return Ok(VarlinkListener::TCP(
                        Some(TcpListener::from_raw_fd(s.into_raw_fd())),
                        true,
                    ));
                }
            } else if address.starts_with("unix:") {
                unsafe {
                    let s = UnixStream::from_raw_fd(l);
                    if timeout != 0 {
                        s.set_read_timeout(Some(Duration::from_secs(timeout)))?;
                    }
                    return Ok(VarlinkListener::UNIX(
                        Some(UnixListener::from_raw_fd(s.into_raw_fd())),
                        true,
                    ));
                }
            } else {
                return Err(ErrorKind::InvalidAddress.into());
            }
        }

        if address.starts_with("tcp:") {
            let l = TcpListener::bind(&address[4..])?;
            unsafe {
                let s = TcpStream::from_raw_fd(l.into_raw_fd());
                if timeout != 0 {
                    s.set_read_timeout(Some(Duration::from_secs(timeout)))?;
                }
                Ok(VarlinkListener::TCP(
                    Some(TcpListener::from_raw_fd(s.into_raw_fd())),
                    false,
                ))
            }
        } else if address.starts_with("unix:") {
            let mut addr = String::from(address[5..].split(";").next().unwrap());
            if addr.starts_with("@") {
                addr = addr.replacen("@", "\0", 1);
                let l = AbstractUnixListener::bind(addr)?;
                unsafe {
                    let s = UnixStream::from_raw_fd(l.into_raw_fd());
                    if timeout != 0 {
                        s.set_read_timeout(Some(Duration::from_secs(timeout)))?;
                    }
                    return Ok(VarlinkListener::UNIX(
                        Some(UnixListener::from_raw_fd(s.into_raw_fd())),
                        false,
                    ));
                }
            }
            // ignore error on non-existant file
            let _ = fs::remove_file(&*addr);
            let l = UnixListener::bind(addr)?;
            unsafe {
                let s = UnixStream::from_raw_fd(l.into_raw_fd());
                if timeout != 0 {
                    s.set_read_timeout(Some(Duration::from_secs(timeout)))?;
                }
                Ok(VarlinkListener::UNIX(
                    Some(UnixListener::from_raw_fd(s.into_raw_fd())),
                    false,
                ))
            }
        } else {
            Err(ErrorKind::InvalidAddress.into())
        }
    }

    pub fn accept(&self) -> Result<VarlinkStream> {
        match self {
            &VarlinkListener::TCP(Some(ref l), _) => {
                let (mut s, _addr) = l.accept()?;
                Ok(VarlinkStream::TCP(s))
            }
            &VarlinkListener::UNIX(Some(ref l), _) => {
                let (mut s, _addr) = l.accept()?;
                Ok(VarlinkStream::UNIX(s))
            }
            _ => Err(ErrorKind::ConnectionClosed.into()),
        }
    }
    pub fn set_nonblocking(&self, b: bool) -> Result<()> {
        match self {
            &VarlinkListener::TCP(Some(ref l), _) => l.set_nonblocking(b)?,
            &VarlinkListener::UNIX(Some(ref l), _) => l.set_nonblocking(b)?,
            _ => Err(ErrorKind::ConnectionClosed)?,
        }
        Ok(())
    }
}

impl Drop for VarlinkListener {
    fn drop(&mut self) {
        match *self {
            VarlinkListener::UNIX(Some(ref listener), false) => {
                if let Ok(local_addr) = listener.local_addr() {
                    if let Some(path) = local_addr.as_pathname() {
                        let _ = fs::remove_file(path);
                    }
                }
            }
            VarlinkListener::UNIX(ref mut listener, true) => {
                if let Some(l) = listener.take() {
                    unsafe {
                        let s = UnixStream::from_raw_fd(l.into_raw_fd());
                        let _ = s.set_read_timeout(None);
                    }
                }
            }
            VarlinkListener::TCP(ref mut listener, true) => {
                if let Some(l) = listener.take() {
                    unsafe {
                        let s = TcpStream::from_raw_fd(l.into_raw_fd());
                        let _ = s.set_read_timeout(None);
                    }
                }
            }
            _ => {}
        }
    }
}

enum Message {
    NewJob(Job),
    Terminate,
}

struct ThreadPool {
    workers: Vec<Worker>,
    num_busy: Arc<RwLock<i64>>,
    sender: mpsc::Sender<Message>,
}

trait FnBox {
    fn call_box(self: Box<Self>);
}

impl<F: FnOnce()> FnBox for F {
    fn call_box(self: Box<F>) {
        (*self)()
    }
}

type Job = Box<FnBox + Send + 'static>;

impl ThreadPool {
    /// Create a new ThreadPool.
    ///
    /// The size is the number of threads in the pool.
    ///
    /// # Panics
    ///
    /// The `new` function will panic if the size is zero.
    pub fn new(size: usize) -> ThreadPool {
        assert!(size > 0);

        let (sender, receiver) = mpsc::channel();

        let receiver = Arc::new(Mutex::new(receiver));

        let mut workers = Vec::with_capacity(size);

        let num_busy = Arc::new(RwLock::new(0 as i64));

        for _ in 0..size {
            workers.push(Worker::new(Arc::clone(&receiver), Arc::clone(&num_busy)));
        }

        ThreadPool {
            workers,
            sender,
            num_busy,
        }
    }

    pub fn execute<F>(&self, f: F)
    where
        F: FnOnce() + Send + 'static,
    {
        let job = Box::new(f);

        self.sender.send(Message::NewJob(job)).unwrap();
    }

    pub fn num_busy(&self) -> i64 {
        let num_busy = self.num_busy.read().unwrap();
        *num_busy
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        for _ in &mut self.workers {
            self.sender.send(Message::Terminate).unwrap();
        }

        for worker in &mut self.workers {
            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
    }
}

struct Worker {
    thread: Option<thread::JoinHandle<()>>,
}

impl Worker {
    fn new(receiver: Arc<Mutex<mpsc::Receiver<Message>>>, num_busy: Arc<RwLock<i64>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let message = receiver.lock().unwrap().recv().unwrap();

            match message {
                Message::NewJob(job) => {
                    {
                        let mut num_busy = num_busy.write().unwrap();
                        *num_busy += 1;
                    }
                    job.call_box();
                    {
                        let mut num_busy = num_busy.write().unwrap();
                        *num_busy -= 1;
                    }
                }
                Message::Terminate => {
                    break;
                }
            }
        });

        Worker {
            thread: Some(thread),
        }
    }
}

/// `listen` creates a server, with `num_worker` threads listening on `varlink_uri`.
///
/// If an `accept_timeout` != 0 is specified, this function returns after the specified
/// amount of seconds, if no new connection is made in that time frame. It still waits for
/// all pending connections to finish.
///
///# Examples
///
///```
/// extern crate failure;
/// extern crate varlink;
/// use failure::Fail;
///
/// let service = varlink::VarlinkService::new(
///     "org.varlink",
///     "test service",
///     "0.1",
///     "http://varlink.org",
///     vec![/* Your varlink interfaces go here */],
/// );
///
/// if let Err(e) = varlink::listen(service, "unix:/tmp/test_listen_timeout", 10, 1) {
///     if e.kind() != varlink::ErrorKind::Timeout {
///         panic!("Error listen: {:?}", e.cause());
///     }
/// }
///```
///# Note
/// You don't have to use this simple server. With the `VarlinkService::handle()` method you
/// can implement your own server model using whatever framework you prefer.
pub fn listen<S: ?Sized + AsRef<str>>(
    service: ::VarlinkService,
    address: &S,
    workers: usize,
    accept_timeout: u64,
) -> Result<()> {
    let service = Arc::new(service);
    let listener = Arc::new(VarlinkListener::new(address, accept_timeout)?);
    listener.set_nonblocking(false)?;
    let pool = ThreadPool::new(workers);

    loop {
        let mut stream: VarlinkStream = match listener.accept() {
            Err(ref e) if e.kind() == ErrorKind::Io(::std::io::ErrorKind::WouldBlock) => {
                if pool.num_busy() == 0 {
                    return Err(ErrorKind::Timeout.into());
                }
                continue;
            }
            r => r?,
        };
        let service = service.clone();

        pool.execute(move || {
            let (mut r, mut w) = stream.split().expect("Could not split stream");

            if let Err(_e) = service.handle(&mut r, &mut w) {
                //println!("Handle Error: {}", e);
                let _ = stream.shutdown();
            }
        });
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! An analysis to determine which locals require allocas and
//! which do not.

use rustc_data_structures::bitvec::BitVector;
use rustc_data_structures::indexed_vec::{Idx, IndexVec};
use rustc::mir::{self, Location, TerminatorKind};
use rustc::mir::visit::{Visitor, LvalueContext};
use rustc::mir::traversal;
use common;
use super::MirContext;
use super::rvalue;

pub fn lvalue_locals<'a, 'tcx>(mircx: &MirContext<'a, 'tcx>) -> BitVector {
    let mir = mircx.mir;
    let mut analyzer = LocalAnalyzer::new(mircx);

    analyzer.visit_mir(mir);

    for (index, ty) in mir.local_decls.iter().map(|l| l.ty).enumerate() {
        let ty = mircx.monomorphize(&ty);
        debug!("local {} has type {:?}", index, ty);
        if ty.is_scalar() ||
            ty.is_box() ||
            ty.is_region_ptr() ||
            ty.is_simd() ||
            common::type_is_zero_size(mircx.ccx, ty)
        {
            // These sorts of types are immediates that we can store
            // in an ValueRef without an alloca.
            assert!(common::type_is_immediate(mircx.ccx, ty) ||
                    common::type_is_fat_ptr(mircx.ccx, ty));
        } else if common::type_is_imm_pair(mircx.ccx, ty) {
            // We allow pairs and uses of any of their 2 fields.
        } else {
            // These sorts of types require an alloca. Note that
            // type_is_immediate() may *still* be true, particularly
            // for newtypes, but we currently force some types
            // (e.g. structs) into an alloca unconditionally, just so
            // that we don't have to deal with having two pathways
            // (gep vs extractvalue etc).
            analyzer.mark_as_lvalue(mir::Local::new(index));
        }
    }

    analyzer.lvalue_locals
}

struct LocalAnalyzer<'mir, 'a: 'mir, 'tcx: 'a> {
    cx: &'mir MirContext<'a, 'tcx>,
    lvalue_locals: BitVector,
    seen_assigned: BitVector
}

impl<'mir, 'a, 'tcx> LocalAnalyzer<'mir, 'a, 'tcx> {
    fn new(mircx: &'mir MirContext<'a, 'tcx>) -> LocalAnalyzer<'mir, 'a, 'tcx> {
        LocalAnalyzer {
            cx: mircx,
            lvalue_locals: BitVector::new(mircx.mir.local_decls.len()),
            seen_assigned: BitVector::new(mircx.mir.local_decls.len())
        }
    }

    fn mark_as_lvalue(&mut self, local: mir::Local) {
        debug!("marking {:?} as lvalue", local);
        self.lvalue_locals.insert(local.index());
    }

    fn mark_assigned(&mut self, local: mir::Local) {
        if !self.seen_assigned.insert(local.index()) {
            self.mark_as_lvalue(local);
        }
    }
}

impl<'mir, 'a, 'tcx> Visitor<'tcx> for LocalAnalyzer<'mir, 'a, 'tcx> {
    fn visit_assign(&mut self,
                    block: mir::BasicBlock,
                    lvalue: &mir::Lvalue<'tcx>,
                    rvalue: &mir::Rvalue<'tcx>,
                    location: Location) {
        debug!("visit_assign(block={:?}, lvalue={:?}, rvalue={:?})", block, lvalue, rvalue);

        if let mir::Lvalue::Local(index) = *lvalue {
            self.mark_assigned(index);
            if !rvalue::rvalue_creates_operand(rvalue) {
                self.mark_as_lvalue(index);
            }
        } else {
            self.visit_lvalue(lvalue, LvalueContext::Store, location);
        }

        self.visit_rvalue(rvalue, location);
    }

    fn visit_terminator_kind(&mut self,
                             block: mir::BasicBlock,
                             kind: &mir::TerminatorKind<'tcx>,
                             location: Location) {
        match *kind {
            mir::TerminatorKind::Call {
                func: mir::Operand::Constant(mir::Constant {
                    literal: mir::Literal::Item { def_id, .. }, ..
                }),
                ref args, ..
            } if Some(def_id) == self.cx.ccx.tcx().lang_items.box_free_fn() => {
                // box_free(x) shares with `drop x` the property that it
                // is not guaranteed to be statically dominated by the
                // definition of x, so x must always be in an alloca.
                if let mir::Operand::Consume(ref lvalue) = args[0] {
                    self.visit_lvalue(lvalue, LvalueContext::Drop, location);
                }
            }
            _ => {}
        }

        self.super_terminator_kind(block, kind, location);
    }

    fn visit_lvalue(&mut self,
                    lvalue: &mir::Lvalue<'tcx>,
                    context: LvalueContext<'tcx>,
                    location: Location) {
        debug!("visit_lvalue(lvalue={:?}, context={:?})", lvalue, context);

        // Allow uses of projections of immediate pair fields.
        if let mir::Lvalue::Projection(ref proj) = *lvalue {
            if let mir::Lvalue::Local(_) = proj.base {
                let ty = proj.base.ty(self.cx.mir, self.cx.ccx.tcx());

                let ty = self.cx.monomorphize(&ty.to_ty(self.cx.ccx.tcx()));
                if common::type_is_imm_pair(self.cx.ccx, ty) {
                    if let mir::ProjectionElem::Field(..) = proj.elem {
                        if let LvalueContext::Consume = context {
                            return;
                        }
                    }
                }
            }
        }

        if let mir::Lvalue::Local(index) = *lvalue {
            match context {
                LvalueContext::Call => {
                    self.mark_assigned(index);
                }

                LvalueContext::StorageLive |
                LvalueContext::StorageDead |
                LvalueContext::Consume => {}

                LvalueContext::Store |
                LvalueContext::Inspect |
                LvalueContext::Borrow { .. } |
                LvalueContext::Projection(..) => {
                    self.mark_as_lvalue(index);
                }

                LvalueContext::Drop => {
                    let ty = lvalue.ty(self.cx.mir, self.cx.ccx.tcx());
                    let ty = self.cx.monomorphize(&ty.to_ty(self.cx.ccx.tcx()));

                    // Only need the lvalue if we're actually dropping it.
                    if self.cx.ccx.shared().type_needs_drop(ty) {
                        self.mark_as_lvalue(index);
                    }
                }
            }
        }

        // A deref projection only reads the pointer, never needs the lvalue.
        if let mir::Lvalue::Projection(ref proj) = *lvalue {
            if let mir::ProjectionElem::Deref = proj.elem {
                return self.visit_lvalue(&proj.base, LvalueContext::Consume, location);
            }
        }

        self.super_lvalue(lvalue, context, location);
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum CleanupKind {
    NotCleanup,
    Funclet,
    Internal { funclet: mir::BasicBlock }
}

pub fn cleanup_kinds<'a, 'tcx>(mir: &mir::Mir<'tcx>) -> IndexVec<mir::BasicBlock, CleanupKind> {
    fn discover_masters<'tcx>(result: &mut IndexVec<mir::BasicBlock, CleanupKind>,
                              mir: &mir::Mir<'tcx>) {
        for (bb, data) in mir.basic_blocks().iter_enumerated() {
            match data.terminator().kind {
                TerminatorKind::Goto { .. } |
                TerminatorKind::Resume |
                TerminatorKind::Return |
                TerminatorKind::Unreachable |
                TerminatorKind::SwitchInt { .. } => {
                    /* nothing to do */
                }
                TerminatorKind::Call { cleanup: unwind, .. } |
                TerminatorKind::Assert { cleanup: unwind, .. } |
                TerminatorKind::DropAndReplace { unwind, .. } |
                TerminatorKind::Drop { unwind, .. } => {
                    if let Some(unwind) = unwind {
                        debug!("cleanup_kinds: {:?}/{:?} registering {:?} as funclet",
                               bb, data, unwind);
                        result[unwind] = CleanupKind::Funclet;
                    }
                }
            }
        }
    }

    fn propagate<'tcx>(result: &mut IndexVec<mir::BasicBlock, CleanupKind>,
                       mir: &mir::Mir<'tcx>) {
        let mut funclet_succs = IndexVec::from_elem(None, mir.basic_blocks());

        let mut set_successor = |funclet: mir::BasicBlock, succ| {
            match funclet_succs[funclet] {
                ref mut s @ None => {
                    debug!("set_successor: updating successor of {:?} to {:?}",
                           funclet, succ);
                    *s = Some(succ);
                },
                Some(s) => if s != succ {
                    span_bug!(mir.span, "funclet {:?} has 2 parents - {:?} and {:?}",
                              funclet, s, succ);
                }
            }
        };

        for (bb, data) in traversal::reverse_postorder(mir) {
            let funclet = match result[bb] {
                CleanupKind::NotCleanup => continue,
                CleanupKind::Funclet => bb,
                CleanupKind::Internal { funclet } => funclet,
            };

            debug!("cleanup_kinds: {:?}/{:?}/{:?} propagating funclet {:?}",
                   bb, data, result[bb], funclet);

            for &succ in data.terminator().successors().iter() {
                let kind = result[succ];
                debug!("cleanup_kinds: propagating {:?} to {:?}/{:?}",
                       funclet, succ, kind);
                match kind {
                    CleanupKind::NotCleanup => {
                        result[succ] = CleanupKind::Internal { funclet: funclet };
                    }
                    CleanupKind::Funclet => {
                        set_successor(funclet, succ);
                    }
                    CleanupKind::Internal { funclet: succ_funclet } => {
                        if funclet != succ_funclet {
                            // `succ` has 2 different funclet going into it, so it must
                            // be a funclet by itself.

                            debug!("promoting {:?} to a funclet and updating {:?}", succ,
                                   succ_funclet);
                            result[succ] = CleanupKind::Funclet;
                            set_successor(succ_funclet, succ);
                            set_successor(funclet, succ);
                        }
                    }
                }
            }
        }
    }

    let mut result = IndexVec::from_elem(CleanupKind::NotCleanup, mir.basic_blocks());

    discover_masters(&mut result, mir);
    propagate(&mut result, mir);
    debug!("cleanup_kinds: result={:?}", result);
    result
}
<commit_msg>Inspect now does not force on-stack Lvalue<commit_after>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! An analysis to determine which locals require allocas and
//! which do not.

use rustc_data_structures::bitvec::BitVector;
use rustc_data_structures::indexed_vec::{Idx, IndexVec};
use rustc::mir::{self, Location, TerminatorKind};
use rustc::mir::visit::{Visitor, LvalueContext};
use rustc::mir::traversal;
use common;
use super::MirContext;
use super::rvalue;

pub fn lvalue_locals<'a, 'tcx>(mircx: &MirContext<'a, 'tcx>) -> BitVector {
    let mir = mircx.mir;
    let mut analyzer = LocalAnalyzer::new(mircx);

    analyzer.visit_mir(mir);

    for (index, ty) in mir.local_decls.iter().map(|l| l.ty).enumerate() {
        let ty = mircx.monomorphize(&ty);
        debug!("local {} has type {:?}", index, ty);
        if ty.is_scalar() ||
            ty.is_box() ||
            ty.is_region_ptr() ||
            ty.is_simd() ||
            common::type_is_zero_size(mircx.ccx, ty)
        {
            // These sorts of types are immediates that we can store
            // in an ValueRef without an alloca.
            assert!(common::type_is_immediate(mircx.ccx, ty) ||
                    common::type_is_fat_ptr(mircx.ccx, ty));
        } else if common::type_is_imm_pair(mircx.ccx, ty) {
            // We allow pairs and uses of any of their 2 fields.
        } else {
            // These sorts of types require an alloca. Note that
            // type_is_immediate() may *still* be true, particularly
            // for newtypes, but we currently force some types
            // (e.g. structs) into an alloca unconditionally, just so
            // that we don't have to deal with having two pathways
            // (gep vs extractvalue etc).
            analyzer.mark_as_lvalue(mir::Local::new(index));
        }
    }

    analyzer.lvalue_locals
}

struct LocalAnalyzer<'mir, 'a: 'mir, 'tcx: 'a> {
    cx: &'mir MirContext<'a, 'tcx>,
    lvalue_locals: BitVector,
    seen_assigned: BitVector
}

impl<'mir, 'a, 'tcx> LocalAnalyzer<'mir, 'a, 'tcx> {
    fn new(mircx: &'mir MirContext<'a, 'tcx>) -> LocalAnalyzer<'mir, 'a, 'tcx> {
        LocalAnalyzer {
            cx: mircx,
            lvalue_locals: BitVector::new(mircx.mir.local_decls.len()),
            seen_assigned: BitVector::new(mircx.mir.local_decls.len())
        }
    }

    fn mark_as_lvalue(&mut self, local: mir::Local) {
        debug!("marking {:?} as lvalue", local);
        self.lvalue_locals.insert(local.index());
    }

    fn mark_assigned(&mut self, local: mir::Local) {
        if !self.seen_assigned.insert(local.index()) {
            self.mark_as_lvalue(local);
        }
    }
}

impl<'mir, 'a, 'tcx> Visitor<'tcx> for LocalAnalyzer<'mir, 'a, 'tcx> {
    fn visit_assign(&mut self,
                    block: mir::BasicBlock,
                    lvalue: &mir::Lvalue<'tcx>,
                    rvalue: &mir::Rvalue<'tcx>,
                    location: Location) {
        debug!("visit_assign(block={:?}, lvalue={:?}, rvalue={:?})", block, lvalue, rvalue);

        if let mir::Lvalue::Local(index) = *lvalue {
            self.mark_assigned(index);
            if !rvalue::rvalue_creates_operand(rvalue) {
                self.mark_as_lvalue(index);
            }
        } else {
            self.visit_lvalue(lvalue, LvalueContext::Store, location);
        }

        self.visit_rvalue(rvalue, location);
    }

    fn visit_terminator_kind(&mut self,
                             block: mir::BasicBlock,
                             kind: &mir::TerminatorKind<'tcx>,
                             location: Location) {
        match *kind {
            mir::TerminatorKind::Call {
                func: mir::Operand::Constant(mir::Constant {
                    literal: mir::Literal::Item { def_id, .. }, ..
                }),
                ref args, ..
            } if Some(def_id) == self.cx.ccx.tcx().lang_items.box_free_fn() => {
                // box_free(x) shares with `drop x` the property that it
                // is not guaranteed to be statically dominated by the
                // definition of x, so x must always be in an alloca.
                if let mir::Operand::Consume(ref lvalue) = args[0] {
                    self.visit_lvalue(lvalue, LvalueContext::Drop, location);
                }
            }
            _ => {}
        }

        self.super_terminator_kind(block, kind, location);
    }

    fn visit_lvalue(&mut self,
                    lvalue: &mir::Lvalue<'tcx>,
                    context: LvalueContext<'tcx>,
                    location: Location) {
        debug!("visit_lvalue(lvalue={:?}, context={:?})", lvalue, context);

        // Allow uses of projections of immediate pair fields.
        if let mir::Lvalue::Projection(ref proj) = *lvalue {
            if let mir::Lvalue::Local(_) = proj.base {
                let ty = proj.base.ty(self.cx.mir, self.cx.ccx.tcx());

                let ty = self.cx.monomorphize(&ty.to_ty(self.cx.ccx.tcx()));
                if common::type_is_imm_pair(self.cx.ccx, ty) {
                    if let mir::ProjectionElem::Field(..) = proj.elem {
                        if let LvalueContext::Consume = context {
                            return;
                        }
                    }
                }
            }
        }

        if let mir::Lvalue::Local(index) = *lvalue {
            match context {
                LvalueContext::Call => {
                    self.mark_assigned(index);
                }

                LvalueContext::StorageLive |
                LvalueContext::StorageDead |
                LvalueContext::Inspect |
                LvalueContext::Consume => {}

                LvalueContext::Store |
                LvalueContext::Borrow { .. } |
                LvalueContext::Projection(..) => {
                    self.mark_as_lvalue(index);
                }

                LvalueContext::Drop => {
                    let ty = lvalue.ty(self.cx.mir, self.cx.ccx.tcx());
                    let ty = self.cx.monomorphize(&ty.to_ty(self.cx.ccx.tcx()));

                    // Only need the lvalue if we're actually dropping it.
                    if self.cx.ccx.shared().type_needs_drop(ty) {
                        self.mark_as_lvalue(index);
                    }
                }
            }
        }

        // A deref projection only reads the pointer, never needs the lvalue.
        if let mir::Lvalue::Projection(ref proj) = *lvalue {
            if let mir::ProjectionElem::Deref = proj.elem {
                return self.visit_lvalue(&proj.base, LvalueContext::Consume, location);
            }
        }

        self.super_lvalue(lvalue, context, location);
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum CleanupKind {
    NotCleanup,
    Funclet,
    Internal { funclet: mir::BasicBlock }
}

pub fn cleanup_kinds<'a, 'tcx>(mir: &mir::Mir<'tcx>) -> IndexVec<mir::BasicBlock, CleanupKind> {
    fn discover_masters<'tcx>(result: &mut IndexVec<mir::BasicBlock, CleanupKind>,
                              mir: &mir::Mir<'tcx>) {
        for (bb, data) in mir.basic_blocks().iter_enumerated() {
            match data.terminator().kind {
                TerminatorKind::Goto { .. } |
                TerminatorKind::Resume |
                TerminatorKind::Return |
                TerminatorKind::Unreachable |
                TerminatorKind::SwitchInt { .. } => {
                    /* nothing to do */
                }
                TerminatorKind::Call { cleanup: unwind, .. } |
                TerminatorKind::Assert { cleanup: unwind, .. } |
                TerminatorKind::DropAndReplace { unwind, .. } |
                TerminatorKind::Drop { unwind, .. } => {
                    if let Some(unwind) = unwind {
                        debug!("cleanup_kinds: {:?}/{:?} registering {:?} as funclet",
                               bb, data, unwind);
                        result[unwind] = CleanupKind::Funclet;
                    }
                }
            }
        }
    }

    fn propagate<'tcx>(result: &mut IndexVec<mir::BasicBlock, CleanupKind>,
                       mir: &mir::Mir<'tcx>) {
        let mut funclet_succs = IndexVec::from_elem(None, mir.basic_blocks());

        let mut set_successor = |funclet: mir::BasicBlock, succ| {
            match funclet_succs[funclet] {
                ref mut s @ None => {
                    debug!("set_successor: updating successor of {:?} to {:?}",
                           funclet, succ);
                    *s = Some(succ);
                },
                Some(s) => if s != succ {
                    span_bug!(mir.span, "funclet {:?} has 2 parents - {:?} and {:?}",
                              funclet, s, succ);
                }
            }
        };

        for (bb, data) in traversal::reverse_postorder(mir) {
            let funclet = match result[bb] {
                CleanupKind::NotCleanup => continue,
                CleanupKind::Funclet => bb,
                CleanupKind::Internal { funclet } => funclet,
            };

            debug!("cleanup_kinds: {:?}/{:?}/{:?} propagating funclet {:?}",
                   bb, data, result[bb], funclet);

            for &succ in data.terminator().successors().iter() {
                let kind = result[succ];
                debug!("cleanup_kinds: propagating {:?} to {:?}/{:?}",
                       funclet, succ, kind);
                match kind {
                    CleanupKind::NotCleanup => {
                        result[succ] = CleanupKind::Internal { funclet: funclet };
                    }
                    CleanupKind::Funclet => {
                        set_successor(funclet, succ);
                    }
                    CleanupKind::Internal { funclet: succ_funclet } => {
                        if funclet != succ_funclet {
                            // `succ` has 2 different funclet going into it, so it must
                            // be a funclet by itself.

                            debug!("promoting {:?} to a funclet and updating {:?}", succ,
                                   succ_funclet);
                            result[succ] = CleanupKind::Funclet;
                            set_successor(succ_funclet, succ);
                            set_successor(funclet, succ);
                        }
                    }
                }
            }
        }
    }

    let mut result = IndexVec::from_elem(CleanupKind::NotCleanup, mir.basic_blocks());

    discover_masters(&mut result, mir);
    propagate(&mut result, mir);
    debug!("cleanup_kinds: result={:?}", result);
    result
}
<|endoftext|>