<commit_before>extern crate core_foundation;
extern crate cgl;
extern crate gleam;

use view::DrawableGeometry;

use cocoa::appkit;
use cocoa::appkit::*;
use cocoa::foundation::*;
use cocoa::base::*;
use objc::declare::ClassDecl;
use objc::runtime::{Class, Object, Sel};
use self::cgl::{CGLEnable, kCGLCECrashOnRemovedFunctions};
use self::core_foundation::base::TCFType;
use self::core_foundation::string::CFString;
use self::core_foundation::bundle::{CFBundleGetBundleWithIdentifier, CFBundleGetFunctionPointerForName};
use std::os::raw::c_void;
use std::str::FromStr;
use view::{ElementState, MouseButton, ViewEvent, TouchPhase, MouseScrollDelta};
use super::utils;

pub fn register() {
    let superclass = Class::get("NSView").unwrap();
    let mut class = ClassDecl::new("NSServoView", superclass).unwrap();

    class.add_ivar::<*mut c_void>("event_queue");

    extern fn store_nsevent(this: &Object, _sel: Sel, event: id) {
        unsafe { msg_send![event, retain] }
        utils::get_event_queue(this).push(event);
    }

    extern fn awake_from_nib(this: &mut Object, _sel: Sel) {
        // FIXME: is that the best way to create a raw pointer?
        let event_queue: Vec<id> = Vec::new();
        let event_queue_ptr = Box::into_raw(Box::new(event_queue));
        unsafe {
            this.set_ivar("event_queue", event_queue_ptr as *mut c_void);
        }
    }

    extern fn accept_first_responder(_this: &Object, _sel: Sel) -> BOOL {
        YES
    }

    unsafe {
        class.add_method(sel!(scrollWheel:), store_nsevent as extern fn(&Object, Sel, id));
        class.add_method(sel!(mouseDown:), store_nsevent as extern fn(&Object, Sel, id));
        class.add_method(sel!(mouseUp:), store_nsevent as extern fn(&Object, Sel, id));
        class.add_method(sel!(mouseMoved:), store_nsevent as extern fn(&Object, Sel, id));

        class.add_method(sel!(acceptsFirstResponder), accept_first_responder as extern fn(&Object, Sel) -> BOOL);

        class.add_method(sel!(awakeFromNib), awake_from_nib as extern fn(&mut Object, Sel));
    }

    class.register();
}

pub struct View {
    nsview: id,
    context: id,
}

impl View {

    pub fn new(nsview: id) -> View {
        let context: id = View::init_gl(nsview);
        View {
            nsview: nsview,
            context: context
        }
    }

    pub fn swap_buffers(&self) {
        unsafe {
            msg_send![self.context, flushBuffer];
        }
    }

    pub fn update_drawable(&self) {
        unsafe {
            msg_send![self.context, update];
        }
    }

    pub fn get_geometry(&self) -> DrawableGeometry {
        unsafe {
            let nswindow: id = msg_send![self.nsview, window];
            let frame: NSRect = msg_send![self.nsview, frame];
            let hidpi_factor: CGFloat = msg_send![nswindow, backingScaleFactor];

            // FIXME: this is going to be wrong as soon as the view is not full contentView
            let visible_rect: NSRect = msg_send![nswindow, contentLayoutRect];

            // FIXME: coordinates are flipped
            let bottom = visible_rect.origin.y - frame.origin.y;
            let left = visible_rect.origin.x - frame.origin.x;
            let right = frame.size.width - left - visible_rect.size.width;
            let top = frame.size.height - bottom - visible_rect.size.height;

            DrawableGeometry {
                view_size: (frame.size.width as u32, frame.size.height as u32),
                margins: (top as u32, right as u32, bottom as u32, left as u32),
                position: (0, 0),
                hidpi_factor: hidpi_factor as f32,
            }
        }
    }

    pub fn get_events(&self) -> Vec<ViewEvent> {
        let nsobject = unsafe { &*self.nsview};
        let event_queue = utils::get_event_queue(nsobject);
        // FIXME: make sure event_queue is empty
        let events = event_queue.drain(..).map(|e| {
            self.nsevent_to_viewevent(e)
        }).filter(|e| {
            e.is_some()
        }).map(|e| {
            e.unwrap()
        }).collect();
        events
    }

    pub fn enter_fullscreen(&self) {
        unsafe {
            msg_send![self.nsview, enterFullScreenMode:nil withOptions:nil];
        }
    }

    pub fn exit_fullscreen(&self) {
        unsafe {
            msg_send![self.nsview, exitFullScreenModeWithOptions:nil];
        }
    }

    fn nsevent_to_viewevent(&self, nsevent: id) -> Option<ViewEvent> {
        unsafe {
            let event_type = nsevent.eventType();
            match event_type {
                NSScrollWheel => {
                    // Stolen from winit
                    use self::MouseScrollDelta::{LineDelta, PixelDelta};
                    let nswindow: id = nsevent.window();
                    let hidpi_factor: CGFloat = msg_send![nswindow, backingScaleFactor];
                    let hidpi_factor = hidpi_factor as f32;
                    let delta = if nsevent.hasPreciseScrollingDeltas() == YES {
                        PixelDelta(hidpi_factor * nsevent.scrollingDeltaX() as f32,
                        hidpi_factor * nsevent.scrollingDeltaY() as f32)
                    } else {
                        LineDelta(hidpi_factor * nsevent.scrollingDeltaX() as f32,
                        hidpi_factor * nsevent.scrollingDeltaY() as f32)
                    };
                    let phase = match nsevent.phase() {
                        appkit::NSEventPhaseMayBegin | appkit::NSEventPhaseBegan => TouchPhase::Started,
                        appkit::NSEventPhaseEnded => TouchPhase::Ended,
                        _ => TouchPhase::Moved,
                    };
                    Some(ViewEvent::MouseWheel(delta, phase))
                },
                NSMouseMoved => {
                    let nswindow = nsevent.window();
                    let window_point = nsevent.locationInWindow();
                    let view_point = self.nsview.convertPoint_fromView_(window_point, nil);
                    let frame = NSView::frame(self.nsview);
                    let hidpi_factor: CGFloat = msg_send![nswindow, backingScaleFactor];
                    let hidpi_factor = hidpi_factor as f32;
                    let x = (hidpi_factor * view_point.x as f32) as i32;
                    let y = (hidpi_factor * (frame.size.height - view_point.y) as f32) as i32;
                    Some(ViewEvent::MouseMoved(x, y))
                }

                NSLeftMouseDown => { Some(ViewEvent::MouseInput(ElementState::Pressed, MouseButton::Left)) },
                NSLeftMouseUp => { Some(ViewEvent::MouseInput(ElementState::Released, MouseButton::Left)) },
                NSRightMouseDown => { Some(ViewEvent::MouseInput(ElementState::Pressed, MouseButton::Right)) },
                NSRightMouseUp => { Some(ViewEvent::MouseInput(ElementState::Released, MouseButton::Right)) },
                NSOtherMouseDown => { Some(ViewEvent::MouseInput(ElementState::Pressed, MouseButton::Middle)) },
                NSOtherMouseUp => { Some(ViewEvent::MouseInput(ElementState::Released, MouseButton::Middle)) },

                _ => None
            }
        }
    }

    fn init_gl(nsview: id) -> id {
        let ctx = unsafe {
            nsview.setWantsBestResolutionOpenGLSurface_(YES);
            let attributes = vec![NSOpenGLPFADoubleBuffer as u32,
                NSOpenGLPFAClosestPolicy as u32,
                NSOpenGLPFAColorSize as u32,
                32,
                NSOpenGLPFAAlphaSize as u32,
                8,
                NSOpenGLPFADepthSize as u32,
                24,
                NSOpenGLPFAStencilSize as u32,
                8,
                NSOpenGLPFAOpenGLProfile as u32,
                NSOpenGLProfileVersion3_2Core as u32,
                0];

            let pixelformat = NSOpenGLPixelFormat::alloc(nil).initWithAttributes_(&attributes);
            let ctx: id = NSOpenGLContext::alloc(nil).initWithFormat_shareContext_(pixelformat, nil);
            msg_send![ctx, setView: nsview];
            let value = 1;
            ctx.setValues_forParameter_(&value, NSOpenGLContextParameter::NSOpenGLCPSwapInterval);
            CGLEnable(ctx.CGLContextObj() as *mut _, kCGLCECrashOnRemovedFunctions);
            ctx
        };

        unsafe {
            msg_send![ctx, update];
            msg_send![ctx, makeCurrentContext];
        };

        gleam::gl::load_with(|addr| {
            let symbol_name: CFString = FromStr::from_str(addr).unwrap();
            let framework_name: CFString = FromStr::from_str("com.apple.opengl").unwrap();
            let framework = unsafe {
                CFBundleGetBundleWithIdentifier(framework_name.as_concrete_TypeRef())
            };
            let symbol = unsafe {
                CFBundleGetFunctionPointerForName(framework, symbol_name.as_concrete_TypeRef())
            };
            symbol as *const c_void
        });

        gleam::gl::clear_color(1.0, 0.0, 0.0, 1.0);
        gleam::gl::clear(gleam::gl::COLOR_BUFFER_BIT);
        gleam::gl::finish();

        ctx
    }
}

<commit_msg>white clear color<commit_after>extern crate core_foundation;
extern crate cgl;
extern crate gleam;

use view::DrawableGeometry;

use cocoa::appkit;
use cocoa::appkit::*;
use cocoa::foundation::*;
use cocoa::base::*;
use objc::declare::ClassDecl;
use objc::runtime::{Class, Object, Sel};
use self::cgl::{CGLEnable, kCGLCECrashOnRemovedFunctions};
use self::core_foundation::base::TCFType;
use self::core_foundation::string::CFString;
use self::core_foundation::bundle::{CFBundleGetBundleWithIdentifier, CFBundleGetFunctionPointerForName};
use std::os::raw::c_void;
use std::str::FromStr;
use view::{ElementState, MouseButton, ViewEvent, TouchPhase, MouseScrollDelta};
use super::utils;

pub fn register() {
    let superclass = Class::get("NSView").unwrap();
    let mut class = ClassDecl::new("NSServoView", superclass).unwrap();

    class.add_ivar::<*mut c_void>("event_queue");

    extern fn store_nsevent(this: &Object, _sel: Sel, event: id) {
        unsafe { msg_send![event, retain] }
        utils::get_event_queue(this).push(event);
    }

    extern fn awake_from_nib(this: &mut Object, _sel: Sel) {
        // FIXME: is that the best way to create a raw pointer?
        let event_queue: Vec<id> = Vec::new();
        let event_queue_ptr = Box::into_raw(Box::new(event_queue));
        unsafe {
            this.set_ivar("event_queue", event_queue_ptr as *mut c_void);
        }
    }

    extern fn accept_first_responder(_this: &Object, _sel: Sel) -> BOOL {
        YES
    }

    unsafe {
        class.add_method(sel!(scrollWheel:), store_nsevent as extern fn(&Object, Sel, id));
        class.add_method(sel!(mouseDown:), store_nsevent as extern fn(&Object, Sel, id));
        class.add_method(sel!(mouseUp:), store_nsevent as extern fn(&Object, Sel, id));
        class.add_method(sel!(mouseMoved:), store_nsevent as extern fn(&Object, Sel, id));

        class.add_method(sel!(acceptsFirstResponder), accept_first_responder as extern fn(&Object, Sel) -> BOOL);

        class.add_method(sel!(awakeFromNib), awake_from_nib as extern fn(&mut Object, Sel));
    }

    class.register();
}

pub struct View {
    nsview: id,
    context: id,
}

impl View {

    pub fn new(nsview: id) -> View {
        let context: id = View::init_gl(nsview);
        View {
            nsview: nsview,
            context: context
        }
    }

    pub fn swap_buffers(&self) {
        unsafe {
            msg_send![self.context, flushBuffer];
        }
    }

    pub fn update_drawable(&self) {
        unsafe {
            msg_send![self.context, update];
        }
    }

    pub fn get_geometry(&self) -> DrawableGeometry {
        unsafe {
            let nswindow: id = msg_send![self.nsview, window];
            let frame: NSRect = msg_send![self.nsview, frame];
            let hidpi_factor: CGFloat = msg_send![nswindow, backingScaleFactor];

            // FIXME: this is going to be wrong as soon as the view is not full contentView
            let visible_rect: NSRect = msg_send![nswindow, contentLayoutRect];

            // FIXME: coordinates are flipped
            let bottom = visible_rect.origin.y - frame.origin.y;
            let left = visible_rect.origin.x - frame.origin.x;
            let right = frame.size.width - left - visible_rect.size.width;
            let top = frame.size.height - bottom - visible_rect.size.height;

            DrawableGeometry {
                view_size: (frame.size.width as u32, frame.size.height as u32),
                margins: (top as u32, right as u32, bottom as u32, left as u32),
                position: (0, 0),
                hidpi_factor: hidpi_factor as f32,
            }
        }
    }

    pub fn get_events(&self) -> Vec<ViewEvent> {
        let nsobject = unsafe { &*self.nsview};
        let event_queue = utils::get_event_queue(nsobject);
        // FIXME: make sure event_queue is empty
        let events = event_queue.drain(..).map(|e| {
            self.nsevent_to_viewevent(e)
        }).filter(|e| {
            e.is_some()
        }).map(|e| {
            e.unwrap()
        }).collect();
        events
    }

    pub fn enter_fullscreen(&self) {
        unsafe {
            msg_send![self.nsview, enterFullScreenMode:nil withOptions:nil];
        }
    }

    pub fn exit_fullscreen(&self) {
        unsafe {
            msg_send![self.nsview, exitFullScreenModeWithOptions:nil];
        }
    }

    fn nsevent_to_viewevent(&self, nsevent: id) -> Option<ViewEvent> {
        unsafe {
            let event_type = nsevent.eventType();
            match event_type {
                NSScrollWheel => {
                    // Stolen from winit
                    use self::MouseScrollDelta::{LineDelta, PixelDelta};
                    let nswindow: id = nsevent.window();
                    let hidpi_factor: CGFloat = msg_send![nswindow, backingScaleFactor];
                    let hidpi_factor = hidpi_factor as f32;
                    let delta = if nsevent.hasPreciseScrollingDeltas() == YES {
                        PixelDelta(hidpi_factor * nsevent.scrollingDeltaX() as f32,
                        hidpi_factor * nsevent.scrollingDeltaY() as f32)
                    } else {
                        LineDelta(hidpi_factor * nsevent.scrollingDeltaX() as f32,
                        hidpi_factor * nsevent.scrollingDeltaY() as f32)
                    };
                    let phase = match nsevent.phase() {
                        appkit::NSEventPhaseMayBegin | appkit::NSEventPhaseBegan => TouchPhase::Started,
                        appkit::NSEventPhaseEnded => TouchPhase::Ended,
                        _ => TouchPhase::Moved,
                    };
                    Some(ViewEvent::MouseWheel(delta, phase))
                },
                NSMouseMoved => {
                    let nswindow = nsevent.window();
                    let window_point = nsevent.locationInWindow();
                    let view_point = self.nsview.convertPoint_fromView_(window_point, nil);
                    let frame = NSView::frame(self.nsview);
                    let hidpi_factor: CGFloat = msg_send![nswindow, backingScaleFactor];
                    let hidpi_factor = hidpi_factor as f32;
                    let x = (hidpi_factor * view_point.x as f32) as i32;
                    let y = (hidpi_factor * (frame.size.height - view_point.y) as f32) as i32;
                    Some(ViewEvent::MouseMoved(x, y))
                }

                NSLeftMouseDown => { Some(ViewEvent::MouseInput(ElementState::Pressed, MouseButton::Left)) },
                NSLeftMouseUp => { Some(ViewEvent::MouseInput(ElementState::Released, MouseButton::Left)) },
                NSRightMouseDown => { Some(ViewEvent::MouseInput(ElementState::Pressed, MouseButton::Right)) },
                NSRightMouseUp => { Some(ViewEvent::MouseInput(ElementState::Released, MouseButton::Right)) },
                NSOtherMouseDown => { Some(ViewEvent::MouseInput(ElementState::Pressed, MouseButton::Middle)) },
                NSOtherMouseUp => { Some(ViewEvent::MouseInput(ElementState::Released, MouseButton::Middle)) },

                _ => None
            }
        }
    }

    fn init_gl(nsview: id) -> id {
        let ctx = unsafe {
            nsview.setWantsBestResolutionOpenGLSurface_(YES);
            let attributes = vec![NSOpenGLPFADoubleBuffer as u32,
                NSOpenGLPFAClosestPolicy as u32,
                NSOpenGLPFAColorSize as u32,
                32,
                NSOpenGLPFAAlphaSize as u32,
                8,
                NSOpenGLPFADepthSize as u32,
                24,
                NSOpenGLPFAStencilSize as u32,
                8,
                NSOpenGLPFAOpenGLProfile as u32,
                NSOpenGLProfileVersion3_2Core as u32,
                0];

            let pixelformat = NSOpenGLPixelFormat::alloc(nil).initWithAttributes_(&attributes);
            let ctx: id = NSOpenGLContext::alloc(nil).initWithFormat_shareContext_(pixelformat, nil);
            msg_send![ctx, setView: nsview];
            let value = 1;
            ctx.setValues_forParameter_(&value, NSOpenGLContextParameter::NSOpenGLCPSwapInterval);
            CGLEnable(ctx.CGLContextObj() as *mut _, kCGLCECrashOnRemovedFunctions);
            ctx
        };

        unsafe {
            msg_send![ctx, update];
            msg_send![ctx, makeCurrentContext];
        };

        gleam::gl::load_with(|addr| {
            let symbol_name: CFString = FromStr::from_str(addr).unwrap();
            let framework_name: CFString = FromStr::from_str("com.apple.opengl").unwrap();
            let framework = unsafe {
                CFBundleGetBundleWithIdentifier(framework_name.as_concrete_TypeRef())
            };
            let symbol = unsafe {
                CFBundleGetFunctionPointerForName(framework, symbol_name.as_concrete_TypeRef())
            };
            symbol as *const c_void
        });

        gleam::gl::clear_color(1.0, 1.0, 1.0, 1.0);
        gleam::gl::clear(gleam::gl::COLOR_BUFFER_BIT);
        gleam::gl::finish();

        ctx
    }
}

<|endoftext|>


<commit_before>// Copyright 2017 CoreOS, Inc.
//
// 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
//
//     http://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.

//! azure metadata fetcher

mod crypto;

use self::crypto::x509;
use errors::*;
use metadata::Metadata;
use pnet;
use retry;
use openssh_keys::PublicKey;

use std::net::{IpAddr, Ipv4Addr};
use std::fs::File;
use std::path::Path;
use std::io::{BufRead, BufReader};

header! {(MSAgentName, "x-ms-agent-name") => [String]}
header! {(MSVersion, "x-ms-version") => [String]}
header! {(MSCipherName, "x-ms-cipher-name") => [String]}
header! {(MSCert, "x-ms-guest-agent-public-x509-cert") => [String]}

const OPTION_245: &str = "OPTION_245";
const MS_AGENT_NAME: &str = "com.coreos.metadata";
const MS_VERSION: &str = "2012-11-30";
const SMIME_HEADER: &str = "\
MIME-Version:1.0
Content-Disposition: attachment; filename=/home/core/encrypted-ssh-cert.pem
Content-Type: application/x-pkcs7-mime; name=/home/core/encrypted-ssh-cert.pem
Content-Transfer-Encoding: base64

";

#[derive(Debug, Deserialize)]
struct GoalState {
    #[serde(rename = "Container")]
    pub container: Container
}

#[derive(Debug, Deserialize)]
struct Container {
    #[serde(rename = "RoleInstanceList")]
    pub role_instance_list: RoleInstanceList
}

#[derive(Debug, Deserialize)]
struct RoleInstanceList {
    #[serde(rename = "RoleInstance")]
    pub role_instances: Vec<RoleInstance>
}

#[derive(Debug, Deserialize)]
struct RoleInstance {
    #[serde(rename = "Configuration")]
    pub configuration: Configuration
}

#[derive(Debug, Deserialize)]
struct Configuration {
    #[serde(rename = "Certificates", default)]
    pub certificates: String
}

#[derive(Debug, Deserialize)]
struct CertificatesFile {
    #[serde(rename = "Data", default)]
    pub data: String
}

#[derive(Debug, Deserialize)]
struct Versions {
    #[serde(rename = "Supported")]
    pub supported: Supported
}

#[derive(Debug, Deserialize)]
struct Supported {
    #[serde(rename = "Version", default)]
    pub versions: Vec<String>
}

struct Azure {
    client: retry::Client,
    endpoint: IpAddr,
}

impl Azure {
    fn new() -> Result<Self> {
        let addr = Azure::get_fabric_address()
            .chain_err(|| "failed to get fabric address")?;
        let client = retry::Client::new()?
            .header(MSAgentName(MS_AGENT_NAME.to_owned()))
            .header(MSVersion(MS_VERSION.to_owned()));
        let azure = Azure {
            client: client,
            endpoint: addr,
        };
        azure.is_fabric_compatible(MS_VERSION)
            .chain_err(|| "failed version compatibility check")?;
        Ok(azure)
    }

    // I don't really understand why this is how we need to get this ip
    // address but this is how it works in the original implementation
    // and nobody complains about it, so w/e
    fn get_fabric_address() -> Result<IpAddr> {
        // get the interfaces on the machine
        let interfaces = pnet::datalink::interfaces();
        trace!("interfaces - {:?}", interfaces);

        for interface in interfaces {
            trace!("looking at interface {:?}", interface);
            let lease_path = format!("/run/systemd/netif/leases/{}", interface.index);
            let lease_path = Path::new(&lease_path);
            if lease_path.exists() {
                debug!("found lease file - {:?}", lease_path);
                let lease = File::open(&lease_path)
                    .chain_err(|| format!("failed to open lease file ({:?})", lease_path))?;
                let lease = BufReader::new(&lease);

                // find the OPTION_245 flag
                for line in lease.lines() {
                    let line = line
                        .chain_err(|| format!("failed to read from lease file ({:?})", lease_path))?;
                    let option: Vec<&str> = line.split('=').collect();
                    if option.len() > 1 && option[0] == OPTION_245 {
                        // value is an 8 digit hex value. convert it to u32 and
                        // then parse that into an ip. Ipv4Addr::from(u32)
                        // performs conversion from big-endian
                        trace!("found fabric address in hex - {:?}", option[1]);
                        let dec = u32::from_str_radix(option[1], 16)
                            .chain_err(|| format!("failed to convert '{}' from hex", option[1]))?;
                        return Ok(IpAddr::V4(Ipv4Addr::from(dec)));
                    }
                }

                debug!("failed to get fabric address from existing lease file '{:?}'", lease_path);
            }
        }

        Err("failed to retrieve fabric address".into())
    }

    fn is_fabric_compatible(&self, version: &str) -> Result<()> {
        let versions: Versions = self.client.get(retry::Xml, format!("http://{}/?comp=versions", self.endpoint)).send()
            .chain_err(|| "failed to get versions")?
            .ok_or_else(|| "failed to get versions: not found")?;

        if versions.supported.versions.iter().any(|v| v == version) {
            Ok(())
        } else {
            Err(format!("fabric version {} not compatible with fabric address {}", MS_VERSION, self.endpoint).into())
        }
    }

    fn get_certs_endpoint(&self) -> Result<String> {
        let goalstate: GoalState = self.client.get(retry::Xml, format!("http://{}/machine/?comp=goalstate", self.endpoint)).send()
            .chain_err(|| "failed to get goal state")?
            .ok_or_else(|| "failed to get goal state: not found response")?;

        // grab the certificates endpoint from the xml and return it
        let cert_endpoint: &str = &goalstate.container.role_instance_list.role_instances[0].configuration.certificates;
        Ok(String::from(cert_endpoint))
    }

    fn get_certs(&self, endpoint: String, mangled_pem: String) -> Result<String> {
        // get the certificates
        let certs: CertificatesFile = self.client.get(retry::Xml, endpoint)
            .header(MSCipherName("DES_EDE3_CBC".to_owned()))
            .header(MSCert(mangled_pem))
            .send()
            .chain_err(|| "failed to get certificates")?
            .ok_or_else(|| "failed to get certificates: not found")?;

        // the cms decryption expects it to have MIME information on the top
        // since cms is really for email attachments...don't tell the cops.
        let mut smime = String::from(SMIME_HEADER);
        smime.push_str(&certs.data);

        Ok(smime)
    }

    // put it all together
    fn get_ssh_pubkey(&self) -> Result<PublicKey> {
        // first we have to get the certificates endoint.
        let certs_endpoint = self.get_certs_endpoint()
            .chain_err(|| "failed to get certs endpoint")?;

        // we have to generate the rsa public/private keypair and the x509 cert
        // that we use to make the request. this is equivalent to
        // `openssl req -x509 -nodes -subj /CN=LinuxTransport -days 365 -newkey rsa:2048 -keyout private.pem -out cert.pem`
        let (x509, pkey) = x509::generate_cert(&x509::Config::new(2048, 365))
            .chain_err(|| "failed to generate keys")?;

        // mangle the pem file for the request
        let mangled_pem = crypto::mangle_pem(&x509)
            .chain_err(|| "failed to mangle pem")?;

        // fetch the encrypted cms blob from the certs endpoint
        let smime = self.get_certs(certs_endpoint, mangled_pem)
            .chain_err(|| "failed to get certs")?;

        // decrypt the cms blob
        let p12 = crypto::decrypt_cms(smime.as_bytes(), &pkey, &x509)
            .chain_err(|| "failed to decrypt cms blob")?;

        // convert that to the OpenSSH public key format
        let ssh_pubkey = crypto::p12_to_ssh_pubkey(&p12)
            .chain_err(|| "failed to convert pkcs12 blob to ssh pubkey")?;

        Ok(ssh_pubkey)
    }
}

pub fn fetch_metadata() -> Result<Metadata> {
    let provider = Azure::new()
        .chain_err(|| "azure: failed create metadata client")?;

    let ssh_pubkey = provider.get_ssh_pubkey()
        .chain_err(|| "azure: failed to get ssh pubkey")?;

    Ok(Metadata::builder()
       .add_publickeys(vec![ssh_pubkey])
       .build())
}
<commit_msg>providers/azure: wait until dhcp lease is aquired<commit_after>// Copyright 2017 CoreOS, Inc.
//
// 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
//
//     http://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.

//! azure metadata fetcher

mod crypto;

use self::crypto::x509;
use errors::*;
use metadata::Metadata;
use pnet;
use retry;
use openssh_keys::PublicKey;

use std::net::IpAddr;
use std::fs::File;
use std::path::Path;
use std::io::{BufRead, BufReader};
use std::time::Duration;
use std::thread;

header! {(MSAgentName, "x-ms-agent-name") => [String]}
header! {(MSVersion, "x-ms-version") => [String]}
header! {(MSCipherName, "x-ms-cipher-name") => [String]}
header! {(MSCert, "x-ms-guest-agent-public-x509-cert") => [String]}

const OPTION_245: &str = "OPTION_245";
const MS_AGENT_NAME: &str = "com.coreos.metadata";
const MS_VERSION: &str = "2012-11-30";
const SMIME_HEADER: &str = "\
MIME-Version:1.0
Content-Disposition: attachment; filename=/home/core/encrypted-ssh-cert.pem
Content-Type: application/x-pkcs7-mime; name=/home/core/encrypted-ssh-cert.pem
Content-Transfer-Encoding: base64

";

#[derive(Debug, Deserialize)]
struct GoalState {
    #[serde(rename = "Container")]
    pub container: Container
}

#[derive(Debug, Deserialize)]
struct Container {
    #[serde(rename = "RoleInstanceList")]
    pub role_instance_list: RoleInstanceList
}

#[derive(Debug, Deserialize)]
struct RoleInstanceList {
    #[serde(rename = "RoleInstance")]
    pub role_instances: Vec<RoleInstance>
}

#[derive(Debug, Deserialize)]
struct RoleInstance {
    #[serde(rename = "Configuration")]
    pub configuration: Configuration
}

#[derive(Debug, Deserialize)]
struct Configuration {
    #[serde(rename = "Certificates", default)]
    pub certificates: String
}

#[derive(Debug, Deserialize)]
struct CertificatesFile {
    #[serde(rename = "Data", default)]
    pub data: String
}

#[derive(Debug, Deserialize)]
struct Versions {
    #[serde(rename = "Supported")]
    pub supported: Supported
}

#[derive(Debug, Deserialize)]
struct Supported {
    #[serde(rename = "Version", default)]
    pub versions: Vec<String>
}

struct Azure {
    client: retry::Client,
    endpoint: IpAddr,
}

impl Azure {
    fn new() -> Result<Self> {
        let addr = Azure::get_fabric_address()
            .chain_err(|| "failed to get fabric address")?;
        let client = retry::Client::new()?
            .header(MSAgentName(MS_AGENT_NAME.to_owned()))
            .header(MSVersion(MS_VERSION.to_owned()));
        let azure = Azure {
            client: client,
            endpoint: addr,
        };
        azure.is_fabric_compatible(MS_VERSION)
            .chain_err(|| "failed version compatibility check")?;
        Ok(azure)
    }

    fn get_fabric_address() -> Result<IpAddr> {
        // get the interfaces on the machine
        let interfaces = pnet::datalink::interfaces();
        trace!("interfaces - {:?}", interfaces);

        // if we don't find the dhcp lease, keep trying
        // TODO(sdemos): eventually this should be a backoff timer instead of
        // looping forever.
        loop {
            for interface in interfaces.clone() {
                trace!("looking at interface {:?}", interface);
                let lease_path = format!("/run/systemd/netif/leases/{}", interface.index);
                let lease_path = Path::new(&lease_path);
                if lease_path.exists() {
                    debug!("found lease file - {:?}", lease_path);
                    let lease = File::open(&lease_path)
                        .chain_err(|| format!("failed to open lease file ({:?})", lease_path))?;
                    let lease = BufReader::new(&lease);

                    // find the OPTION_245 flag
                    for line in lease.lines() {
                        let line = line
                            .chain_err(|| format!("failed to read from lease file ({:?})", lease_path))?;
                        let option: Vec<&str> = line.split('=').collect();
                        if option.len() > 1 && option[0] == OPTION_245 {
                            // value is an 8 digit hex value. convert it to u32 and
                            // then parse that into an ip. Ipv4Addr::from(u32)
                            // performs conversion from big-endian
                            trace!("found fabric address in hex - {:?}", option[1]);
                            let dec = u32::from_str_radix(option[1], 16)
                                .chain_err(|| format!("failed to convert '{}' from hex", option[1]))?;
                            return Ok(IpAddr::V4(dec.into()));
                        }
                    }

                    debug!("failed to get fabric address from existing lease file '{:?}'", lease_path);
                }
            }

            // sleep for a bit
            thread::sleep(Duration::from_millis(500));
        }
        // Err("failed to retrieve fabric address".into())
    }

    fn is_fabric_compatible(&self, version: &str) -> Result<()> {
        let versions: Versions = self.client.get(retry::Xml, format!("http://{}/?comp=versions", self.endpoint)).send()
            .chain_err(|| "failed to get versions")?
            .ok_or_else(|| "failed to get versions: not found")?;

        if versions.supported.versions.iter().any(|v| v == version) {
            Ok(())
        } else {
            Err(format!("fabric version {} not compatible with fabric address {}", MS_VERSION, self.endpoint).into())
        }
    }

    fn get_certs_endpoint(&self) -> Result<String> {
        let goalstate: GoalState = self.client.get(retry::Xml, format!("http://{}/machine/?comp=goalstate", self.endpoint)).send()
            .chain_err(|| "failed to get goal state")?
            .ok_or_else(|| "failed to get goal state: not found response")?;

        // grab the certificates endpoint from the xml and return it
        let cert_endpoint: &str = &goalstate.container.role_instance_list.role_instances[0].configuration.certificates;
        Ok(String::from(cert_endpoint))
    }

    fn get_certs(&self, endpoint: String, mangled_pem: String) -> Result<String> {
        // get the certificates
        let certs: CertificatesFile = self.client.get(retry::Xml, endpoint)
            .header(MSCipherName("DES_EDE3_CBC".to_owned()))
            .header(MSCert(mangled_pem))
            .send()
            .chain_err(|| "failed to get certificates")?
            .ok_or_else(|| "failed to get certificates: not found")?;

        // the cms decryption expects it to have MIME information on the top
        // since cms is really for email attachments...don't tell the cops.
        let mut smime = String::from(SMIME_HEADER);
        smime.push_str(&certs.data);

        Ok(smime)
    }

    // put it all together
    fn get_ssh_pubkey(&self) -> Result<PublicKey> {
        // first we have to get the certificates endoint.
        let certs_endpoint = self.get_certs_endpoint()
            .chain_err(|| "failed to get certs endpoint")?;

        // we have to generate the rsa public/private keypair and the x509 cert
        // that we use to make the request. this is equivalent to
        // `openssl req -x509 -nodes -subj /CN=LinuxTransport -days 365 -newkey rsa:2048 -keyout private.pem -out cert.pem`
        let (x509, pkey) = x509::generate_cert(&x509::Config::new(2048, 365))
            .chain_err(|| "failed to generate keys")?;

        // mangle the pem file for the request
        let mangled_pem = crypto::mangle_pem(&x509)
            .chain_err(|| "failed to mangle pem")?;

        // fetch the encrypted cms blob from the certs endpoint
        let smime = self.get_certs(certs_endpoint, mangled_pem)
            .chain_err(|| "failed to get certs")?;

        // decrypt the cms blob
        let p12 = crypto::decrypt_cms(smime.as_bytes(), &pkey, &x509)
            .chain_err(|| "failed to decrypt cms blob")?;

        // convert that to the OpenSSH public key format
        let ssh_pubkey = crypto::p12_to_ssh_pubkey(&p12)
            .chain_err(|| "failed to convert pkcs12 blob to ssh pubkey")?;

        Ok(ssh_pubkey)
    }
}

pub fn fetch_metadata() -> Result<Metadata> {
    let provider = Azure::new()
        .chain_err(|| "azure: failed create metadata client")?;

    let ssh_pubkey = provider.get_ssh_pubkey()
        .chain_err(|| "azure: failed to get ssh pubkey")?;

    Ok(Metadata::builder()
       .add_publickeys(vec![ssh_pubkey])
       .build())
}
<|endoftext|>


<commit_before>use getopts::{optflag,optopt,getopts,usage};
use std::os;
use std::collections::HashMap;
use std::io;
use std::io::fs::PathExtensions;
use std::borrow::ToOwned;
use std::str;

use vobject::parse_component;

macro_rules! main_try {
    ($result: expr, $errmsg: expr) => (
        match $result {
            Ok(m) => m,
            Err(e) => {
                println!("{}: {}", $errmsg, e);
                os::set_exit_status(1);
                return;
            }
        }
    )
}

fn get_env() -> HashMap<String, String> {
    let mut env = HashMap::new();
    env.extend(os::env().into_iter().filter(|&(ref key, ref value)| {
        key.as_slice().starts_with("MATES_") && value.len() > 0
    }));
    env
}

fn expect_env<'a>(env: &'a HashMap<String, String>, key: &str) -> &'a String {
    env.get(key).expect(
        format!("The {} environment variable must be set.", key).as_slice()
    )
}

fn build_index(outfile: &Path, dir: &Path) -> io::IoResult<()> {
    if !dir.is_dir() {
        return Err(io::IoError {
            kind: io::MismatchedFileTypeForOperation,
            desc: "MATES_DIR must be a directory.",
            detail: None
        });
    };

    let mut outf = io::File::create(outfile);
    let entries = try!(io::fs::readdir(dir));
    for entry in entries.iter() {
        if !entry.is_file() {
            continue;
        }

        print!("Processing {}\n", entry.display());

        let itemstr = try!(io::File::open(entry).read_to_string());
        let item = match parse_component(&itemstr) {
            Ok(item) => item,
            Err(e) => {
                println!("Error: Failed to parse item {}: {}\n", entry.display(), e);
                os::set_exit_status(1);
                continue;
            }
        };

        let name = match item.single_prop("FN") {
            Some(name) => name.get_raw_value(),
            None => {
                print!("Warning: No name in {}, skipping.\n", entry.display());
                continue;
            }
        };

        let emails = item.all_props("EMAIL");
        for email in emails.iter() {
            try!(outf.write_str(
                format!("{}\t{}\t{}\n", email.get_raw_value(), name, entry.display()).as_slice()
            ))
        };
    };
    return Ok(());
}

pub fn cli_main() {
    let args = os::args();

    let program = args[0].as_slice();
    let opts = [
        optflag("i", "index", "Rewrite/create index."),
        optflag("h", "help", "Print help."),
        optopt("", "mutt-query", "Search for contact, for mutt's query_command.", ""),
        optopt("", "file-query", "Search for contact, return filepaths.", ""),
        optopt("", "email-query", "Search for contact, return \"name <email>\".", "")
    ];

    let matches = main_try!(getopts(args.tail(), &opts), "Failed to parse arguments");

    let env = get_env();

    let print_usage = |&:| {
        println!("{}", usage(program, &opts));
    };

    if matches.opt_present("h") {
        print_usage();

    } else if matches.opt_present("index") {
        let index_file = expect_env(&env, "MATES_INDEX");
        let mates_dir = expect_env(&env, "MATES_DIR");
        println!("Rebuilding index file \"{}\"...", index_file);
        main_try!(build_index(
            &Path::new(index_file.as_slice()),
            &Path::new(mates_dir.as_slice())
        ), "Failed to build index");

    } else if matches.opt_present("mutt-query") {
        // FIXME: Better way to write this? We already checked for presence of mutt-search before
        let query = matches.opt_str("mutt-query").expect("This should never happen and yet it did.");
        main_try!(mutt_query(env, query), "Failed to execute grep");
    } else if matches.opt_present("file-query") {
        // FIXME: Better way to write this? We already checked for presence of mutt-search before
        let query = matches.opt_str("file-query").expect("This should never happen and yet it did.");
        main_try!(file_query(env, query), "Failed to execute grep");
    } else if matches.opt_present("email-query") {
        // FIXME: Better way to write this? We already checked for presence of mutt-search before
        let query = matches.opt_str("email-query").expect("This should never happen and yet it did.");
        main_try!(email_query(env, query), "Failed to execute grep");
    } else {
        print_usage();
    };
}

fn mutt_query<'a>(env: HashMap<String, String>, query: String) -> io::IoResult<()> {
    println!("");  // For some reason mutt requires an empty line
    for item in try!(index_query(env, query)) {
        if(item.email.len() > 0 && item.name.len() > 0) {
            println!("{}\t{}\t{}", item.email, item.name, item.filepath);
        };
    };
    Ok(())
}

fn file_query<'a>(env: HashMap<String, String>, query: String) -> io::IoResult<()> {
    for item in try!(index_query(env, query)) {
        if(item.filepath.len() > 0) {
            println!("{}", item.filepath)
        };
    };
    Ok(())
}

fn email_query<'a>(env: HashMap<String, String>, query: String) -> io::IoResult<()> {
    for item in try!(index_query(env, query)) {
        if(item.name.len() > 0 && item.email.len() > 0) {
            println!("{} <{}>", item.name, item.email)
        };
    };
    Ok(())
}

fn index_query<'a>(env: HashMap<String, String>, query: String) -> io::IoResult<IndexIterator<'a>> {
    let index_file = expect_env(&env, "MATES_INDEX");
    let default_grep = "grep".to_owned();
    let grep_cmd = match env.get("MATES_GREP") {
        Some(x) => x,
        None => &default_grep
    };

    let mut process = try!(io::Command::new(grep_cmd.as_slice())
                           .arg(query.as_slice())
                           .arg(index_file.as_slice())
                           .stderr(io::process::InheritFd(2))
                           .spawn());

    let stream = match process.stdout.as_mut() {
        Some(x) => x,
        None => return Err(io::IoError {
            kind: io::IoUnavailable,
            desc: "Failed to get stdout from grep process.",
            detail: None
        })
    };

    let output = try!(stream.read_to_string());
    Ok(IndexIterator::new(&output))
}

struct IndexItem<'a> {
    pub email: String,
    pub name: String,
    pub filepath: String
}

impl<'a> IndexItem<'a> {
    fn new(line: String) -> IndexItem<'a> {
        let mut parts = line.split('\t');

        IndexItem {
            email: parts.next().unwrap_or("").to_string(),
            name: parts.next().unwrap_or("").to_string(),
            filepath: parts.next().unwrap_or("").to_string()
        }
    }
}

struct IndexIterator<'a> {
    linebuffer: Vec<String>
}

impl<'a> IndexIterator<'a> {
    fn new(output: &String) -> IndexIterator<'a> {

        let rv = output.split('\n').map(|x: &str| x.to_string()).collect();
        IndexIterator {
            linebuffer: rv
        }
    }
}

impl<'a> Iterator for IndexIterator<'a> {
    type Item = IndexItem<'a>;

    fn next(&mut self) -> Option<IndexItem<'a>> {
        match self.linebuffer.pop() {
            Some(x) => Some(IndexItem::new(x)),
            None => None
        }
    }
}
<commit_msg>whoops<commit_after>use getopts::{optflag,optopt,getopts,usage};
use std::os;
use std::collections::HashMap;
use std::io;
use std::io::fs::PathExtensions;
use std::borrow::ToOwned;
use std::str;

use vobject::parse_component;

macro_rules! main_try {
    ($result: expr, $errmsg: expr) => (
        match $result {
            Ok(m) => m,
            Err(e) => {
                println!("{}: {}", $errmsg, e);
                os::set_exit_status(1);
                return;
            }
        }
    )
}

fn get_env() -> HashMap<String, String> {
    let mut env = HashMap::new();
    env.extend(os::env().into_iter().filter(|&(ref key, ref value)| {
        key.as_slice().starts_with("MATES_") && value.len() > 0
    }));
    env
}

fn expect_env<'a>(env: &'a HashMap<String, String>, key: &str) -> &'a String {
    env.get(key).expect(
        format!("The {} environment variable must be set.", key).as_slice()
    )
}

fn build_index(outfile: &Path, dir: &Path) -> io::IoResult<()> {
    if !dir.is_dir() {
        return Err(io::IoError {
            kind: io::MismatchedFileTypeForOperation,
            desc: "MATES_DIR must be a directory.",
            detail: None
        });
    };

    let mut outf = io::File::create(outfile);
    let entries = try!(io::fs::readdir(dir));
    for entry in entries.iter() {
        if !entry.is_file() {
            continue;
        }

        print!("Processing {}\n", entry.display());

        let itemstr = try!(io::File::open(entry).read_to_string());
        let item = match parse_component(&itemstr) {
            Ok(item) => item,
            Err(e) => {
                println!("Error: Failed to parse item {}: {}\n", entry.display(), e);
                os::set_exit_status(1);
                continue;
            }
        };

        let name = match item.single_prop("FN") {
            Some(name) => name.get_raw_value(),
            None => {
                print!("Warning: No name in {}, skipping.\n", entry.display());
                continue;
            }
        };

        let emails = item.all_props("EMAIL");
        for email in emails.iter() {
            try!(outf.write_str(
                format!("{}\t{}\t{}\n", email.get_raw_value(), name, entry.display()).as_slice()
            ))
        };
    };
    return Ok(());
}

pub fn cli_main() {
    let args = os::args();

    let program = args[0].as_slice();
    let opts = [
        optflag("i", "index", "Rewrite/create index."),
        optflag("h", "help", "Print help."),
        optopt("", "mutt-query", "Search for contact, for mutt's query_command.", ""),
        optopt("", "file-query", "Search for contact, return filepaths.", ""),
        optopt("", "email-query", "Search for contact, return \"name <email>\".", "")
    ];

    let matches = main_try!(getopts(args.tail(), &opts), "Failed to parse arguments");

    let env = get_env();

    let print_usage = |&:| {
        println!("{}", usage(program, &opts));
        println!("Environment variables:");
        println!("- MATES_INDEX: Path to index file, which is basically a cache of all");
        println!("               contacts.");
        println!("- MATES_DIR:   The vdir to use.");
        println!("- MATES_GREP:  The grep executable to use.");
    };

    if matches.opt_present("h") {
        print_usage();

    } else if matches.opt_present("index") {
        let index_file = expect_env(&env, "MATES_INDEX");
        let mates_dir = expect_env(&env, "MATES_DIR");
        println!("Rebuilding index file \"{}\"...", index_file);
        main_try!(build_index(
            &Path::new(index_file.as_slice()),
            &Path::new(mates_dir.as_slice())
        ), "Failed to build index");

    } else if matches.opt_present("mutt-query") {
        // FIXME: Better way to write this? We already checked for presence of mutt-search before
        let query = matches.opt_str("mutt-query").expect("This should never happen and yet it did.");
        main_try!(mutt_query(env, query), "Failed to execute grep");
    } else if matches.opt_present("file-query") {
        // FIXME: Better way to write this? We already checked for presence of mutt-search before
        let query = matches.opt_str("file-query").expect("This should never happen and yet it did.");
        main_try!(file_query(env, query), "Failed to execute grep");
    } else if matches.opt_present("email-query") {
        // FIXME: Better way to write this? We already checked for presence of mutt-search before
        let query = matches.opt_str("email-query").expect("This should never happen and yet it did.");
        main_try!(email_query(env, query), "Failed to execute grep");
    } else {
        print_usage();
    };
}

fn mutt_query<'a>(env: HashMap<String, String>, query: String) -> io::IoResult<()> {
    println!("");  // For some reason mutt requires an empty line
    for item in try!(index_query(env, query)) {
        if(item.email.len() > 0 && item.name.len() > 0) {
            println!("{}\t{}\t{}", item.email, item.name, item.filepath);
        };
    };
    Ok(())
}

fn file_query<'a>(env: HashMap<String, String>, query: String) -> io::IoResult<()> {
    for item in try!(index_query(env, query)) {
        if(item.filepath.len() > 0) {
            println!("{}", item.filepath)
        };
    };
    Ok(())
}

fn email_query<'a>(env: HashMap<String, String>, query: String) -> io::IoResult<()> {
    for item in try!(index_query(env, query)) {
        if(item.name.len() > 0 && item.email.len() > 0) {
            println!("{} <{}>", item.name, item.email)
        };
    };
    Ok(())
}

fn index_query<'a>(env: HashMap<String, String>, query: String) -> io::IoResult<IndexIterator<'a>> {
    let index_file = expect_env(&env, "MATES_INDEX");
    let default_grep = "grep".to_owned();
    let grep_cmd = match env.get("MATES_GREP") {
        Some(x) => x,
        None => &default_grep
    };

    let mut process = try!(io::Command::new(grep_cmd.as_slice())
                           .arg(query.as_slice())
                           .arg(index_file.as_slice())
                           .stderr(io::process::InheritFd(2))
                           .spawn());

    let stream = match process.stdout.as_mut() {
        Some(x) => x,
        None => return Err(io::IoError {
            kind: io::IoUnavailable,
            desc: "Failed to get stdout from grep process.",
            detail: None
        })
    };

    let output = try!(stream.read_to_string());
    Ok(IndexIterator::new(&output))
}

struct IndexItem<'a> {
    pub email: String,
    pub name: String,
    pub filepath: String
}

impl<'a> IndexItem<'a> {
    fn new(line: String) -> IndexItem<'a> {
        let mut parts = line.split('\t');

        IndexItem {
            email: parts.next().unwrap_or("").to_string(),
            name: parts.next().unwrap_or("").to_string(),
            filepath: parts.next().unwrap_or("").to_string()
        }
    }
}

struct IndexIterator<'a> {
    linebuffer: Vec<String>
}

impl<'a> IndexIterator<'a> {
    fn new(output: &String) -> IndexIterator<'a> {

        let rv = output.split('\n').map(|x: &str| x.to_string()).collect();
        IndexIterator {
            linebuffer: rv
        }
    }
}

impl<'a> Iterator for IndexIterator<'a> {
    type Item = IndexItem<'a>;

    fn next(&mut self) -> Option<IndexItem<'a>> {
        match self.linebuffer.pop() {
            Some(x) => Some(IndexItem::new(x)),
            None => None
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2019 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use slog::{Drain, Key, OwnedKVList, Record, KV};
use std::{fmt, io, sync::Mutex};

#[derive(Debug, Default)]
pub struct MemDrain {
    buffer: Mutex<Vec<String>>,
    logging: bool,
}

impl MemDrain {
    pub fn clear(&self) {
        self.buffer.lock().unwrap().clear();
    }

    pub fn start_logging(&mut self) {
        self.logging = true;
    }

    pub fn stop_logging(&mut self) {
        self.logging = false;
    }

    pub fn to_string(&self) -> String {
        let buffer = self.buffer.lock().unwrap();
        buffer.join("\n")
    }
}

impl Drain for MemDrain {
    type Err = io::Error;
    type Ok = ();

    fn log(&self, record: &Record, kvs: &OwnedKVList) -> io::Result<()> {
        if self.logging {
            let mut buffer = self.buffer.lock().unwrap();
            const TIMESTAMP_FORMAT: &'static str = "%b %d %H:%M:%S%.3f";
            let mut serializer = Serializer::new(fmt::format(*record.msg()));

            record.kv().serialize(record, &mut serializer)?;
            kvs.serialize(record, &mut serializer)?;
            let line = format!(
                "{} {} {}",
                chrono::Local::now().format(TIMESTAMP_FORMAT),
                record.level().as_short_str(),
                serializer.text,
            );

            buffer.push(line);
        }
        Ok(())
    }
}

struct Serializer {
    text: String,
}

impl Serializer {
    fn new(text: String) -> Self {
        Serializer { text }
    }
}

impl slog::ser::Serializer for Serializer {
    fn emit_arguments(&mut self, key: Key, val: &fmt::Arguments) -> slog::Result {
        self.text = self.text.replace(key, &format!("{:?}", val));
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use slog::{o, slog_debug, slog_info, Logger};
    use std::sync::Arc;

    #[test]
    fn drain_storage_log() {
        let s1 = "Multiple log messages should";
        let s2 = "all be find inside log string";
        let drain = Arc::new(Mutex::new(MemDrain::default()));
        let r_vec = drain.clone();
        drain.lock().unwrap().start_logging();
        let log = Logger::root(drain.fuse(), o!());
        slog_info!(log, "{}", s1);
        slog_info!(log, "{}", s2);
        let result = format!("{}", r_vec.lock().unwrap().to_string());
        assert!(result.contains(s1));
        assert!(result.contains(s2));
    }

    #[test]
    fn drain_format() {
        let s1 = "Log should contain message type";
        let s2 = "Type strings are shorten";
        let drain = Arc::new(Mutex::new(MemDrain::default()));
        let r_vec = drain.clone();
        drain.lock().unwrap().start_logging();
        let log = Logger::root(drain.fuse(), o!());
        slog_info!(log, "{}", s1);
        slog_debug!(log, "{}", s2);
        let result = format!("{}", r_vec.lock().unwrap().to_string());
        assert!(result.contains("INFO"));
        assert!(result.contains("DEBG"));
    }

    #[test]
    fn drain_key_values() {
        let txt = "Key values should be swapped, LOGGER and RECORD";
        let logger_value = "when defined on logger";
        let macro_value = "when defined on record";
        let drain = Arc::new(Mutex::new(MemDrain::default()));
        let r_vec = drain.clone();
        drain.lock().unwrap().start_logging();
        let log = Logger::root(drain.fuse(), o!("LOGGER" => logger_value));
        slog_info!(log, "{}", txt; "RECORD" => macro_value);
        let result = format!("{}", r_vec.lock().unwrap().to_string());
        assert!(result.contains(logger_value));
        assert!(result.contains(macro_value));
    }
}
<commit_msg>mem_drain: Avoid explicit `'static' lifetime for `const &str`<commit_after>// Copyright (C) 2019 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use slog::{Drain, Key, OwnedKVList, Record, KV};
use std::{fmt, io, sync::Mutex};

#[derive(Debug, Default)]
pub struct MemDrain {
    buffer: Mutex<Vec<String>>,
    logging: bool,
}

impl MemDrain {
    pub fn clear(&self) {
        self.buffer.lock().unwrap().clear();
    }

    pub fn start_logging(&mut self) {
        self.logging = true;
    }

    pub fn stop_logging(&mut self) {
        self.logging = false;
    }

    pub fn to_string(&self) -> String {
        let buffer = self.buffer.lock().unwrap();
        buffer.join("\n")
    }
}

impl Drain for MemDrain {
    type Err = io::Error;
    type Ok = ();

    fn log(&self, record: &Record, kvs: &OwnedKVList) -> io::Result<()> {
        if self.logging {
            let mut buffer = self.buffer.lock().unwrap();
            const TIMESTAMP_FORMAT: &str = "%b %d %H:%M:%S%.3f";
            let mut serializer = Serializer::new(fmt::format(*record.msg()));

            record.kv().serialize(record, &mut serializer)?;
            kvs.serialize(record, &mut serializer)?;
            let line = format!(
                "{} {} {}",
                chrono::Local::now().format(TIMESTAMP_FORMAT),
                record.level().as_short_str(),
                serializer.text,
            );

            buffer.push(line);
        }
        Ok(())
    }
}

struct Serializer {
    text: String,
}

impl Serializer {
    fn new(text: String) -> Self {
        Serializer { text }
    }
}

impl slog::ser::Serializer for Serializer {
    fn emit_arguments(&mut self, key: Key, val: &fmt::Arguments) -> slog::Result {
        self.text = self.text.replace(key, &format!("{:?}", val));
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use slog::{o, slog_debug, slog_info, Logger};
    use std::sync::Arc;

    #[test]
    fn drain_storage_log() {
        let s1 = "Multiple log messages should";
        let s2 = "all be find inside log string";
        let drain = Arc::new(Mutex::new(MemDrain::default()));
        let r_vec = drain.clone();
        drain.lock().unwrap().start_logging();
        let log = Logger::root(drain.fuse(), o!());
        slog_info!(log, "{}", s1);
        slog_info!(log, "{}", s2);
        let result = format!("{}", r_vec.lock().unwrap().to_string());
        assert!(result.contains(s1));
        assert!(result.contains(s2));
    }

    #[test]
    fn drain_format() {
        let s1 = "Log should contain message type";
        let s2 = "Type strings are shorten";
        let drain = Arc::new(Mutex::new(MemDrain::default()));
        let r_vec = drain.clone();
        drain.lock().unwrap().start_logging();
        let log = Logger::root(drain.fuse(), o!());
        slog_info!(log, "{}", s1);
        slog_debug!(log, "{}", s2);
        let result = format!("{}", r_vec.lock().unwrap().to_string());
        assert!(result.contains("INFO"));
        assert!(result.contains("DEBG"));
    }

    #[test]
    fn drain_key_values() {
        let txt = "Key values should be swapped, LOGGER and RECORD";
        let logger_value = "when defined on logger";
        let macro_value = "when defined on record";
        let drain = Arc::new(Mutex::new(MemDrain::default()));
        let r_vec = drain.clone();
        drain.lock().unwrap().start_logging();
        let log = Logger::root(drain.fuse(), o!("LOGGER" => logger_value));
        slog_info!(log, "{}", txt; "RECORD" => macro_value);
        let result = format!("{}", r_vec.lock().unwrap().to_string());
        assert!(result.contains(logger_value));
        assert!(result.contains(macro_value));
    }
}
<|endoftext|>


<commit_before>extern crate uuid;
use coin_cbc;

use dsl::LpExpression::*;
use dsl::*;
use solvers::{Solution, SolverTrait, Status, WithMaxSeconds, WithNbThreads};
use std::collections::HashMap;

/// Solver that calls cbc through [rust bindings](https://github.com/KardinalAI/coin_cbc)
#[derive(Debug, Clone, Default)]
pub struct NativeCbcSolver {
    name: String,
    threads: Option<u32>,
    seconds: Option<u32>,
}

impl NativeCbcSolver {
    pub fn new() -> NativeCbcSolver {
        NativeCbcSolver {
            name: "CbcNative".to_string(),
            threads: None,
            seconds: None,
        }
    }
}

impl WithMaxSeconds<NativeCbcSolver> for NativeCbcSolver {
    fn max_seconds(&self) -> Option<u32> {
        self.seconds
    }
    fn with_max_seconds(&self, seconds: u32) -> NativeCbcSolver {
        NativeCbcSolver {
            seconds: Some(seconds),
            ..self.clone()
        }
    }
}
impl WithNbThreads<NativeCbcSolver> for NativeCbcSolver {
    fn nb_threads(&self) -> Option<u32> {
        self.threads
    }
    fn with_nb_threads(&self, threads: u32) -> NativeCbcSolver {
        NativeCbcSolver {
            threads: Some(threads),
            ..self.clone()
        }
    }
}

/// Recursively unwrap an expression in a Vec of variables and literals.
fn var_lit(expr: &LpExpression, lst: &mut Vec<(String, f32)>, mul: Option<f32>) {
    let mul = match mul {
        Some(lit) => lit,
        None => 1.,
    };
    match expr {
        &ConsBin(LpBinary { ref name, .. })
        | &ConsInt(LpInteger { ref name, .. })
        | &ConsCont(LpContinuous { ref name, .. }) => {
            lst.push((name.clone(), mul * split_constant_and_expr(expr).0));
        }

        MulExpr(val, ref e) => match **e {
            ConsBin(LpBinary { ref name, .. })
            | ConsInt(LpInteger { ref name, .. })
            | ConsCont(LpContinuous { ref name, .. }) => {
                if let LitVal(lit) = *val.clone() {
                    lst.push((name.clone(), mul * lit))
                }
            }
            MulExpr(..) | AddExpr(..) | SubExpr(..) => {
                let next_mul = match *val.clone() {
                    LitVal(lit) => Some(lit * mul),
                    _ => None,
                };
                var_lit(&*e, lst, next_mul)
            }
            _ => (),
        },
        &AddExpr(ref e1, ref e2) => {
            var_lit(&*e1, lst, None);
            var_lit(&*e2, lst, None);
        }
        &SubExpr(ref e1, ref e2) => {
            var_lit(&*e1, lst, None);
            var_lit(&*e2, lst, Some(-1.));
        }
        _ => (),
    }
}

fn always_literal(expr: &LpExpression) -> f64 {
    match *expr {
        LitVal(num) => num as f64,
        _ => panic!("wrong generalization"),
    }
}

fn add_variable(m: &mut coin_cbc::Model, expr: &LpExpression) -> coin_cbc::Col {
    match expr {
        ConsInt(LpInteger {
            name: _,
            lower_bound,
            upper_bound,
        }) => {
            let col = m.add_integer();
            if let Some(lb) = lower_bound {
                m.set_col_lower(col, *lb as f64)
            }
            if let Some(ub) = upper_bound {
                m.set_col_lower(col, *ub as f64)
            }
            col
        }
        ConsCont(LpContinuous {
            name: _,
            lower_bound,
            upper_bound,
        }) => {
            let col = m.add_col();
            if let Some(lb) = lower_bound {
                m.set_col_lower(col, *lb as f64)
            }
            if let Some(ub) = upper_bound {
                m.set_col_lower(col, *ub as f64)
            }
            col
        }
        ConsBin(_) => m.add_binary(),
        _ => panic!("Unexpected LpExpression on LpProblem.variables()!"),
    }
}

impl SolverTrait for NativeCbcSolver {
    type P = LpProblem;

    fn run<'a>(&self, problem: &'a Self::P) -> Result<Solution<'a>, String> {
        let mut m = coin_cbc::Model::default();
        // columns (variables)
        let cols: HashMap<String, coin_cbc::Col> = problem
            .variables()
            .iter()
            .map(|(name, expr)| (name.clone(), add_variable(&mut m, expr)))
            .collect();
        // rows (constraints)
        problem.constraints.iter().for_each(|cons| {
            let row = m.add_row();
            let general = cons.generalize();
            match general.1 {
                Constraint::GreaterOrEqual => m.set_row_lower(row, always_literal(&general.2)),
                Constraint::LessOrEqual => m.set_row_upper(row, always_literal(&general.2)),
                Constraint::Equal => m.set_row_equal(row, always_literal(&general.2)),
            }
            let mut lst: Vec<_> = Vec::new();
            var_lit(&general.0, &mut lst, None);
            lst.iter()
                .for_each(|(n, lit)| m.set_weight(row, cols[n], *lit as f64));
        });
        // objective
        if let Some(objective) = &problem.obj_expr {
            let mut lst: Vec<_> = Vec::new();
            var_lit(&objective, &mut lst, None);
            lst.iter()
                .for_each(|(n, lit)| m.set_obj_coeff(cols[n], *lit as f64))
        }
        m.set_obj_sense(match problem.objective_type {
            LpObjective::Maximize => coin_cbc::Sense::Maximize,
            LpObjective::Minimize => coin_cbc::Sense::Minimize,
        });

        let sol = m.solve();

        let mut lst: Vec<_> = Vec::new();
        var_lit(&(problem.obj_expr.clone().unwrap()), &mut lst, None);

        Ok(Solution {
            status: match sol.raw().status() {
                coin_cbc::raw::Status::Finished => Status::Optimal,
                coin_cbc::raw::Status::Abandoned => Status::Infeasible,
                _ => Status::NotSolved,
            },
            results: cols
                .iter()
                .map(|(name, col)| (name.to_owned(), sol.col(*col) as f32))
                .collect(),
            related_problem: Some(problem),
        })
    }
}
<commit_msg>Fix adding upper bounds in Native Cbc<commit_after>extern crate uuid;
use coin_cbc;

use dsl::LpExpression::*;
use dsl::*;
use solvers::{Solution, SolverTrait, Status, WithMaxSeconds, WithNbThreads};
use std::collections::HashMap;

/// Solver that calls cbc through [rust bindings](https://github.com/KardinalAI/coin_cbc)
#[derive(Debug, Clone, Default)]
pub struct NativeCbcSolver {
    name: String,
    threads: Option<u32>,
    seconds: Option<u32>,
}

impl NativeCbcSolver {
    pub fn new() -> NativeCbcSolver {
        NativeCbcSolver {
            name: "CbcNative".to_string(),
            threads: None,
            seconds: None,
        }
    }
}

impl WithMaxSeconds<NativeCbcSolver> for NativeCbcSolver {
    fn max_seconds(&self) -> Option<u32> {
        self.seconds
    }
    fn with_max_seconds(&self, seconds: u32) -> NativeCbcSolver {
        NativeCbcSolver {
            seconds: Some(seconds),
            ..self.clone()
        }
    }
}
impl WithNbThreads<NativeCbcSolver> for NativeCbcSolver {
    fn nb_threads(&self) -> Option<u32> {
        self.threads
    }
    fn with_nb_threads(&self, threads: u32) -> NativeCbcSolver {
        NativeCbcSolver {
            threads: Some(threads),
            ..self.clone()
        }
    }
}

/// Recursively unwrap an expression in a Vec of variables and literals.
fn var_lit(expr: &LpExpression, lst: &mut Vec<(String, f32)>, mul: Option<f32>) {
    let mul = match mul {
        Some(lit) => lit,
        None => 1.,
    };
    match expr {
        &ConsBin(LpBinary { ref name, .. })
        | &ConsInt(LpInteger { ref name, .. })
        | &ConsCont(LpContinuous { ref name, .. }) => {
            lst.push((name.clone(), mul * split_constant_and_expr(expr).0));
        }

        MulExpr(val, ref e) => match **e {
            ConsBin(LpBinary { ref name, .. })
            | ConsInt(LpInteger { ref name, .. })
            | ConsCont(LpContinuous { ref name, .. }) => {
                if let LitVal(lit) = *val.clone() {
                    lst.push((name.clone(), mul * lit))
                }
            }
            MulExpr(..) | AddExpr(..) | SubExpr(..) => {
                let next_mul = match *val.clone() {
                    LitVal(lit) => Some(lit * mul),
                    _ => None,
                };
                var_lit(&*e, lst, next_mul)
            }
            _ => (),
        },
        &AddExpr(ref e1, ref e2) => {
            var_lit(&*e1, lst, None);
            var_lit(&*e2, lst, None);
        }
        &SubExpr(ref e1, ref e2) => {
            var_lit(&*e1, lst, None);
            var_lit(&*e2, lst, Some(-1.));
        }
        _ => (),
    }
}

fn always_literal(expr: &LpExpression) -> f64 {
    match *expr {
        LitVal(num) => num as f64,
        _ => panic!("wrong generalization"),
    }
}

fn add_variable(m: &mut coin_cbc::Model, expr: &LpExpression) -> coin_cbc::Col {
    match expr {
        ConsInt(LpInteger {
            name: _,
            lower_bound,
            upper_bound,
        }) => {
            let col = m.add_integer();
            if let Some(lb) = lower_bound {
                m.set_col_lower(col, *lb as f64)
            }
            if let Some(ub) = upper_bound {
                m.set_col_upper(col, *ub as f64)
            }
            col
        }
        ConsCont(LpContinuous {
            name: _,
            lower_bound,
            upper_bound,
        }) => {
            let col = m.add_col();
            if let Some(lb) = lower_bound {
                m.set_col_lower(col, *lb as f64)
            }
            if let Some(ub) = upper_bound {
                m.set_col_upper(col, *ub as f64)
            }
            col
        }
        ConsBin(_) => m.add_binary(),
        _ => panic!("Unexpected LpExpression on LpProblem.variables()!"),
    }
}

impl SolverTrait for NativeCbcSolver {
    type P = LpProblem;

    fn run<'a>(&self, problem: &'a Self::P) -> Result<Solution<'a>, String> {
        let mut m = coin_cbc::Model::default();
        // columns (variables)
        let cols: HashMap<String, coin_cbc::Col> = problem
            .variables()
            .iter()
            .map(|(name, expr)| (name.clone(), add_variable(&mut m, expr)))
            .collect();
        // rows (constraints)
        problem.constraints.iter().for_each(|cons| {
            let row = m.add_row();
            let general = cons.generalize();
            match general.1 {
                Constraint::GreaterOrEqual => m.set_row_lower(row, always_literal(&general.2)),
                Constraint::LessOrEqual => m.set_row_upper(row, always_literal(&general.2)),
                Constraint::Equal => m.set_row_equal(row, always_literal(&general.2)),
            }
            let mut lst: Vec<_> = Vec::new();
            var_lit(&general.0, &mut lst, None);
            lst.iter()
                .for_each(|(n, lit)| m.set_weight(row, cols[n], *lit as f64));
        });
        // objective
        if let Some(objective) = &problem.obj_expr {
            let mut lst: Vec<_> = Vec::new();
            var_lit(&objective, &mut lst, None);
            lst.iter().for_each(|(n, lit)| m.set_obj_coeff(cols[n], *lit as f64))
        }
        m.set_obj_sense(match problem.objective_type {
            LpObjective::Maximize => coin_cbc::Sense::Maximize,
            LpObjective::Minimize => coin_cbc::Sense::Minimize,
        });

        let sol = m.solve();

        Ok(Solution {
            status: match sol.raw().status() {
                coin_cbc::raw::Status::Finished => Status::Optimal,
                coin_cbc::raw::Status::Abandoned => Status::Infeasible,
                _ => Status::NotSolved,
            },
            results: cols
                .iter()
                .map(|(name, col)| (name.to_owned(), sol.col(*col) as f32))
                .collect(),
            related_problem: Some(problem),
        })
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Helper functions for implementing `RngCore` functions.
//!
//! For cross-platform reproducibility, these functions all use Little Endian:
//! least-significant part first. For example, `next_u64_via_u32` takes `u32`
//! values `x, y`, then outputs `(y << 32) | x`. To implement `next_u32`
//! from `next_u64` in little-endian order, one should use `next_u64() as u32`.
//!
//! Byte-swapping (like the std `to_le` functions) is only needed to convert
//! to/from byte sequences, and since its purpose is reproducibility,
//! non-reproducible sources (e.g. `OsRng`) need not bother with it.

use crate::RngCore;
use core::cmp::min;

/// Implement `next_u64` via `next_u32`, little-endian order.
pub fn next_u64_via_u32<R: RngCore + ?Sized>(rng: &mut R) -> u64 {
    // Use LE; we explicitly generate one value before the next.
    let x = u64::from(rng.next_u32());
    let y = u64::from(rng.next_u32());
    (y << 32) | x
}

/// Implement `fill_bytes` via `next_u64` and `next_u32`, little-endian order.
///
/// The fastest way to fill a slice is usually to work as long as possible with
/// integers. That is why this method mostly uses `next_u64`, and only when
/// there are 4 or less bytes remaining at the end of the slice it uses
/// `next_u32` once.
pub fn fill_bytes_via_next<R: RngCore + ?Sized>(rng: &mut R, dest: &mut [u8]) {
    let mut left = dest;
    while left.len() >= 8 {
        let (l, r) = { left }.split_at_mut(8);
        left = r;
        let chunk: [u8; 8] = rng.next_u64().to_le_bytes();
        l.copy_from_slice(&chunk);
    }
    let n = left.len();
    if n > 4 {
        let chunk: [u8; 8] = rng.next_u64().to_le_bytes();
        left.copy_from_slice(&chunk[..n]);
    } else if n > 0 {
        let chunk: [u8; 4] = rng.next_u32().to_le_bytes();
        left.copy_from_slice(&chunk[..n]);
    }
}

macro_rules! fill_via_chunks {
    ($src:expr, $dst:expr, $ty:ty) => {{
        const SIZE: usize = core::mem::size_of::<$ty>();
        let chunk_size_u8 = min($src.len() * SIZE, $dst.len());
        let chunk_size = (chunk_size_u8 + SIZE - 1) / SIZE;

        let mut iter_src = $src.iter();
        let mut iter_chunks = $dst.chunks_exact_mut(SIZE);
        let mut n;
        loop {
            n = iter_src.next();
            if let (Some(n_), Some(chunk)) = (n, iter_chunks.next()) {
                chunk.copy_from_slice(&n_.to_le_bytes());
            } else {
                break;
            }
        }
        let rem = iter_chunks.into_remainder();
        if let Some(n_) = n {
            rem.copy_from_slice(&n_.to_le_bytes()[..rem.len()]);
        }

        (chunk_size, chunk_size_u8)
    }};
}

/// Implement `fill_bytes` by reading chunks from the output buffer of a block
/// based RNG.
///
/// The return values are `(consumed_u32, filled_u8)`.
///
/// `filled_u8` is the number of filled bytes in `dest`, which may be less than
/// the length of `dest`.
/// `consumed_u32` is the number of words consumed from `src`, which is the same
/// as `filled_u8 / 4` rounded up.
///
/// # Example
/// (from `IsaacRng`)
///
/// ```ignore
/// fn fill_bytes(&mut self, dest: &mut [u8]) {
///     let mut read_len = 0;
///     while read_len < dest.len() {
///         if self.index >= self.rsl.len() {
///             self.isaac();
///         }
///
///         let (consumed_u32, filled_u8) =
///             impls::fill_via_u32_chunks(&mut self.rsl[self.index..],
///                                        &mut dest[read_len..]);
///
///         self.index += consumed_u32;
///         read_len += filled_u8;
///     }
/// }
/// ```
pub fn fill_via_u32_chunks(src: &[u32], dest: &mut [u8]) -> (usize, usize) {
    fill_via_chunks!(src, dest, u32)
}

/// Implement `fill_bytes` by reading chunks from the output buffer of a block
/// based RNG.
///
/// The return values are `(consumed_u64, filled_u8)`.
/// `filled_u8` is the number of filled bytes in `dest`, which may be less than
/// the length of `dest`.
/// `consumed_u64` is the number of words consumed from `src`, which is the same
/// as `filled_u8 / 8` rounded up.
///
/// See `fill_via_u32_chunks` for an example.
pub fn fill_via_u64_chunks(src: &[u64], dest: &mut [u8]) -> (usize, usize) {
    fill_via_chunks!(src, dest, u64)
}

/// Implement `next_u32` via `fill_bytes`, little-endian order.
pub fn next_u32_via_fill<R: RngCore + ?Sized>(rng: &mut R) -> u32 {
    let mut buf = [0; 4];
    rng.fill_bytes(&mut buf);
    u32::from_ne_bytes(buf)
}

/// Implement `next_u64` via `fill_bytes`, little-endian order.
pub fn next_u64_via_fill<R: RngCore + ?Sized>(rng: &mut R) -> u64 {
    let mut buf = [0; 8];
    rng.fill_bytes(&mut buf);
    u64::from_ne_bytes(buf)
}

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

    #[test]
    fn test_fill_via_u32_chunks() {
        let src = [1, 2, 3];
        let mut dst = [0u8; 11];
        assert_eq!(fill_via_u32_chunks(&src, &mut dst), (3, 11));
        assert_eq!(dst, [1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0]);

        let mut dst = [0u8; 13];
        assert_eq!(fill_via_u32_chunks(&src, &mut dst), (3, 12));
        assert_eq!(dst, [1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0]);

        let mut dst = [0u8; 5];
        assert_eq!(fill_via_u32_chunks(&src, &mut dst), (2, 5));
        assert_eq!(dst, [1, 0, 0, 0, 2]);
    }

    #[test]
    fn test_fill_via_u64_chunks() {
        let src = [1, 2];
        let mut dst = [0u8; 11];
        assert_eq!(fill_via_u64_chunks(&src, &mut dst), (2, 11));
        assert_eq!(dst, [1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0]);

        let mut dst = [0u8; 17];
        assert_eq!(fill_via_u64_chunks(&src, &mut dst), (2, 16));
        assert_eq!(dst, [1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0]);

        let mut dst = [0u8; 5];
        assert_eq!(fill_via_u64_chunks(&src, &mut dst), (1, 5));
        assert_eq!(dst, [1, 0, 0, 0, 0]);
    }
}
<commit_msg>Simplify code<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Helper functions for implementing `RngCore` functions.
//!
//! For cross-platform reproducibility, these functions all use Little Endian:
//! least-significant part first. For example, `next_u64_via_u32` takes `u32`
//! values `x, y`, then outputs `(y << 32) | x`. To implement `next_u32`
//! from `next_u64` in little-endian order, one should use `next_u64() as u32`.
//!
//! Byte-swapping (like the std `to_le` functions) is only needed to convert
//! to/from byte sequences, and since its purpose is reproducibility,
//! non-reproducible sources (e.g. `OsRng`) need not bother with it.

use crate::RngCore;
use core::cmp::min;

/// Implement `next_u64` via `next_u32`, little-endian order.
pub fn next_u64_via_u32<R: RngCore + ?Sized>(rng: &mut R) -> u64 {
    // Use LE; we explicitly generate one value before the next.
    let x = u64::from(rng.next_u32());
    let y = u64::from(rng.next_u32());
    (y << 32) | x
}

/// Implement `fill_bytes` via `next_u64` and `next_u32`, little-endian order.
///
/// The fastest way to fill a slice is usually to work as long as possible with
/// integers. That is why this method mostly uses `next_u64`, and only when
/// there are 4 or less bytes remaining at the end of the slice it uses
/// `next_u32` once.
pub fn fill_bytes_via_next<R: RngCore + ?Sized>(rng: &mut R, dest: &mut [u8]) {
    let mut left = dest;
    while left.len() >= 8 {
        let (l, r) = { left }.split_at_mut(8);
        left = r;
        let chunk: [u8; 8] = rng.next_u64().to_le_bytes();
        l.copy_from_slice(&chunk);
    }
    let n = left.len();
    if n > 4 {
        let chunk: [u8; 8] = rng.next_u64().to_le_bytes();
        left.copy_from_slice(&chunk[..n]);
    } else if n > 0 {
        let chunk: [u8; 4] = rng.next_u32().to_le_bytes();
        left.copy_from_slice(&chunk[..n]);
    }
}

macro_rules! fill_via_chunks {
    ($src:expr, $dst:expr, $ty:ty) => {{
        const SIZE: usize = core::mem::size_of::<$ty>();
        let chunk_size_u8 = min($src.len() * SIZE, $dst.len());
        let chunk_size = (chunk_size_u8 + SIZE - 1) / SIZE;

        let mut iter_src = $src.iter();
        let mut chunks = $dst.chunks_exact_mut(SIZE);
        for (chunk, n) in (&mut chunks).zip(&mut iter_src) {
            chunk.copy_from_slice(&n.to_le_bytes());
        }
        let rem = chunks.into_remainder();
        if let Some(n) = iter_src.next() {
            rem.copy_from_slice(&n.to_le_bytes()[..rem.len()]);
        }

        (chunk_size, chunk_size_u8)
    }};
}

/// Implement `fill_bytes` by reading chunks from the output buffer of a block
/// based RNG.
///
/// The return values are `(consumed_u32, filled_u8)`.
///
/// `filled_u8` is the number of filled bytes in `dest`, which may be less than
/// the length of `dest`.
/// `consumed_u32` is the number of words consumed from `src`, which is the same
/// as `filled_u8 / 4` rounded up.
///
/// # Example
/// (from `IsaacRng`)
///
/// ```ignore
/// fn fill_bytes(&mut self, dest: &mut [u8]) {
///     let mut read_len = 0;
///     while read_len < dest.len() {
///         if self.index >= self.rsl.len() {
///             self.isaac();
///         }
///
///         let (consumed_u32, filled_u8) =
///             impls::fill_via_u32_chunks(&mut self.rsl[self.index..],
///                                        &mut dest[read_len..]);
///
///         self.index += consumed_u32;
///         read_len += filled_u8;
///     }
/// }
/// ```
pub fn fill_via_u32_chunks(src: &[u32], dest: &mut [u8]) -> (usize, usize) {
    fill_via_chunks!(src, dest, u32)
}

/// Implement `fill_bytes` by reading chunks from the output buffer of a block
/// based RNG.
///
/// The return values are `(consumed_u64, filled_u8)`.
/// `filled_u8` is the number of filled bytes in `dest`, which may be less than
/// the length of `dest`.
/// `consumed_u64` is the number of words consumed from `src`, which is the same
/// as `filled_u8 / 8` rounded up.
///
/// See `fill_via_u32_chunks` for an example.
pub fn fill_via_u64_chunks(src: &[u64], dest: &mut [u8]) -> (usize, usize) {
    fill_via_chunks!(src, dest, u64)
}

/// Implement `next_u32` via `fill_bytes`, little-endian order.
pub fn next_u32_via_fill<R: RngCore + ?Sized>(rng: &mut R) -> u32 {
    let mut buf = [0; 4];
    rng.fill_bytes(&mut buf);
    u32::from_ne_bytes(buf)
}

/// Implement `next_u64` via `fill_bytes`, little-endian order.
pub fn next_u64_via_fill<R: RngCore + ?Sized>(rng: &mut R) -> u64 {
    let mut buf = [0; 8];
    rng.fill_bytes(&mut buf);
    u64::from_ne_bytes(buf)
}

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

    #[test]
    fn test_fill_via_u32_chunks() {
        let src = [1, 2, 3];
        let mut dst = [0u8; 11];
        assert_eq!(fill_via_u32_chunks(&src, &mut dst), (3, 11));
        assert_eq!(dst, [1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0]);

        let mut dst = [0u8; 13];
        assert_eq!(fill_via_u32_chunks(&src, &mut dst), (3, 12));
        assert_eq!(dst, [1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0]);

        let mut dst = [0u8; 5];
        assert_eq!(fill_via_u32_chunks(&src, &mut dst), (2, 5));
        assert_eq!(dst, [1, 0, 0, 0, 2]);
    }

    #[test]
    fn test_fill_via_u64_chunks() {
        let src = [1, 2];
        let mut dst = [0u8; 11];
        assert_eq!(fill_via_u64_chunks(&src, &mut dst), (2, 11));
        assert_eq!(dst, [1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0]);

        let mut dst = [0u8; 17];
        assert_eq!(fill_via_u64_chunks(&src, &mut dst), (2, 16));
        assert_eq!(dst, [1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0]);

        let mut dst = [0u8; 5];
        assert_eq!(fill_via_u64_chunks(&src, &mut dst), (1, 5));
        assert_eq!(dst, [1, 0, 0, 0, 0]);
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 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.

#[inline(unknown)] //~ ERROR E0535
pub fn something() {}

fn main() {}
<commit_msg>Update invalid inline arg test.<commit_after>// Copyright 2016 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.

#[inline(unknown)] //~ ERROR E0535
pub fn something() {}

fn main() {
    something();
}
<|endoftext|>


<commit_before>use ram::chunk::Chunk;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::collections::HashSet;
use ram::segs::{Segment, MAX_SEGMENT_SIZE};
use ram::entry::{EntryContent, Entry, EntryType};
use itertools::Itertools;

use libc;

pub struct CombinedCleaner;

#[derive(Clone)]
struct DummyEntry {
    size: usize,
    addr: usize,
    cell_hash: Option<u64>,
    timestamp: u32,
}

struct DummySegment {
    head: usize,
    entries: Vec<DummyEntry>
}

impl DummySegment {
    fn new() -> DummySegment {
        DummySegment {
            head: 0, entries: Vec::new()
        }
    }
}

// Combine small segments into larger segments
// This cleaner will perform a greedy approach to relocate entries from old segments to fewer
// new segments to reduce number or segments and reclaim spaces from tombstones

// for higher hit rate for fetching cells in segments, we need to put data with close timestamp together
// this optimization is intended for enabling neb to contain data more than it's memory

impl CombinedCleaner {
    pub fn combine_segments(chunk: &Chunk, segments: &Vec<Arc<Segment>>) {
        if segments.len() < 2 { return }
        debug!("Combining segments");

        let segment_ids_to_combine: HashSet<_> = segments.iter().map(|seg| seg.id).collect();

        debug!("get all entries in segments to combine");
        let entries: Vec<_> = segments
            .iter()
            .flat_map(|seg| {
                chunk.live_entries(seg)
            })
            .filter(|entry| {
                // live entries have done a lot of filtering work already
                // but we still need to remove those tombstones that pointed to segments we are about to combine
                if let EntryContent::Tombstone(ref tombstone) = entry.content {
                    return !segment_ids_to_combine.contains(&tombstone.segment_id);
                }
                return true;
            })
            .map(|entry| {
                let entry_size = entry.meta.entry_size;
                let entry_addr = entry.meta.entry_pos;
                let cell_header = match entry.content {
                    EntryContent::Cell(header) => Some(header),
                    _ => None
                };
                DummyEntry {
                    size: entry_size,
                    addr: entry_addr,
                    timestamp: cell_header.map(|h| h.timestamp).unwrap_or(0),
                    cell_hash: cell_header.map(|h| h.hash),
                }
            })
            .group_by(|entry| {
                entry.timestamp
            })
            .into_iter()
            .map(|(t, group)| {
                let mut group: Vec<_> = group.collect();
                group.sort_by(|entry1,entry2| {
                    let size1 = entry1.size;
                    let size2 = entry2.size;
                    size1.cmp(&size2).reverse()
                });
                return (t, group.into_iter());
            })
            .sorted_by_key(|&(t, _)| t)
            .into_iter()
            .map(|(_, group)| group)
            .flatten()
            .collect();

        // provide additional state for whether entry have been claimed on simulation
        let mut entries: Vec<_> = entries.into_iter().map(|e| (e, false)).collect();

        debug!("simulate the combine process to determine the efficiency");
        let mut pending_segments = Vec::with_capacity(segments.len());
        let entries_num = entries.len();
        let mut entries_to_claim = entries_num;
        let mut cursor = 0;
        pending_segments.push(DummySegment::new());
        while entries_to_claim > 0 {
            let segment_space_remains =
                MAX_SEGMENT_SIZE - pending_segments.last().unwrap().head;
            let index = cursor % entries_num;
            let entry_pair = entries.get_mut(index).unwrap();
            if entry_pair.1 {
                // entry claimed
                cursor += 1;
                continue;
            }

            let entry = &entry_pair.0;
            let entry_size = entry.size;
            if entry_size > segment_space_remains {
                if index == entries_num - 1 {
                    // iterated to the last one, which means no more entries can be placed
                    // in the segment, then create a new segment
                    pending_segments.push(DummySegment::new());
                }
                cursor += 1;
                continue;
            }
            let mut last_segment = pending_segments.last_mut().unwrap();
            last_segment.entries.push(entry.clone());

            // pump dummy segment head pointer
            last_segment.head += entry_size;

            // mark entry claimed
            entry_pair.1 = true;
            entries_to_claim -= 1;

            // move to next entry
            cursor += 1;
        }

        debug!("Checking combine feasibility");
        let pending_segments_len = pending_segments.len();
        let segments_to_combine_len = segments.len();
        if pending_segments_len >= segments_to_combine_len  {
            warn!("Trying to combine segments but resulting segments still does not go down {}/{}",
                  pending_segments_len, segments_to_combine_len);
        }

        debug!("Updating cell reference");
        pending_segments
            .iter()
            .map(|dummy_seg| {
                let new_seg_id = chunk.next_segment_id();
                let new_seg = Segment::new(new_seg_id, dummy_seg.head, &chunk.backup_storage, &chunk.wal_storage);
                let mut cell_mapping = Vec::with_capacity(dummy_seg.entries.len());
                let mut seg_cursor = new_seg.addr;
                debug!("Combining segment to new one with id {}", new_seg_id);
                for entry in &dummy_seg.entries {
                    let entry_addr = entry.addr;
                    unsafe {
                        libc::memcpy(
                            seg_cursor as *mut libc::c_void,
                            entry_addr as *mut libc::c_void,
                            entry.size);
                    }
                    if let Some(cell_hash) = entry.cell_hash {
                        debug!("Marked cell relocation hash {}, addr {} to segment {}", cell_hash, entry_addr, new_seg_id);
                        cell_mapping.push((seg_cursor, entry_addr, cell_hash));
                    }
                    seg_cursor += entry.size;
                }
                new_seg.append_header.store(seg_cursor, Ordering::Relaxed);
                return (new_seg, cell_mapping);
            })
            .flat_map(|(segment, cells)| {
                debug!("Putting new segment {}", segment.id);
                let seg_ref = Arc::new(segment);
                chunk.put_segment(seg_ref.clone());
                seg_ref.archive();
                return cells;
            })
            .for_each(|(new, old, hash)| {
                debug!("Reset cell {} ptr from {} to {}", hash, old, new);
                if let Some(mut actual_addr) = chunk.index.get_mut(&hash) {
                    if *actual_addr == old {
                        *actual_addr = new
                    } else {
                        warn!("cell {} with address {}, have been changed to {} on combine", hash, old, *actual_addr);
                    }
                } else {
                    warn!("cell {} address {} have been removed on combine", hash, old);
                }
            });

        debug!("Removing old segments");
        for old_seg in segments {
            chunk.remove_segment(old_seg.id);
        }
    }

}<commit_msg>order from grater to lesser for greedy layout<commit_after>use ram::chunk::Chunk;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::collections::HashSet;
use ram::segs::{Segment, MAX_SEGMENT_SIZE};
use ram::entry::{EntryContent, Entry, EntryType};
use itertools::Itertools;

use libc;

pub struct CombinedCleaner;

#[derive(Clone)]
struct DummyEntry {
    size: usize,
    addr: usize,
    cell_hash: Option<u64>,
    timestamp: u32,
}

struct DummySegment {
    head: usize,
    entries: Vec<DummyEntry>
}

impl DummySegment {
    fn new() -> DummySegment {
        DummySegment {
            head: 0, entries: Vec::new()
        }
    }
}

// Combine small segments into larger segments
// This cleaner will perform a greedy approach to relocate entries from old segments to fewer
// new segments to reduce number or segments and reclaim spaces from tombstones

// for higher hit rate for fetching cells in segments, we need to put data with close timestamp together
// this optimization is intended for enabling neb to contain data more than it's memory

impl CombinedCleaner {
    pub fn combine_segments(chunk: &Chunk, segments: &Vec<Arc<Segment>>) {
        if segments.len() < 2 { return }
        debug!("Combining segments");

        let segment_ids_to_combine: HashSet<_> = segments.iter().map(|seg| seg.id).collect();

        debug!("get all entries in segments to combine and order them by data tempreture and size");
        let mut entries: Vec<_> = segments
            .iter()
            .flat_map(|seg| {
                chunk.live_entries(seg)
            })
            .filter(|entry| {
                // live entries have done a lot of filtering work already
                // but we still need to remove those tombstones that pointed to segments we are about to combine
                if let EntryContent::Tombstone(ref tombstone) = entry.content {
                    return !segment_ids_to_combine.contains(&tombstone.segment_id);
                }
                return true;
            })
            .map(|entry| {
                let entry_size = entry.meta.entry_size;
                let entry_addr = entry.meta.entry_pos;
                let cell_header = match entry.content {
                    EntryContent::Cell(header) => Some(header),
                    _ => None
                };
                DummyEntry {
                    size: entry_size,
                    addr: entry_addr,
                    timestamp: cell_header.map(|h| h.timestamp).unwrap_or(0),
                    cell_hash: cell_header.map(|h| h.hash),
                }
            })
            .group_by(|entry| {
                entry.timestamp
            })
            .into_iter()
            .map(|(t, group)| {
                let mut group: Vec<_> = group.collect();
                group.sort_by(|entry1,entry2| {
                    let size1 = entry1.size;
                    let size2 = entry2.size;
                    size1.cmp(&size2).reverse()
                });
                return (t, group.into_iter());
            })
            .sorted_by_key(|&(t, _)| t)
            .into_iter()
            .map(|(_, group)| group)
            .flatten()
            .collect();

        // order by temperature and size from greater to lesser
        entries.reverse();

        // provide additional state for whether entry have been claimed on simulation
        let mut entries: Vec<_> = entries.into_iter().map(|e| (e, false)).collect();

        debug!("simulate the combine process to determine the efficiency");
        let mut pending_segments = Vec::with_capacity(segments.len());
        let entries_num = entries.len();
        let mut entries_to_claim = entries_num;
        let mut cursor = 0;
        pending_segments.push(DummySegment::new());
        while entries_to_claim > 0 {
            let segment_space_remains =
                MAX_SEGMENT_SIZE - pending_segments.last().unwrap().head;
            let index = cursor % entries_num;
            let entry_pair = entries.get_mut(index).unwrap();
            if entry_pair.1 {
                // entry claimed
                cursor += 1;
                continue;
            }

            let entry = &entry_pair.0;
            let entry_size = entry.size;
            if entry_size > segment_space_remains {
                if index == entries_num - 1 {
                    // iterated to the last one, which means no more entries can be placed
                    // in the segment, then create a new segment
                    pending_segments.push(DummySegment::new());
                }
                cursor += 1;
                continue;
            }
            let mut last_segment = pending_segments.last_mut().unwrap();
            last_segment.entries.push(entry.clone());

            // pump dummy segment head pointer
            last_segment.head += entry_size;

            // mark entry claimed
            entry_pair.1 = true;
            entries_to_claim -= 1;

            // move to next entry
            cursor += 1;
        }

        debug!("Checking combine feasibility");
        let pending_segments_len = pending_segments.len();
        let segments_to_combine_len = segments.len();
        if pending_segments_len >= segments_to_combine_len  {
            warn!("Trying to combine segments but resulting segments still does not go down {}/{}",
                  pending_segments_len, segments_to_combine_len);
        }

        debug!("Updating cell reference");
        pending_segments
            .iter()
            .map(|dummy_seg| {
                let new_seg_id = chunk.next_segment_id();
                let new_seg = Segment::new(new_seg_id, dummy_seg.head, &chunk.backup_storage, &chunk.wal_storage);
                let mut cell_mapping = Vec::with_capacity(dummy_seg.entries.len());
                let mut seg_cursor = new_seg.addr;
                debug!("Combining segment to new one with id {}", new_seg_id);
                for entry in &dummy_seg.entries {
                    let entry_addr = entry.addr;
                    unsafe {
                        libc::memcpy(
                            seg_cursor as *mut libc::c_void,
                            entry_addr as *mut libc::c_void,
                            entry.size);
                    }
                    if let Some(cell_hash) = entry.cell_hash {
                        debug!("Marked cell relocation hash {}, addr {} to segment {}", cell_hash, entry_addr, new_seg_id);
                        cell_mapping.push((seg_cursor, entry_addr, cell_hash));
                    }
                    seg_cursor += entry.size;
                }
                new_seg.append_header.store(seg_cursor, Ordering::Relaxed);
                return (new_seg, cell_mapping);
            })
            .flat_map(|(segment, cells)| {
                debug!("Putting new segment {}", segment.id);
                let seg_ref = Arc::new(segment);
                chunk.put_segment(seg_ref.clone());
                seg_ref.archive();
                return cells;
            })
            .for_each(|(new, old, hash)| {
                debug!("Reset cell {} ptr from {} to {}", hash, old, new);
                if let Some(mut actual_addr) = chunk.index.get_mut(&hash) {
                    if *actual_addr == old {
                        *actual_addr = new
                    } else {
                        warn!("cell {} with address {}, have been changed to {} on combine", hash, old, *actual_addr);
                    }
                } else {
                    warn!("cell {} address {} have been removed on combine", hash, old);
                }
            });

        debug!("Removing old segments");
        for old_seg in segments {
            chunk.remove_segment(old_seg.id);
        }
    }

}<|endoftext|>